#include "jiddb.h"
#include <QDir>
#include <QSqlQuery>
#include <QSqlError>
#include <QStandardPaths>
#include <QTimeZone>
#include <QVariant>
#include <iostream>
#include <stdexcept>

JidDb::JidDb(const QString &jid) :
    jid(jid),
    db(QSqlDatabase::addDatabase("QSQLITE", jid))
{
    const auto path = QStandardPaths::writableLocation(
            QStandardPaths::AppDataLocation);
    const QDir dir;

    if (!dir.exists(path) && !dir.mkdir(path))
        throw std::runtime_error("Failed to create app dir");

    const auto abspath = path + "/" + jid + ".db";

    db.setDatabaseName(abspath);

    if (!db.open())
        throw std::runtime_error(qPrintable("database" + abspath
            + "could not be created"));

    QSqlQuery q(db);

    if (!q.exec("create table if not exists roster "
        "(jid TEXT unique) strict;"))
        std::cerr << "JidDb::JidDb: query exec failed" << std::endl;
}

int JidDb::addToRoster(const QString &jid)
{
    QSqlQuery q(db);

    if (!q.exec("insert or ignore into roster (jid) values ('" + jid + "');"))
    {
        std::cerr << "JidDb::addToRoster: query exec failed" << std::endl;
        return -1;
    }

    Q_EMIT addedToRoster(jid);
    return 0;
}

int JidDb::addToRoster(const QStringList &roster)
{
    for (const auto &r : roster)
        if (addToRoster(r))
            return -1;

    return 0;
}

QStringList JidDb::roster() const
{
    QStringList ret;
    QSqlQuery q(db);

    if (!q.exec("select jid from roster;"))
        std::cerr << "query exec failed" << std::endl;
    else
        while (q.next())
            ret.append(q.value("jid").toString());

    return ret;
}

int JidDb::ensureContactTable(const QString &jid) const
{
    QSqlQuery q(db);

    if (!q.exec("create table if not exists '" + jid +
        "' (direction TEXT, time INTEGER, body TEXT) strict;"))
    {
        std::cerr << "JidDb::storeMessage: query exec failed" << std::endl;
        return -1;
    }

    return 0;
}

static QString securityPolicyTableName(const QString &encryption)
{
    return "securityPolicy/" + encryption;
}

int JidDb::ensureSecurityPolicyTable(const QString &encryption,
    QString &table) const
{
    QSqlQuery q(db);

    table = securityPolicyTableName(encryption);

    if (!q.exec("create table if not exists '" + table
        + "' (policy TEXT unique) strict;"))
    {
        std::cerr << "JidDb::ensureSecurityPolicyTable: query exec failed"
            << std::endl;
        return -1;
    }

    return 0;
}

static QString keysTableName(const QString &encryption)
{
    return "keys/" + encryption;
}

int JidDb::ensureKeysTable(const QString &encryption, QString &table) const
{
    QSqlQuery q(db);

    table = keysTableName(encryption);

    if (!q.exec("create table if not exists '" + table
        + "' (owner TEXT, trustlevel TEXT, key TEXT) strict;"))
    {
        std::cerr << "JidDb::ensureKeysTable: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return -1;
    }

    return 0;
}

QList<JidDb::Message> JidDb::messages(const QString &jid,
    const int tail) const
{
    QSqlQuery q(db);

    if (ensureContactTable(jid))
        return QList<Message>();
    else if (tail < 0)
    {
        if (!q.exec("select * from '" + jid + "' order by time;"))
        {
            std::cerr << "JidDb::messages: query exec failed: "
                << qPrintable(q.lastError().text()) << std::endl;
            return QList<Message>();
        }
    }
    else if (!q.exec("select * from '" + jid + "' order by time desc limit "
        + QString::number(tail) + ";"))
    {
        std::cerr << "JidDb::messages: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return QList<Message>();
    }

    QList<Message> ret;

    while (q.next())
    {
        Message m;
        bool ok;
        qlonglong t = q.value("time").toLongLong(&ok);

        if (!ok)
        {
            std::cerr << "JidDb::messages: invalid time" << std::endl;
            // Attempt to read other messages.
            continue;
        }

        m.body = q.value("body").toString();
        m.dt = QDateTime::fromSecsSinceEpoch(t, QTimeZone::systemTimeZone());

        const auto dirstr = q.value("direction").toString();

        if (dirstr == "in")
            m.direction = Direction::In;
        else if (dirstr == "out")
            m.direction = Direction::Out;
        else
        {
            std::cerr << "JidDb::messages: invalid direction" << std::endl;
            // Attempt to read other messages.
            continue;
        }

        ret << m;
    }

    return ret;
}

