#include "UpdateTransaction.h"
#include <QDateTime>
#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QJsonArray>
#include <QJsonDocument>
#include <QSaveFile>
#include <QSet>
#include <QUuid>

UpdateTransaction::UpdateTransaction(const QString& installDir, const QString& updateDir, const QString& fromVersion,
                                     const QString& toVersion, int manifestId)
    : m_installDir(QDir::cleanPath(installDir)),
      m_updateDir(QDir::cleanPath(updateDir.isEmpty() ? QDir(installDir).filePath("update") : updateDir)),
      m_stateFile(QDir(m_updateDir).filePath("upgrade_state.json")),
      m_transactionId(QUuid::createUuid().toString(QUuid::WithoutBraces)),
      m_fromVersion(fromVersion), m_toVersion(toVersion), m_manifestId(manifestId)
{
    m_stagingDir = QDir(m_updateDir).filePath("staging/" + m_transactionId);
    m_backupDir = QDir(m_updateDir).filePath("backup/" + m_transactionId);
}

bool UpdateTransaction::safeRelativePath(const QString& path)
{
    const QString clean = QDir::cleanPath(QDir::fromNativeSeparators(path));
    return !clean.isEmpty() && !QDir::isAbsolutePath(clean) && clean != ".."
        && !clean.startsWith("../") && !clean.contains(":");
}

bool UpdateTransaction::copyOverwrite(const QString& source, const QString& destination) const
{
    if (!QDir().mkpath(QFileInfo(destination).path())) return false;
    if (QFile::exists(destination) && !QFile::remove(destination)) return false;
    return QFile::copy(source, destination);
}

bool UpdateTransaction::writeState(const QString& status, const QString& errorCode, const QString& message)
{
    m_state["transaction_id"] = m_transactionId;
    m_state["from_version"] = m_fromVersion;
    m_state["to_version"] = m_toVersion;
    m_state["status"] = status;
    m_state["manifest_id"] = m_manifestId;
    m_state["updated_at"] = QDateTime::currentDateTimeUtc().toString(Qt::ISODate);
    m_state["staging_dir"] = m_stagingDir;
    m_state["backup_dir"] = m_backupDir;
    m_state["error_code"] = errorCode;
    m_state["message"] = message;
    QJsonArray paths;
    for (const QString& path : m_changedPaths) paths.append(path);
    m_state["changed_paths"] = paths;
    QJsonArray obsoletePaths;
    for (const QString& path : m_obsoletePaths) obsoletePaths.append(path);
    m_state["obsolete_paths"] = obsoletePaths;

    QDir().mkpath(m_updateDir);
    QSaveFile file(m_stateFile);
    if (!file.open(QIODevice::WriteOnly)) return false;
    const QByteArray payload = QJsonDocument(m_state).toJson(QJsonDocument::Indented);
    if (file.write(payload) != payload.size()) { file.cancelWriting(); return false; }
    return file.commit();
}

bool UpdateTransaction::initialize()
{
    QDir(m_stagingDir).removeRecursively();
    QDir(m_backupDir).removeRecursively();
    if (!QDir().mkpath(m_stagingDir) || !QDir().mkpath(m_backupDir)) return false;
    m_state["started_at"] = QDateTime::currentDateTimeUtc().toString(Qt::ISODate);
    return writeState("prepared");
}

