// SPDX-FileCopyrightText: 2023 Linus Jahn <lnj@kaidan.im>
// SPDX-FileCopyrightText: 2024 Filipe Azevedo <pasnox@gmail.com>
//
// SPDX-License-Identifier: LGPL-2.1-or-later

#include "QXmppAccountMigrationManager.h"

#include "QXmppClient.h"
#include "QXmppConstants_p.h"
#include "QXmppPromise.h"
#include "QXmppTask.h"
#include "QXmppUtils_p.h"

#include "Algorithms.h"
#include "StringLiterals.h"
#include "XmlWriter.h"

#include <QDomElement>

using namespace QXmpp;
using namespace QXmpp::Private;

// Utilities for account data extensions

struct XmlElementId {
    QString tagName;
    QString xmlns;

    static XmlElementId fromDom(const QDomElement &el) { return { el.tagName(), el.namespaceURI() }; }

    bool operator==(const XmlElementId &other) const
    {
        return tagName == other.tagName && xmlns == other.xmlns;
    }

    bool operator<(const XmlElementId &other) const
    {
        const auto result = xmlns.compare(other.xmlns);

        if (result == 0) {
            return tagName.compare(other.tagName) < 0;
        }

        return result < 0;
    }
};

#ifndef QXMPP_DOC
template<>
struct std::hash<XmlElementId> {
    std::size_t operator()(const XmlElementId &m) const noexcept
    {
        std::size_t h1 = std::hash<QString> {}(m.tagName);
        std::size_t h2 = std::hash<QString> {}(m.xmlns);
        return h1 ^ (h2 << 1);
    }
};
#endif

using AnyParser = QXmppExportData::ExtensionParser<std::any>;
using AnySerializer = QXmppExportData::ExtensionSerializer<std::any>;

static std::unordered_map<std::type_index, XmlElementId> &accountDataMapping()
{
    thread_local static std::unordered_map<std::type_index, XmlElementId> registry;
    return registry;
}

static std::unordered_map<XmlElementId, AnyParser> &accountDataParsers()
{
    thread_local static std::unordered_map<XmlElementId, AnyParser> registry;
    return registry;
}

static std::unordered_map<std::type_index, AnySerializer> &accountDataSerializers()
{
    thread_local static std::unordered_map<std::type_index, AnySerializer> registry;
    return registry;
}

struct QXmppExportDataPrivate : QSharedData {
    QString accountJid;
    std::unordered_map<std::type_index, std::any> extensions;
};

QXmppExportData::QXmppExportData()
    : d(new QXmppExportDataPrivate())
{
}

QXMPP_PRIVATE_DEFINE_RULE_OF_SIX(QXmppExportData)

std::variant<QXmppExportData, QXmppError> QXmppExportData::fromDom(const QDomElement &el)
{
    if (el.tagName() != u"account-data" || el.namespaceURI() != ns_qxmpp_export) {
        return QXmppError { u"Invalid XML document provided."_s, {} };
    }

    const auto &parsers = accountDataParsers();

    QXmppExportData data;
    data.setAccountJid(el.attribute(u"jid"_s));

    for (const auto &extension : iterChildElements(el)) {
        const auto parser = parsers.find(XmlElementId::fromDom(extension));
        if (parser != parsers.end()) {
            const auto &[_, parse] = *parser;

            auto result = parse(extension);
            if (hasError(result)) {
                return getError(std::move(result));
            }

            auto extensionValue = getValue(std::move(result));
            data.d->extensions.emplace(extensionValue.type(), std::move(extensionValue));
        }
    }

    return data;
}

void QXmppExportData::toXml(QXmlStreamWriter *writer) const
{
    // We need to generate the xml file with nodes always in the same order.
    // This is needed for our unit tests which are based on xml generation.
    const auto sortedExtensionsKeys = [this]() {
        auto keys = transform<std::vector<std::pair<std::type_index, XmlElementId>>>(accountDataMapping(), [](auto pair) {
            return pair;
        });
        std::ranges::stable_sort(keys, [](const auto &left, const auto &right) {
            return left.second < right.second;
        });
        return keys;
    }();

    XmlWriter w(writer);

    writer->writeStartDocument();
    w.write(Element {
        { u"account-data", ns_qxmpp_export },
        Attribute { u"jid", d->accountJid },
        [&] {
            const auto &serializers = accountDataSerializers();
            for (const auto &sortedKey : sortedExtensionsKeys) {
                const auto &typeIndex = sortedKey.first;

                if (!d->extensions.contains(typeIndex)) {
                    continue;
                }

                const auto serializer = serializers.find(typeIndex);
                if (serializer != serializers.end()) {
                    const auto &[_, serialize] = *serializer;

                    serialize(d->extensions.at(typeIndex), *writer);
                }
            }
        },
    });
    writer->writeEndDocument();
}

const QString &QXmppExportData::accountJid() const
{
    return d->accountJid;
}

void QXmppExportData::setAccountJid(const QString &jid)
{
    d->accountJid = jid;
}

const std::unordered_map<std::type_index, std::any> &QXmppExportData::extensions() const
{
    return d->extensions;
}

void QXmppExportData::setExtension(std::any value)
{
    d->extensions.emplace(std::type_index(value.type()), std::move(value));
}