int JidDb::store(const JidDb::Message &msg) const
{
    QSqlQuery q(db);
    QString dir;

    switch (msg.direction)
    {
        case Direction::In:
            dir = "in";
            break;

        case Direction::Out:
            dir = "out";
            break;
    }

    ensureContactTable(msg.contact);

    if (!q.exec("insert into '" + msg.contact
        + "' (direction, time, body) values "
        "('" + dir + "', "
        + QString::number(msg.dt.toSecsSinceEpoch()) + ", "
        "'" + msg.body + "')"))
    {
        std::cerr << "JidDb::storeMessage: query exec 2 failed" << std::endl;
        return -1;
    }

    return 0;
}

QStringList JidDb::tables() const
{
    QSqlQuery q(db);

    if (!q.exec("select name from sqlite_schema where "
        "type = 'table' and name not like 'sqlite_%'"))
    {
        std::cerr << "JidDb::conversations: query failed" << std::endl;
        return QStringList();
    }

    QStringList ret;

    while (q.next())
        ret << q.value("name").toString();

    return ret;
}

QList<JidDb::Conversation> JidDb::conversations() const
{
    const auto conversations = tables();
    QList<Conversation> ret;

    qDebug() << conversations;

    for (const auto &jid : conversations)
        if (jid.contains('@'))
        {
            const auto msgs = messages(jid, 1);

            if (!msgs.isEmpty())
            {
                const auto &m = msgs.first();

                ret << Conversation(jid, m.body, m.dt);
            }
        }

    return ret;
}

QString JidDb::securityPolicy(const QString &encryption) const
{
    const auto tb = tables();
    const auto name = securityPolicyTableName(encryption);

    for (const auto &t : tb)
    {
        if (t == name)
        {
            QStringList ret;
            QSqlQuery q(db);

            if (!q.exec("select policy from '" + name + "';"))
                std::cerr << "query exec failed" << std::endl;
            else
                while (q.next())
                    return q.value("policy").toString();

            break;
        }
    }

    return QString();
}

std::optional<QXmppOmemoStorage::OwnDevice> JidDb::omemoOwnDevice() const
{
    return std::optional<QXmppOmemoStorage::OwnDevice>();
}

int JidDb::store(const QString &encryption, const QString &policy) const
{
    QString table;

    if (ensureSecurityPolicyTable(encryption, table))
        return -1;

    QSqlQuery q(db);

    if (!q.exec("insert or ignore into '" + table
        + "' (policy) values ('" + policy + "');"))
    {
        std::cerr << "JidDb::store: query exec failed"
            << std::endl;
        return -1;
    }

    return 0;
}

QList<JidDb::Keys> JidDb::keys(const QString &encryption) const
{
    QList<Keys> ret;
    QString table;

    if (ensureKeysTable(encryption, table))
    {
        std::cerr << "JidDb::keys: ensureKeysTable failed" << std::endl;
        return QList<Keys>();
    }

    QSqlQuery q(db);

    if (!q.exec("select owner, trustlevel, key from '" + table + "';"))
        std::cerr << "JidDb::keys: query exec failed" << std::endl;
    else
        while (q.next())
        {
            const auto owner = q.value("owner").toString();
            const auto trust_level = q.value("trustlevel").toString();
            const auto key_b64 = q.value("key").toString();
            const auto key = QByteArray::fromBase64(key_b64.toLatin1());
            bool found = false;

            for (auto &k : ret)
                if (k.owner == owner)
                {
                    k.keys.append(key);
                    found = true;
                    break;
                }

            if (!found)
            {
                Keys keys;

                keys.owner = owner;
                keys.trust_level = trust_level;
                keys.keys.append(key);
                ret.append(keys);
            }
        }

    return ret;
}

