// 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;

    std::map<uint32_t, QByteArray> pkeys = preKeyPairs->readAll();
    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();
    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();
    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;
    bool had = true;
    try {
        devs = devices->getRecord(jid);
    } catch (const LMDBAL::NotFound& error) {
        had = false;
    }

    devs.insert(deviceId, device);
    if (had) {
        devices->changeRecord(jid, devs);
    } else {
        devices->addRecord(jid, devs);
    }

    return Core::makeReadyTask();
}

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

    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) {
    QHash<uint32_t, Device> devs = devices->getRecord(jid);
    devs.remove(deviceId);
    if (devs.isEmpty()) {
        devices->removeRecord(jid);
    } else {
        devices->changeRecord(jid, devs);
    }
    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) {
    preKeyPairs->removeRecord(keyId);
    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 {
        devs = devices->getRecord(jid);
    } 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, QDateTime(), Shared::TrustLevel::undecided, Shared::EncryptionProtocol::omemo2, false);
    }
}

void Core::OmemoHandler::requestBundles(const QString& jid) {
    QXmppTask<void> task = acc->om->buildMissingSessions({jid});
    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;
    acc->oh->getDevices(jid, keys);
    std::map<QByteArray, Shared::TrustLevel> trustLevels = acc->th->getKeys(Shared::EncryptionProtocol::omemo2, jid);

    qDebug() << "OMEMO info for " << jid << " devices:" << keys.size() << ", trustLevels:" << trustLevels.size();
    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;
            qDebug() << "Found a trust level for a device!";
        }
    }

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

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

    qDebug() << "OMEMO info for " << jid << " devices:" << keys.size() << ", trustLevels:" << trustLevels.size();
    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;
            qDebug() << "Found a trust level for a device!";
        }
    }

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

void Core::OmemoHandler::onOmemoDeviceAdded(const QString& jid, uint32_t 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;
}