bool UpdateTransaction::resumeExisting(QString* errorMessage)
{
    QFile file(m_stateFile);
    if (!file.open(QIODevice::ReadOnly)) {
        if (errorMessage) *errorMessage = "cannot open upgrade_state.json";
        return false;
    }
    const QJsonDocument doc = QJsonDocument::fromJson(file.readAll());
    file.close();
    if (!doc.isObject()) {
        if (errorMessage) *errorMessage = "invalid upgrade_state.json";
        return false;
    }
    m_state = doc.object();
    const QString expectedToVersion = m_toVersion;
    const int expectedManifestId = m_manifestId;
    const QString stateStatus = m_state.value("status").toString();
    m_transactionId = m_state.value("transaction_id").toString();
    m_fromVersion = m_state.value("from_version").toString();
    m_toVersion = m_state.value("to_version").toString();
    m_manifestId = m_state.value("manifest_id").toInt();
    if (m_toVersion != expectedToVersion || m_manifestId != expectedManifestId
        || (stateStatus != "awaiting_bootstrap" && stateStatus != "post_verify"
            && stateStatus != "rollback_required")) {
        if (errorMessage) *errorMessage = "upgrade state does not match resume request";
        return false;
    }
    m_stagingDir = m_state.value("staging_dir").toString();
    m_backupDir = m_state.value("backup_dir").toString();
    m_changedPaths.clear();
    m_obsoletePaths.clear();
    for (const QJsonValue& value : m_state.value("changed_paths").toArray()) {
        const QString path = value.toString();
        if (!safeRelativePath(path)) {
            if (errorMessage) *errorMessage = "unsafe path in upgrade state";
            return false;
        }
        m_changedPaths.append(path);
    }
    for (const QJsonValue& value : m_state.value("obsolete_paths").toArray()) {
        const QString path = value.toString();
        if (!safeRelativePath(path)) {
            if (errorMessage) *errorMessage = "unsafe obsolete path in upgrade state";
            return false;
        }
        m_obsoletePaths.append(path);
    }
    if (m_transactionId.isEmpty() || m_stagingDir.isEmpty() || m_backupDir.isEmpty()) {
        if (errorMessage) *errorMessage = "incomplete upgrade state";
        return false;
    }
    return true;
}

bool UpdateTransaction::recordVerifiedFiles(const QStringList& changedPaths,
                                                   const QStringList& obsoletePaths)
{
    m_changedPaths.clear();
    m_obsoletePaths.clear();
    QSet<QString> changedKeys;
    for (const QString& path : changedPaths) {
        if (!safeRelativePath(path)) return false;
        const QString normalized = QDir::fromNativeSeparators(path);
        changedKeys.insert(normalized.toCaseFolded());
        m_changedPaths.append(normalized);
    }
    for (const QString& path : obsoletePaths) {
        if (!safeRelativePath(path)) return false;
        const QString normalized = QDir::fromNativeSeparators(path);
        if (changedKeys.contains(normalized.toCaseFolded())) return false;
        m_obsoletePaths.append(normalized);
    }
    return writeState("verified");
}

bool UpdateTransaction::backupCurrentFiles()
{
    if (!writeState("waiting_mainapp_exit")) return false;
    QStringList paths = m_changedPaths;
    paths.append(m_obsoletePaths);
    for (const QString& path : paths) {
        const QString installed = QDir(m_installDir).filePath(path);
        if (!QFile::exists(installed)) continue;
        const QString backup = QDir(m_backupDir).filePath(path);
        if (!copyOverwrite(installed, backup))
            return markFailed("backup_failed", path);
    }
    return writeState("backed_up");
}

bool UpdateTransaction::installStagedFiles(QString* failedPath)
{
    if (!writeState("replacing")) return false;
    for (const QString& path : m_changedPaths) {
        const QString source = QDir(m_stagingDir).filePath(path);
        const QString destination = QDir(m_installDir).filePath(path);
        if (!copyOverwrite(source, destination)) {
            if (failedPath) *failedPath = path;
            writeState("rollback_required", "replace_failed", path);
            return false;
        }
    }
    return writeState("replaced");
}

bool UpdateTransaction::markAwaitingBootstrap() { return writeState("awaiting_bootstrap"); }

bool UpdateTransaction::markRollbackRequired(const QString& reason)
{ return writeState("rollback_required", "post_install_failed", reason); }

bool UpdateTransaction::markRolledBack() { return writeState("rolled_back"); }

bool UpdateTransaction::markPostVerify() { return writeState("post_verify"); }