void QXmppExportData::registerExtensionInternal(std::type_index type, ExtensionParser<std::any> parse, ExtensionSerializer<std::any> serialize, QStringView tagName, QStringView xmlns)
{
    const auto id = XmlElementId { tagName.toString(), xmlns.toString() };
    accountDataMapping().emplace(type, id);
    accountDataParsers().emplace(id, parse);
    accountDataSerializers().emplace(type, serialize);
}

bool QXmppExportData::isExtensionRegistered(std::type_index type)
{
    return accountDataSerializers().contains(type);
}

struct QXmppAccountMigrationManagerPrivate {
    template<typename T = Success>
    using Result = std::variant<T, QXmppError>;

    struct ExtensionData {
        std::function<QXmppTask<Result<>>(std::any)> importFunction;
        std::function<QXmppTask<Result<std::any>>()> exportFunction;
    };

    std::unordered_map<std::type_index, ExtensionData> extensions;
};

/*!
    \class QXmppAccountMigrationManager
    \inmodule QXmpp

    \brief Allows to export and import account data.

    This manager helps migrating a user account into another server.

    You can use exportData() to start a data export. Afterwards you can use the exported data to
    start a data import on another account using importData().

    The data that is exported (or imported) is determined by the other registered client
    extensions. They can register callbacks for export and import using registerExtension().

    \ingroup Managers
    \since QXmpp 1.8
*/

/*!
    \typealias QXmppAccountMigrationManager::Result

    Contains T or QXmppError.
*/

/*!
    \fn QXmppAccountMigrationManager::errorOccurred(const QXmppError &error)

    Emitted when an error occured during export or import. \a error is the occured error.
*/

/*!
    \fn template <typename DataType, typename ImportFunc, typename ExportFunc> void QXmppAccountMigrationManager::registerExportData(ImportFunc importFunc, ExportFunc exportFunc)

    Registers a data type that can be imported to an account using \a importFunc and generated using
    \a exportFunc.

    The functions are used when importData() or exportData() is called. You can unregister the
    functions using unregisterExportData().

    The data type MUST also be registered via QXmppExportData::registerExtension() so it can be
    serialized.
*/

/*!
    \fn template <typename DataType> void QXmppAccountMigrationManager::unregisterExportData()

    Unregisters a data type.
*/

/*!
    Constructs an account migration manager.

    \note You would need the `QXmppClient(QXmppClient::NoExtensions, this)` approach to use this manager
    because it needs to be instantiated before others using it.
*/
QXmppAccountMigrationManager::QXmppAccountMigrationManager()
    : d(std::make_unique<QXmppAccountMigrationManagerPrivate>())
{
}

QXmppAccountMigrationManager::~QXmppAccountMigrationManager() = default;

/*! Imports QXmppExportData to the currently connected \a account. */
QXmppTask<QXmppAccountMigrationManager::Result<>> QXmppAccountMigrationManager::importData(const QXmppExportData &account)
{
    if (account.extensions().empty()) {
        return makeReadyTask<Result<>>({});
    }

    auto promise = std::make_shared<QXmppPromise<Result<>>>();
    auto task = promise->task();

    auto counter = std::make_shared<int>(account.extensions().size());

    for (const auto &[typeIndex, value] : account.extensions()) {
        auto extensionType = d->extensions.find(typeIndex);

        // There is no registered extension to import this data type
        if (extensionType == d->extensions.cend()) {
            if (--(*counter) == 0) {
                promise->finish(Success());
            }
            continue;
        }

        auto &[_, extension] = *extensionType;
        extension.importFunction(value).then(this, [promise, counter](auto &&result) mutable {
            if (promise->task().isFinished()) {
                return;
            }

            if (hasError(result)) {
                return promise->finish(getError(std::move(result)));
            }

            if ((--(*counter)) == 0) {
                promise->finish(Success());
            }
        });
    }

    return task;
}

/*! Creates a data export of the current account. */
QXmppTask<QXmppAccountMigrationManager::Result<QXmppExportData>> QXmppAccountMigrationManager::exportData()
{
    struct State {
        QXmppPromise<Result<QXmppExportData>> p;
        QXmppExportData data;
        size_t counter = 0;
    };

    auto state = std::make_shared<State>();
    state->data.setAccountJid(client()->configuration().jidBare());
    state->counter = d->extensions.size();

    // early exit
    if (d->extensions.empty()) {
        return makeReadyTask<Result<QXmppExportData>>(std::move(state->data));
    }

    for (const auto &[dataTypeIndex, extension] : std::as_const(d->extensions)) {
        extension.exportFunction().then(this, [state](auto &&result) mutable {
            auto &[p, data, counter] = *state;

            if (p.task().isFinished()) {
                return;
            }

            if (hasError(result)) {
                return p.finish(getError(std::move(result)));
            }

            data.setExtension(getValue(std::move(result)));

            if (--counter == 0) {
                p.finish(std::move(data));
            }
        });
    }

    return state->p.task();
}

void QXmppAccountMigrationManager::registerMigrationDataInternal(
    std::type_index dataType,
    std::function<QXmppTask<Result<>>(std::any)> importFunc,
    std::function<QXmppTask<Result<std::any>>()> exportFunc)
{
    d->extensions.emplace(dataType, QXmppAccountMigrationManagerPrivate::ExtensionData { importFunc, exportFunc });
}

void QXmppAccountMigrationManager::unregisterMigrationDataInternal(std::type_index dataType)
{
    d->extensions.erase(dataType);
}
