/*
 * Squawk messenger.
 * Copyright (C) 2019  Yury Gubich <blue@macaw.me>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <QDebug>
#include "omemohandler.h"
#include "core/account.h"
#include "core/adapterfunctions.h"

Core::OmemoHandler::OmemoHandler(Account* account) :
    QObject(),
    QXmppOmemoStorage(),
    acc(account),
    ownDevice(std::nullopt),
    db(acc->getName() + "/omemo"),
    meta(db.addCache<QString, QVariant>("meta")),
    devices(db.addCache<QString, QHash<uint32_t, Device>>("devices")),
    preKeyPairs(db.addCache<uint32_t, QByteArray>("preKeyPairs")),
    signedPreKeyPairs(db.addCache<uint32_t, SignedPreKeyPair>("signedPreKeyPairs"))
{
    db.open();
    try {
        QVariant own = meta->getRecord("ownDevice");
        ownDevice = own.value<OwnDevice>();
        qDebug() << "Successfully found own device omemo data for account" << acc->getName();
    } catch (const LMDBAL::NotFound& e) {
        qDebug() << "No device omemo data was found for account" << acc->getName();
    }
}

Core::OmemoHandler::~OmemoHandler() {
    db.close();
}

bool Core::OmemoHandler::hasOwnDevice() {
    return ownDevice.has_value();
}

QXmppTask<QXmppOmemoStorage::OmemoData> Core::OmemoHandler::allData() {
    OmemoData data;
    data.ownDevice = ownDevice;

    LMDBAL::Transaction txn = db.beginReadOnlyTransaction();
    std::map<uint32_t, QByteArray> pkeys = preKeyPairs->readAll(txn);
    for (const std::pair<const uint32_t, QByteArray>& pair : pkeys)
        data.preKeyPairs.insert(pair.first, pair.second);

    std::map<uint32_t, SignedPreKeyPair> spre = signedPreKeyPairs->readAll(txn);
    for (const std::pair<const uint32_t, SignedPreKeyPair>& pair : spre) {
        QXmppOmemoStorage::SignedPreKeyPair qxpair = {pair.second.first, pair.second.second};
        data.signedPreKeyPairs.insert(pair.first, qxpair);
    }

    std::map<QString, QHash<uint32_t, Device>> devs = devices->readAll(txn);
    for (const std::pair<const QString, QHash<uint32_t, Device>>& pair : devs)
        data.devices.insert(pair.first, pair.second);

    return Core::makeReadyTask(std::move(data));
}

QXmppTask<void> Core::OmemoHandler::addDevice(const QString& jid, uint32_t deviceId, const QXmppOmemoStorage::Device& device) {
    QHash<uint32_t, Device> devs;
    LMDBAL::WriteTransaction txn = db.beginTransaction();
    bool had = true;
    try {
        devices->getRecord(jid, devs, txn);
    } catch (const LMDBAL::NotFound& error) {
        had = false;
    }

    devs.insert(deviceId, device);  //overwrites
    if (had)
        devices->changeRecord(jid, devs, txn);
    else
        devices->addRecord(jid, devs, txn);

    txn.commit();
    return Core::makeReadyTask();
}

QXmppTask<void> Core::OmemoHandler::addPreKeyPairs(const QHash<uint32_t, QByteArray>& keyPairs) {
    LMDBAL::WriteTransaction txn = db.beginTransaction();
    for (QHash<uint32_t, QByteArray>::const_iterator itr = keyPairs.begin(), end = keyPairs.end(); itr != end; ++itr)
        preKeyPairs->forceRecord(itr.key(), itr.value(), txn);

    txn.commit();
    return Core::makeReadyTask();
}

QXmppTask<void> Core::OmemoHandler::addSignedPreKeyPair(uint32_t keyId, const QXmppOmemoStorage::SignedPreKeyPair& keyPair) {
    signedPreKeyPairs->forceRecord(keyId, std::make_pair(keyPair.creationDate, keyPair.data));
    return Core::makeReadyTask();
}

QXmppTask<void> Core::OmemoHandler::removeDevice(const QString& jid, uint32_t deviceId) {
    LMDBAL::WriteTransaction txn = db.beginTransaction();
    QHash<uint32_t, Device> devs = devices->getRecord(jid, txn);
    devs.remove(deviceId);
    if (devs.isEmpty())
        devices->removeRecord(jid, txn);
    else
        devices->changeRecord(jid, devs, txn);

    txn.commit();
    return Core::makeReadyTask();
}

QXmppTask<void> Core::OmemoHandler::removeDevices(const QString& jid) {
    devices->removeRecord(jid);
    return Core::makeReadyTask();
}

QXmppTask<void> Core::OmemoHandler::removePreKeyPair(uint32_t keyId) {
    try {
        preKeyPairs->removeRecord(keyId);
    } catch (const LMDBAL::NotFound& e) {
        qDebug() << "Couldn't remove preKeyPair " << e.what();
    }
    return Core::makeReadyTask();
}

QXmppTask<void> Core::OmemoHandler::removeSignedPreKeyPair(uint32_t keyId) {
    try {
        signedPreKeyPairs->removeRecord(keyId);
    } catch (const LMDBAL::NotFound& e) {}
    return Core::makeReadyTask();
}

QXmppTask<void> Core::OmemoHandler::setOwnDevice(const std::optional<OwnDevice>& device) {
    bool had = ownDevice.has_value();
    ownDevice = device;
    if (ownDevice.has_value()) {
        if (had)
            meta->changeRecord("ownDevice", QVariant::fromValue(ownDevice.value()));
        else
            meta->addRecord("ownDevice", QVariant::fromValue(ownDevice.value()));
    } else if (had) {
        meta->removeRecord("ownDevice");
    }
    return Core::makeReadyTask();
}

QXmppTask<void> Core::OmemoHandler::resetAll() {
    ownDevice = std::nullopt;
    db.drop();

    return Core::makeReadyTask();
}

void Core::OmemoHandler::getDevices(const QString& jid, std::list<Shared::KeyInfo>& out) const {
    QHash<uint32_t, Device> devs;
    try {
        devices->getRecord(jid, devs);
    } catch (const LMDBAL::NotFound& error) {}

    for (QHash<uint32_t, Device>::const_iterator itr = devs.begin(), end = devs.end(); itr != end; ++itr) {
        const Device& dev = itr.value();
        out.emplace_back(
            itr.key(),
            dev.keyId,
            dev.label,
            dev.removalFromDeviceListDate,
            Shared::TrustLevel::undecided,
            Shared::EncryptionProtocol::omemo2,
            false
        );
    }
}

void Core::OmemoHandler::requestBundles(const QString& jid) {
    QXmppTask<void> task = acc->om->buildMissingSessions({jid});
    Contact* cnt = acc->rh->getContact(jid);
    if (cnt)
        cnt->omemoBundles = Shared::Possible::discovering;

    task.then(this, std::bind(&OmemoHandler::onBundlesReceived, this, jid));
}

void Core::OmemoHandler::requestOwnBundles() {
    QXmppTask<void> task = acc->om->buildMissingSessions({acc->getBareJid()});
    task.then(this, std::bind(&OmemoHandler::onOwnBundlesReceived, this));
}

void Core::OmemoHandler::onBundlesReceived(const QString& jid) {
    std::list<Shared::KeyInfo> keys = readKeys(jid);

    Contact* cnt = acc->rh->getContact(jid);
    if (cnt)
        cnt->omemoBundles = Shared::Possible::present;

    acc->delay->receivedBundles(jid, keys);
}

void Core::OmemoHandler::onOwnBundlesReceived() {
    std::list<Shared::KeyInfo> keys = readKeys(acc->getBareJid());
    if (ownDevice)
        keys.emplace_front(
            ownDevice->id,
            ownDevice->publicIdentityKey,
            ownDevice->label,
            QDateTime::currentDateTime(),
            Shared::TrustLevel::authenticated,
            Shared::EncryptionProtocol::omemo2,
            true
        );

    acc->delay->receivedOwnBundles(keys);
}

std::list<Shared::KeyInfo> Core::OmemoHandler::readKeys(const QString& jid) {
    std::list<Shared::KeyInfo> keys;
    getDevices(jid, keys);
    std::map<QByteArray, Shared::TrustLevel> trustLevels = acc->th->getKeys(Shared::EncryptionProtocol::omemo2, jid);

    for (Shared::KeyInfo& key : keys) {
        std::map<QByteArray, Shared::TrustLevel>::const_iterator itr = trustLevels.find(key.fingerPrint);
        if (itr != trustLevels.end())
            key.trustLevel = itr->second;
    }

    return keys;
}

void Core::OmemoHandler::onOmemoDeviceAdded(const QString& jid, uint32_t id) {
    SHARED_UNUSED(id);
    qDebug() << "OMEMO device added for" << jid;
}

QDataStream & operator >> (QDataStream& in, QXmppOmemoStorage::Device& device) {
    in >> device.label;
    in >> device.keyId;
    in >> device.session;
    in >> device.unrespondedSentStanzasCount;
    in >> device.unrespondedReceivedStanzasCount;
    in >> device.removalFromDeviceListDate;

    return in;
}

QDataStream & operator << (QDataStream& out, const QXmppOmemoStorage::Device& device) {
    out << device.label;
    out << device.keyId;
    out << device.session;
    out << device.unrespondedSentStanzasCount;
    out << device.unrespondedReceivedStanzasCount;
    out << device.removalFromDeviceListDate;

    return out;
}

QDataStream & operator >> (QDataStream& in, QXmppOmemoStorage::OwnDevice& device) {
    in >> device.id;
    in >> device.label;
    in >> device.privateIdentityKey;
    in >> device.publicIdentityKey;
    in >> device.latestSignedPreKeyId;
    in >> device.latestPreKeyId;

    return in;
}

QDataStream & operator << (QDataStream& out, const QXmppOmemoStorage::OwnDevice& device) {
    out << device.id;
    out << device.label;
    out << device.privateIdentityKey;
    out << device.publicIdentityKey;
    out << device.latestSignedPreKeyId;
    out << device.latestPreKeyId;

    return out;
}