int JidDb::store(const QString &encryption, const JidDb::Keys &keys) const
{
    QString table;

    if (ensureKeysTable(encryption, table))
    {
        std::cerr << "JidDb::store: ensureKeysTable failed" << std::endl;
        return -1;
    }

    for (const auto &k : keys.keys)
    {
        QSqlQuery q(db);

        const auto kstr = k.isEmpty() ? "\"\""
            : "'" + QString::fromLatin1(k.toBase64()) + "'";

        if (!q.exec("insert or replace into '" + table
            + "' (owner, trustlevel, key) values ("
            "'" + keys.owner + "', "
            "'" + keys.trust_level + "',"
            + kstr
            + ");"))
        {
            std::cerr << "JidDb::store: query exec failed: "
                << qPrintable(q.lastError().text()) << std::endl;
            return -1;
        }
    }

    return 0;
}

int JidDb::ensureOmemoDeviceTable() const
{
    QSqlQuery q(db);

    if (!q.exec("create table if not exists 'omemo/devices' "
        "(jid TEXT, deviceID INTEGER, keyID TEXT, session TEXT, "
        "unresp_sent INTEGER, unresp_recv INTEGER, removal INTEGER) "
        "strict;"))
    {
        std::cerr << "JidDb::ensureOmemoTable: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return -1;
    }

    return 0;
}

int JidDb::store(const QString &jid, const uint32_t id,
    const QXmppOmemoStorage::Device &d) const
{
    if (ensureOmemoDeviceTable())
    {
        std::cerr << "JidDb::store: ensureOmemoTable failed" << std::endl;
        return -1;
    }

    QSqlQuery q(db);
    const auto keyId = d.keyId.isEmpty() ? "\"\""
            : "'" + QString::fromLatin1(d.keyId.toBase64()) + "'",
        session = d.session.isEmpty() ? "\"\""
            : "'" + QString::fromLatin1(d.session.toBase64()) + "'",
        query = "insert or replace into 'omemo/devices' "
            "(jid, deviceID, keyID, session, "
            "unresp_sent, unresp_recv , removal) values ("
            "'" + jid + "', "
            + QString::number(id) + ", "
            + keyId + ", "
            + session + ", "
            + QString::number(d.unrespondedSentStanzasCount) + ", "
            + QString::number(d.unrespondedReceivedStanzasCount) + ", "
            + QString::number(d.removalFromDeviceListDate.toSecsSinceEpoch())
            + ");";

    if (!q.exec(query))
    {
        std::cerr << "JidDb::store(id, jid, d): query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return -1;
    }

    return 0;
}

int JidDb::ensureOmemoSignedPreKeyTable() const
{
    QSqlQuery q(db);

    if (!q.exec("create table if not exists 'omemo/signedprekeys' "
        "(id INTEGER, creation INTEGER, data TEXT) strict;"))
    {
        std::cerr << "JidDb::ensureOmemoSignedPreKeyTable: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return -1;
    }

    return 0;
}

int JidDb::store(uint32_t id, const QXmppOmemoStorage::SignedPreKeyPair &spk)
    const
{
    if (ensureOmemoSignedPreKeyTable())
    {
        std::cerr << "JidDb::store: ensureOmemoSignedPreKeyTable failed"
            << std::endl;
        return -1;
    }

    QSqlQuery q(db);
    const auto data = spk.data.isEmpty() ? "\"\""
        : "'" + QString::fromLatin1(spk.data.toBase64()) + "'",
        query = "insert or replace into 'omemo/signedprekeys' "
            "(id, creation, data) values ("
            + QString::number(id) + ", "
            + QString::number(spk.creationDate.toSecsSinceEpoch()) + ", "
            + data
            + ");";

    if (!q.exec(query))
    {
        std::cerr << "JidDb::store(spk): query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return -1;
    }

    return 0;
}