bool UpdateTransaction::rollback(QString* failedPath)
{
    writeState("rolling_back");
    bool ok = true;
    QStringList paths = m_changedPaths;
    paths.append(m_obsoletePaths);
    for (const QString& path : paths) {
        const QString installed = QDir(m_installDir).filePath(path);
        const QString backup = QDir(m_backupDir).filePath(path);
        if (QFile::exists(backup)) {
            if (!copyOverwrite(backup, installed)) { ok = false; if (failedPath) *failedPath = path; }
        } else if (QFile::exists(installed) && !QFile::remove(installed)) {
            ok = false; if (failedPath) *failedPath = path;
        }
    }
    writeState(ok ? "rolled_back" : "failed", ok ? QString() : "rollback_failed",
               failedPath ? *failedPath : QString());
    return ok;
}

bool UpdateTransaction::commit()
{
    if (!writeState("committed")) return false;
    QDir(m_stagingDir).removeRecursively();
    QDir(m_backupDir).removeRecursively();
    return true;
}

bool UpdateTransaction::markFailed(const QString& errorCode, const QString& message)
{ return writeState("failed", errorCode, message); }

QString UpdateTransaction::transactionId() const { return m_transactionId; }
QString UpdateTransaction::fromVersion() const { return m_fromVersion; }
QString UpdateTransaction::stagingDir() const { return m_stagingDir; }
QString UpdateTransaction::backupDir() const { return m_backupDir; }
QStringList UpdateTransaction::obsoletePaths() const { return m_obsoletePaths; }
QString UpdateTransaction::healthFile() const { return QDir(m_updateDir).filePath("health_" + m_transactionId + ".ok"); }

bool UpdateTransaction::recoverInterrupted(const QString& installDir, const QString& updateDir,
                                           QString* restoredVersion, QString* errorMessage)
{
    QString stateFile = QDir(updateDir.isEmpty() ? QDir(installDir).filePath("update") : updateDir)
                            .filePath("upgrade_state.json");
    const QString legacyStateFile = QDir(installDir).filePath("update/upgrade_state.json");
    if (!QFile::exists(stateFile) && QFile::exists(legacyStateFile))
        stateFile = legacyStateFile;
    QFile file(stateFile);
    if (!file.exists()) return true;
    if (!file.open(QIODevice::ReadOnly)) { if (errorMessage) *errorMessage = "cannot open upgrade_state.json"; return false; }
    const QJsonDocument doc = QJsonDocument::fromJson(file.readAll());
    file.close();
    if (!doc.isObject()) { if (errorMessage) *errorMessage = "invalid upgrade_state.json"; return false; }
    QJsonObject state = doc.object();
    const QString status = state.value("status").toString();
    if (status == "committed") return true;
    if (status == "rolled_back") {
        if (restoredVersion) *restoredVersion = state.value("from_version").toString();
        return true;
    }
    const QString backupDir = state.value("backup_dir").toString();
    QJsonArray paths = state.value("changed_paths").toArray();
    for (const QJsonValue& value : state.value("obsolete_paths").toArray()) paths.append(value);
    const QString errorCode = state.value("error_code").toString();
    const bool mayContainNewFiles = status == "replacing" || status == "replaced"
        || status == "post_verify" || status == "rollback_required"
        || status == "awaiting_bootstrap" || status == "rolling_back" || errorCode == "rollback_failed";
    bool ok = true;
    for (const QJsonValue& value : paths) {
        const QString path = value.toString();
        if (!safeRelativePath(path)) { ok = false; continue; }
        const QString installed = QDir(installDir).filePath(path);
        const QString backup = QDir(backupDir).filePath(path);
        if (QFile::exists(backup)) {
            QDir().mkpath(QFileInfo(installed).path());
            if (QFile::exists(installed)) QFile::remove(installed);
            if (!QFile::copy(backup, installed)) ok = false;
        } else if (mayContainNewFiles) {
            if (QFile::exists(installed) && !QFile::remove(installed)) ok = false;
        }
    }
    if (restoredVersion) *restoredVersion = state.value("from_version").toString();
    state["status"] = ok ? "rolled_back" : "failed";
    QSaveFile output(stateFile);
    if (output.open(QIODevice::WriteOnly)) { output.write(QJsonDocument(state).toJson(QJsonDocument::Indented)); output.commit(); }
    if (!ok && errorMessage) *errorMessage = "failed to restore one or more files";
    return ok;
}