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

Core::OmemoHandler::OmemoHandler(Account* account) :
    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, QXmppOmemoStorage::SignedPreKeyPair>("signedPreKeyPairs"))
{
    db.open();
    try {
        QVariant own = meta->getRecord("ownDevice");
        ownDevice.value() = own.value<OwnDevice>();
        qDebug() << "Successfully found own device omemo data for account" << acc->getName();
    } catch (const DataBase::NotFound& e) {
        qDebug() << "No device omemo data was found for account" << acc->getName();
    }
}

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

QFuture<void> Core::OmemoHandler::emptyVoidFuture() {
    QFutureInterface<QXmppOmemoStorage::OmemoData> result(QFutureInterfaceBase::Started);
    result.reportFinished();
    return result.future();
}

QFuture<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, QXmppOmemoStorage::SignedPreKeyPair> spre = signedPreKeyPairs->readAll();
    for (const std::pair<const uint32_t, QXmppOmemoStorage::SignedPreKeyPair>& pair : spre) {
        data.signedPreKeyPairs.insert(pair.first, pair.second);
    }

    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);
    }

    QFutureInterface<QXmppOmemoStorage::OmemoData> result(QFutureInterfaceBase::Started);
    result.reportResult(std::move(data));
    result.reportFinished();
    return result.future();
}

QFuture<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 DataBase::NotFound& error) {
        had = false;
    }

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

    return emptyVoidFuture();
}

QFuture<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 emptyVoidFuture();
}

QFuture<void> Core::OmemoHandler::addSignedPreKeyPair(uint32_t keyId, const QXmppOmemoStorage::SignedPreKeyPair& keyPair) {
    signedPreKeyPairs->addRecord(keyId, keyPair);
    return emptyVoidFuture();
}

QFuture<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 emptyVoidFuture();
}

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

QFuture<void> Core::OmemoHandler::removePreKeyPair(uint32_t keyId) {
    preKeyPairs->removeRecord(keyId);
    return emptyVoidFuture();
}

QFuture<void> Core::OmemoHandler::removeSignedPreKeyPair(uint32_t keyId) {
    signedPreKeyPairs->removeRecord(keyId);
    return emptyVoidFuture();
}

QFuture<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 emptyVoidFuture();
}

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

    return emptyVoidFuture();
}

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

QDataStream & operator >> (QDataStream& in, QXmppOmemoStorage::SignedPreKeyPair& pair) {
    in >> pair.creationDate;
    in >> pair.data;

    return in;
}

QDataStream & operator << (QDataStream& out, const QXmppOmemoStorage::SignedPreKeyPair& pair) {
    out << pair.creationDate;
    out << pair.data;

    return out;
}