int JidDb::ensureOmemoPreKeyTable() const
{
    QSqlQuery q(db);

    if (!q.exec("create table if not exists 'omemo/prekeys' "
        "(id INTEGER, data TEXT) strict;"))
    {
        std::cerr << "JidDb::ensureOmemoPreKeyTable: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return -1;
    }

    return 0;
}

int JidDb::store(const QHash<uint32_t, QByteArray> &pairs) const
{
    if (ensureOmemoPreKeyTable())
    {
        std::cerr << "JidDb::store: ensureOmemoPreKeyTable failed"
            << std::endl;
        return -1;
    }

    for (const auto id : pairs.keys())
    {
        const auto data = pairs.value(id);
        const auto datastr = data.isEmpty() ? "\"\""
            : "'" + QString::fromLatin1(data.toBase64()) + "'",
            query = "insert or replace into 'omemo/prekeys' "
                "(id, data) values ("
                + QString::number(id) + ", "
                + datastr
                + ");";

        QSqlQuery q(db);

        if (!q.exec(query))
        {
            std::cerr << "JidDb::store(pairs): query exec failed: "
                << qPrintable(q.lastError().text()) << std::endl;
            return -1;
        }
    }

    return 0;
}

int JidDb::ensureOmemoOwnDeviceTable() const
{
    QSqlQuery q(db);

    if (!q.exec("create table if not exists 'omemo/owndevice' "
        "(id INTEGER, label TEXT, privkey TEXT, pubkey TEXT, "
        "latestSignedPreKeyId INTEGER, latestPreKeyId INTEGER) strict;"))
    {
        std::cerr << "JidDb::ensureOmemoOwnDeviceTable: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return -1;
    }

    return 0;
}

int JidDb::store(const QXmppOmemoStorage::OwnDevice &own) const
{
    if (ensureOmemoOwnDeviceTable())
    {
        std::cerr << "JidDb::store: ensureOmemoOwnDeviceTable failed"
            << std::endl;
        return -1;
    }

    QSqlQuery q(db);
    const auto privkey = own.privateIdentityKey.isEmpty() ? "\"\""
            : "'" + QString::fromLatin1(own.privateIdentityKey.toBase64()) + "'",
        pubkey = own.publicIdentityKey.isEmpty() ? "\"\""
            : "'" + QString::fromLatin1(own.publicIdentityKey.toBase64()) + "'",
        query = "insert or replace into 'omemo/owndevice' "
            "(id, label, privkey, pubkey, latestSignedPreKeyId, "
            "latestPreKeyId) values ("
            + QString::number(own.id) + ", "
            "'" + own.label + "', "
            + privkey + ", "
            + pubkey + ", "
            + QString::number(own.latestSignedPreKeyId) + ", "
            + QString::number(own.latestPreKeyId)
            + ");";

    if (!q.exec(query))
    {
        std::cerr << "JidDb::store(own): query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        std::cerr << "Query was: " << qPrintable(query) << std::endl;
        return -1;
    }

    return 0;
}

QList<JidDb::OmemoDevice> JidDb::omemoDevices() const
{
    QList<OmemoDevice> ret;
    QSqlQuery q(db);

    if (!q.exec("select * from 'omemo/devices';"))
    {
        std::cerr << "JidDb::omemoDevices: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return QList<OmemoDevice>();
    }

    while (q.next())
    {
        QXmppOmemoStorage::Device device;
        bool ok;
        const uint32_t deviceId = q.value("deviceID").toULongLong(&ok);

        if (!ok)
        {
            std::cerr << "JidDb::messages: invalid deviceID" << std::endl;
            // Attempt to read other messages.
            continue;
        }

        device.unrespondedReceivedStanzasCount =
            q.value("unresp_recv").toInt(&ok);

        if (!ok)
        {
            std::cerr << "JidDb::messages: invalid unresp_recv" << std::endl;
            // Attempt to read other messages.
            continue;
        }

        device.unrespondedSentStanzasCount = q.value("unresp_sent").toInt(&ok);

        if (!ok)
        {
            std::cerr << "JidDb::messages: invalid unresp_sent" << std::endl;
            // Attempt to read other messages.
            continue;
        }

        const QByteArray keyId = q.value("keyID").toByteArray(),
            session = q.value("session").toByteArray();
        const auto jid = q.value("jid").toString();
        bool found = false;

        for (auto &r : ret)
            if (jid == r.jid)
            {
                r.devices.insert(deviceId, device);
                break;
            }

        if (!found)
        {
            OmemoDevice d;

            d.jid = jid;
            d.devices.insert(deviceId, device);
            ret << d;
        }
    }

    return ret;
}

QHash<uint32_t, QXmppOmemoStorage::SignedPreKeyPair> JidDb::signedPreKeyPairs()
    const
{
    QHash<uint32_t, QXmppOmemoStorage::SignedPreKeyPair> ret;

    QSqlQuery q(db);

    if (!q.exec("select * from 'omemo/signedprekeys';"))
    {
        std::cerr << "JidDb::signedPreKeyPairs: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return QHash<uint32_t, QXmppOmemoStorage::SignedPreKeyPair>();
    }

    while (q.next())
    {
        QXmppOmemoStorage::SignedPreKeyPair spk;
        bool ok;
        const uint32_t id = q.value("id").toULongLong(&ok);

        if (!ok)
        {
            std::cerr << "JidDb::messages: invalid id" << std::endl;
            // Attempt to read other entries.
            continue;
        }

        const auto creation = q.value("creation").toInt(&ok);

        if (!ok)
        {
            std::cerr << "JidDb::messages: invalid creation date" << std::endl;
            // Attempt to read other entries.
            continue;
        }

        spk.creationDate = QDateTime::fromSecsSinceEpoch(creation);
        spk.data = q.value("data").toString().toLatin1().toBase64();
        ret.insert(id, spk);
    }

    return ret;
}

QHash<uint32_t, QByteArray> JidDb::preKeyPairs() const
{
    QHash<uint32_t, QByteArray> ret;

    QSqlQuery q(db);

    if (!q.exec("select * from 'omemo/signedprekeys';"))
    {
        std::cerr << "JidDb::signedPreKeyPairs: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return QHash<uint32_t, QByteArray>();
    }

    while (q.next())
    {
        bool ok;
        const uint32_t id = q.value("id").toULongLong(&ok);

        if (!ok)
        {
            std::cerr << "JidDb::messages: invalid id" << std::endl;
            // Attempt to read other entries.
            continue;
        }

        const auto data = q.value("data").toString().toLatin1().toBase64();
        ret.insert(id, data);
    }

    return ret;
}

void JidDb::removeKeys(const QString &encryption) const
{
    QString table;

    if (ensureKeysTable(encryption, table))
        std::cerr << "JidDb::removeKeys: ensureKeysTable failed" << std::endl;
    else
    {
        QSqlQuery q(db);

        if (!q.exec("delete from '" + table + "'"))
            std::cerr << "JidDb::removeKeys: query exec failed: "
                << qPrintable(q.lastError().text()) << std::endl;
    }
}

void JidDb::removeOmemoDevice(const QString &jid, const uint32_t id) const
{
    if (ensureOmemoDeviceTable())
        std::cerr << "JidDb::removeOmemoDevice: ensureOmemoDeviceTable failed"
            << std::endl;
    else
    {
        QSqlQuery q(db);
        const auto idstr = QString::number(id);

        if (!q.exec("delete from 'omemo/devices' "
            "where (jid = '" + jid + "' "
            "and id = " + idstr + ")"))
            std::cerr << "JidDb::removeOmemoDevice: query exec failed: "
                << qPrintable(q.lastError().text()) << std::endl;
    }
}

void JidDb::removeOmemoDevices(const QString &jid) const
{
    if (ensureOmemoDeviceTable())
        std::cerr << "JidDb::removeOmemoDevices: ensureOmemoDeviceTable failed"
            << std::endl;
    else
    {
        QSqlQuery q(db);

        if (!q.exec("delete from 'omemo/devices' "
            "where jid = '" + jid + "'"))
            std::cerr << "JidDb::removeOmemoDevices: query exec failed: "
                << qPrintable(q.lastError().text()) << std::endl;
    }
}

int JidDb::ensureOwnIdTable(const QString &encryption, QString &table)
    const
{
    QSqlQuery q(db);

    table = encryption + "/ownID";

    if (!q.exec("create table if not exists '" + table + "' "
        "(ID TEXT) strict;"))
    {
        std::cerr << "JidDb::ensureOwnIdTable: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return -1;
    }

    return 0;
}

QByteArray JidDb::ownKeyId(const QString &encryption) const
{
    QString table;

    if (ensureOwnIdTable(encryption, table))
    {
        std::cerr << "JidDb::store: ensureOwnIdTable failed" << std::endl;
        return QByteArray();
    }

    QSqlQuery q(db);

    if (!q.exec("select * from '" + table + "';"))
    {
        std::cerr << "JidDb::ownKeyId: query exec failed: "
            << qPrintable(q.lastError().text()) << std::endl;
        return QByteArray();
    }

    while (q.next())
    {
        const auto id = q.value("id").toString();

        return QByteArray::fromBase64(id.toLatin1());
    }

    return QByteArray();
}

int JidDb::store(const QString &encryption, const QByteArray &ownId) const
{
    QString table;

    if (ensureOwnIdTable(encryption, table))
    {
        std::cerr << "JidDb::store: ensureOmemoOwnDeviceTable failed"
            << std::endl;
        return -1;
    }

    QSqlQuery q(db);
    const auto id = ownId.isEmpty() ? "\"\""
        : "'" + QString::fromLatin1(ownId.toBase64()) + "'",
        query = "insert or replace into '" + table + "' "
            "(id) values ("
            + id
            + ");";

    return 0;
}

void JidDb::removeSignedPreKeyPair(const uint32_t id) const
{
    if (ensureOmemoSignedPreKeyTable())
        std::cerr << "JidDb::removeSignedPreKeyPair: "
            "ensureOmemoSignedPreKeyTable failed" << std::endl;
    else
    {
        QSqlQuery q(db);
        const auto idstr = QString::number(id);

        if (!q.exec("delete from 'omemo/signedprekeys' where id = " + idstr))
            std::cerr << "JidDb::removeSignedPreKeyPair: query exec failed: "
                << qPrintable(q.lastError().text()) << std::endl;
    }
}

void JidDb::removePreKeyPair(const uint32_t id) const
{
    if (ensureOmemoPreKeyTable())
        std::cerr << "JidDb::removePreKeyPair: "
            "ensureOmemoPreKeyTable failed" << std::endl;
    else
    {
        QSqlQuery q(db);
        const auto idstr = QString::number(id);

        if (!q.exec("delete from 'omemo/prekeys' where id = " + idstr))
            std::cerr << "JidDb::removePreKeyPair: query exec failed: "
                << qPrintable(q.lastError().text()) << std::endl;
    }
}

void JidDb::removeOmemo() const
{
    if (ensureOmemoSignedPreKeyTable())
        std::cerr << "JidDb::removeOmemo: "
            "ensureOmemoSignedPreKeyTable failed" << std::endl;
    else if (ensureOmemoPreKeyTable())
        std::cerr << "JidDb::removeOmemo: "
            "ensureOmemoPreKeyTable failed" << std::endl;
    else
    {
        QSqlQuery q(db);

        if (!q.exec("delete from 'omemo/signedprekeys'"))
            std::cerr << "JidDb::removeOmemo: query exec failed: "
                << qPrintable(q.lastError().text()) << std::endl;
        else
        {
            QSqlQuery q(db);

            if (!q.exec("delete from 'omemo/prekeys'"))
                std::cerr << "JidDb::removeOmemo: query exec failed: "
                    << qPrintable(q.lastError().text()) << std::endl;
        }
    }
}