#include <QApplication>
#include <QCoreApplication>
#include <QDebug>
#include <QDir>
#include <QDirIterator>
#include <QElapsedTimer>
#include <QFile>
#include <QMessageBox>
#include <QProcess>
#include <QProgressDialog>
#include <QSaveFile>
#include <QStorageInfo>
#include <QThread>
#include "UpdaterLogic.h"
#include "UpdateTransaction.h"
#include "FileHelper.h"
#include "ConfigHelper.h"
#include "TicketHelper.h"
#include "PolicyHelper.h"
#include "TranslationHelper.h"
#include <QVersionNumber>

int main(int argc, char* argv[])
{
    QApplication app(argc, argv);
    QApplication::setApplicationName("Marsco Updater");

    TranslationHelper translation;
    translation.installForSystemLocale();

    UpdaterLogic logic;
    QString offlinePackagePath;
    QString appId;
    QString channel;
    QString targetVersion;
    int targetVersionId = 0;
    if (argc >= 2 && QString(argv[1]).startsWith("--offline-package=")) {
        offlinePackagePath = QString(argv[1]).mid(QString("--offline-package=").size());
        if (!logic.loadOfflinePackage(offlinePackagePath)) {
            QMessageBox::critical(nullptr, QCoreApplication::translate("Updater", "Invalid Offline Package"), logic.offlineError());
            return -1;
        }
        const QJsonObject packageManifest = logic.getManifest();
        appId = packageManifest.value("app_id").toString();
        channel = packageManifest.value("channel").toString();
        targetVersion = packageManifest.value("version").toString();
        targetVersionId = packageManifest.value("manifest_seq").toInt();
    } else {
        if (argc < 5) {
            QMessageBox::critical(
                nullptr,
                QCoreApplication::translate("Updater", "Invalid Updater Arguments"),
                QCoreApplication::translate("Updater", "The updater requires online update arguments or an offline update package."));
            return -1;
        }
        appId = argv[1]; channel = argv[2]; targetVersion = argv[3]; targetVersionId = QString(argv[4]).toInt();
    }
    QString bootstrapResult;
    for (int i = 5; i < argc; ++i) {
        const QString arg = argv[i];
        if (arg.startsWith("--bootstrap-resume="))
            bootstrapResult = arg.mid(QString("--bootstrap-resume=").size());
    }
    const bool resumingFromBootstrap = !bootstrapResult.isEmpty();
    const QString runtimeDir = QApplication::applicationDirPath();

    ConfigHelper& config = ConfigHelper::instance();
    const QString targetDir = config.installRoot();
    const QString updateDir = config.updateRoot();
    QDir().mkpath(updateDir);
    if (appId != config.getValue("App", "app_id") || channel != config.getValue("App", "channel")) {
        QMessageBox::critical(
            nullptr,
            QCoreApplication::translate("Updater", "Offline Package Not Applicable"),
            QCoreApplication::translate("Updater", "The update package application or channel does not match this installation."));
        return -1;
    }
    QString fromVersion = config.getValue("App", "current_version");
    PolicyHelper updatePolicy(runtimeDir);
    const bool policyApplicable =
        updatePolicy.loadPolicy() && updatePolicy.isValid()
        && updatePolicy.isApplicable(appId, channel, fromVersion)
        && !updatePolicy.isExpired()
        && updatePolicy.isVersionAllowed(targetVersion);
    if (!policyApplicable
        || (!offlinePackagePath.isEmpty() && !updatePolicy.isOfflineAllowed()))
    {
        QMessageBox::critical(
            nullptr,
            QCoreApplication::translate("Updater", "Update Denied"),
            QCoreApplication::translate(
                "Updater",
                "The local signed policy is invalid, expired, not applicable to this installation, or does not allow the requested update."));
        return -1;
    }
    if (QVersionNumber::compare(QVersionNumber::fromString(targetVersion),
                                QVersionNumber::fromString(fromVersion)) < 0
        && !updatePolicy.allowRollback())
    {
        QMessageBox::critical(
            nullptr,
            QCoreApplication::translate("Updater", "Downgrade Prohibited"),
            QCoreApplication::translate("Updater", "The current signed policy does not allow installing an older version."));
        return -1;
    }
    QString deviceId = config.getValue("Update", "device_id");
    if (deviceId.isEmpty()) deviceId = "unknown_device";

    if (!resumingFromBootstrap) {
        QString restoredVersion;
        QString recoveryError;
        if (!UpdateTransaction::recoverInterrupted(targetDir, updateDir, &restoredVersion, &recoveryError)) {
            QMessageBox::critical(
                nullptr,
                QCoreApplication::translate("Updater", "Update Recovery Failed"),
                QCoreApplication::translate("Updater", "The previous update did not finish, and the old version could not be restored: %1\nDo not continue running the software. Contact an administrator.")
                    .arg(recoveryError));
            return -1;
        }
        if (!restoredVersion.isEmpty() && restoredVersion != fromVersion) {
            if (!config.setValue("App", "current_version", restoredVersion)) {
                QMessageBox::critical(
                    nullptr,
                    QCoreApplication::translate("Updater", "Update Recovery Failed"),
                    QCoreApplication::translate("Updater", "The old files were restored, but the version state could not be restored. Check access to the system state store."));
                return -1;
            }
            fromVersion = restoredVersion;
        }
    }

    QProgressDialog progress(QCoreApplication::translate("Updater", "Preparing the update..."), QString(), 0, 100);
    progress.setWindowTitle(QCoreApplication::translate("Updater", "Updating to %1").arg(targetVersion));
    progress.setCancelButton(nullptr);
    progress.setWindowModality(Qt::ApplicationModal);
    progress.setMinimumDuration(0);
    progress.setAutoClose(false);
    progress.setValue(2);
    progress.show();
    QApplication::processEvents();

    UpdateTransaction transaction(targetDir, updateDir, fromVersion, targetVersion, targetVersionId);
    const auto setProgress = [&](int value, const QString& message) {
        progress.setValue(value);
        progress.setLabelText(message);
        QApplication::processEvents();
    };
    const auto formatBytes = [](qint64 bytes) {
        const double value = static_cast<double>(bytes);
        if (bytes >= 1024LL * 1024 * 1024)
            return QString::number(value / (1024.0 * 1024 * 1024), 'f', 2) + " GB";
        if (bytes >= 1024LL * 1024)
            return QString::number(value / (1024.0 * 1024), 'f', 2) + " MB";
        if (bytes >= 1024)
            return QString::number(value / 1024.0, 'f', 1) + " KB";
        return QString::number(bytes) + " B";
    };
    QObject::connect(&logic, &UpdaterLogic::downloadProgress,
        [&](qint64 received, qint64 total, const QString& path, double bytesPerSecond) {
            const int value = total > 0 ? 30 + static_cast<int>(30 * received / total) : 60;
            const QString fileText = path.isEmpty()
                                         ? QCoreApplication::translate("Updater", "Preparing the download...")
                                         : QCoreApplication::translate("Updater", "Downloading: %1").arg(path);
            progress.setValue(value);
            progress.setLabelText(QString("%1\n%2 / %3 | %4/s")
                .arg(fileText, formatBytes(received), formatBytes(total),
                     formatBytes(static_cast<qint64>(bytesPerSecond))));
            QApplication::processEvents();
        });
    const auto reportUpdateResult = [&](bool success) {
        if (offlinePackagePath.isEmpty())
            logic.reportResult(deviceId, fromVersion, targetVersion, success);
    };
    const auto fail = [&](const QString& title, const QString& message,
                          const QString& errorCode = QString("update_failed")) {
        transaction.markFailed(errorCode, message);
        reportUpdateResult(false);
        progress.close();
        QMessageBox::critical(nullptr, title, message);
        return -1;
    };
    const auto configuredName = [&](const QString& key, const QString& fallback) {
        const QString value = config.getValue("Runtime", key).trimmed();
        return value.isEmpty() ? fallback : value;
    };
    const QString mainExecutable = configuredName("main_executable", "MainApp.exe");
    const QString updaterExecutable = configuredName("updater_executable", "Updater.exe");
    const QString bootstrapExecutable = configuredName("bootstrap_executable", "Bootstrap.exe");
    bool timeoutOk = false;
    int healthCheckTimeoutMs = config.getValue("Runtime", "health_check_timeout_ms").toInt(&timeoutOk);
    if (!timeoutOk || healthCheckTimeoutMs < 1000) healthCheckTimeoutMs = 15000;
    const QString mainAppPath = QDir(runtimeDir).filePath(mainExecutable);
    const QString updaterPath = QDir(runtimeDir).filePath(updaterExecutable);
    const QString bootstrapPath = QDir(runtimeDir).filePath(bootstrapExecutable);
    const QString launchToken = config.getValue("App", "launch_token");
    const auto launchMainApp = [&](const QString& healthFile = QString()) {
        QString ticketPath;
        QString ticketError;
        const QString launchVersion = config.getValue("App", "current_version");
        if (!TicketHelper::createTicket(appId, deviceId, launchVersion, launchToken,
                                        &ticketPath, &ticketError)) {
            qDebug() << "Cannot create launch ticket:" << ticketError;
            return false;
        }
        QStringList args{QString("--ticket-file=%1").arg(ticketPath)};
        if (!healthFile.isEmpty()) args.append(QString("--health-file=%1").arg(healthFile));
        const bool started = QProcess::startDetached(mainAppPath, args);
        if (!started) QFile::remove(ticketPath);
        return started;
    };
    const auto bootstrapPlanFile = [&]() {
        return QDir(updateDir).filePath("bootstrap_plan_" + transaction.transactionId() + ".txt");
    };
    const auto launchBootstrap = [&](const QString& mode) {
        const QStringList args{
            bootstrapPlanFile(), targetDir, transaction.stagingDir(), transaction.backupDir(), updaterPath,
            QString::number(QCoreApplication::applicationPid()), appId, channel,
            targetVersion, QString::number(targetVersionId), mode
        };
        return QFile::exists(bootstrapPath) && QProcess::startDetached(bootstrapPath, args);
    };
    const auto delegateRollback = [&](const QString& title, const QString& reason) {
        FileHelper::killProcess(mainExecutable);
        transaction.markRollbackRequired(reason);
        progress.setLabelText(QCoreApplication::translate("Updater", "Handing rollback over to Bootstrap..."));
        QApplication::processEvents();
        if (launchBootstrap("rollback")) {
            progress.close();
            return 0;
        }
        reportUpdateResult(false);
        progress.close();
        QMessageBox::critical(nullptr, title,
            reason + QCoreApplication::translate("Updater", "\n\nBootstrap could not be started to perform the rollback. Do not continue running the software. Contact an administrator."));
        return -1;
    };

    if (resumingFromBootstrap) {
        QString resumeError;
        if (!transaction.resumeExisting(&resumeError)) {
            reportUpdateResult(false);
            progress.close();
            QMessageBox::critical(
                nullptr,
                QCoreApplication::translate("Updater", "Transaction Resume Failed"),
                QCoreApplication::translate("Updater", "Cannot read the Bootstrap update transaction: %1").arg(resumeError));
            return -1;
        }
        fromVersion = transaction.fromVersion();
        if (QFile::exists(QDir(transaction.backupDir()).filePath(".offline_mode")))
            offlinePackagePath = "__bootstrap_resumed_offline__";
        if (bootstrapResult == "rolledback") {
            const bool stateOk = config.setValue("App", "current_version", fromVersion);
            transaction.markRolledBack();
            reportUpdateResult(false);
            if (stateOk) launchMainApp();
            progress.close();
            QMessageBox::warning(
                nullptr,
                QCoreApplication::translate("Updater", "Update Rolled Back"),
                stateOk
                    ? QCoreApplication::translate("Updater", "The new version could not be installed or started. The old version was restored and started automatically.")
                    : QCoreApplication::translate("Updater", "The old files were restored, but the old version number could not be written back. Check permissions on the system state store."));
            return stateOk ? 0 : -1;
        }
        if (bootstrapResult != "success") {
            reportUpdateResult(false);
            progress.close();
            QMessageBox::critical(
                nullptr,
                QCoreApplication::translate("Updater", "Automatic Rollback Failed"),
                QCoreApplication::translate("Updater", "Bootstrap could not fully restore the old version. Do not continue running the software. Contact an administrator."));
            return -1;
        }
    } else if (!transaction.initialize()) {
        return fail(
            QCoreApplication::translate("Updater", "Update Preparation Failed"),
            QCoreApplication::translate("Updater", "Cannot create the update transaction directory or save its state."),
            "transaction_init_failed");
    } else if (!offlinePackagePath.isEmpty()) {
        QFile marker(QDir(transaction.backupDir()).filePath(".offline_mode"));
        if (!marker.open(QIODevice::WriteOnly) || marker.write("offline\n") != 8)
            return fail(
                QCoreApplication::translate("Updater", "Update Preparation Failed"),
                QCoreApplication::translate("Updater", "Cannot save the offline transaction marker."),
                "offline_marker_failed");
    }

    setProgress(
        resumingFromBootstrap ? 72 : 10,
        offlinePackagePath.isEmpty()
            ? QCoreApplication::translate("Updater", "Fetching and validating the version manifest...")
            : QCoreApplication::translate("Updater", "Validating the offline update package..."));
    const QString manifestCacheDir = QDir(updateDir).filePath("manifest_cache");
    if (resumingFromBootstrap) {
        if (!logic.loadManifestCache(manifestCacheDir, targetVersion))
            return delegateRollback(
                QCoreApplication::translate("Updater", "Manifest Cache Failed"),
                QCoreApplication::translate("Updater", "The signed manifest cache could not be read after Bootstrap installation."));
    } else if (offlinePackagePath.isEmpty()) {
        logic.getManifest(appId, channel, targetVersion, targetVersionId);
    }
    if (!logic.verifyManifestSignature()) {
        if (resumingFromBootstrap)
            return delegateRollback(
                QCoreApplication::translate("Updater", "Security Validation Failed"),
                QCoreApplication::translate("Updater", "The version manifest signature could not be revalidated after Bootstrap installation."));
        return fail(
            QCoreApplication::translate("Updater", "Security Validation Failed"),
            QCoreApplication::translate("Updater", "The version manifest signature is invalid. The update has stopped. Contact an administrator."),
            "manifest_signature_invalid");
    }
    QStringList obsoletePaths;
    if (!resumingFromBootstrap && fromVersion != targetVersion) {
        UpdaterLogic oldManifestLogic;
        if (oldManifestLogic.loadManifestCache(manifestCacheDir, fromVersion)
            && oldManifestLogic.getManifest().value("version").toString() == fromVersion
            && oldManifestLogic.verifyManifestSignature()) {
            obsoletePaths = logic.obsoleteFilesComparedTo(oldManifestLogic.getManifest());
            qDebug() << "Obsolete files from signed previous manifest:" << obsoletePaths;
        } else {
            qDebug() << "No valid signed previous manifest cache; obsolete deletion skipped for safety";
        }
    }
    if (!logic.saveManifestCache(manifestCacheDir)) {
        if (resumingFromBootstrap)
            return delegateRollback(
                QCoreApplication::translate("Updater", "Manifest Cache Failed"),
                QCoreApplication::translate("Updater", "Cannot save the new version manifest cache."));
        return fail(
            QCoreApplication::translate("Updater", "Manifest Cache Failed"),
            QCoreApplication::translate("Updater", "Cannot save the new version manifest cache. The update has stopped."),
            "manifest_cache_failed");
    }

    if (!resumingFromBootstrap) {
        setProgress(
            25,
            offlinePackagePath.isEmpty()
                ? QCoreApplication::translate("Updater", "Fetching secure download URLs...")
                : QCoreApplication::translate("Updater", "Preparing offline package files..."));
        if (offlinePackagePath.isEmpty())
            logic.getDownloadUrl(appId, channel, targetVersion, targetVersionId);
        if (logic.getFileList().isEmpty())
            return fail(
                QCoreApplication::translate("Updater", "No Update Files"),
                QCoreApplication::translate("Updater", "The server returned no version files. The update has stopped."),
                "empty_file_list");

        QStorageInfo storage(updateDir);
        storage.refresh();
        const qint64 requiredBytes = logic.estimateAdditionalDiskBytes(targetDir, obsoletePaths);
        const qint64 availableBytes = storage.bytesAvailable();
        if (!storage.isValid() || !storage.isReady() || availableBytes < requiredBytes) {
            return fail(
                QCoreApplication::translate("Updater", "Insufficient Disk Space"),
                QCoreApplication::translate("Updater", "The update requires at least %1 of free space, but the installation drive has only %2 available.\nThe required space includes downloaded files, the old-version backup, and a safety margin.")
                    .arg(formatBytes(requiredBytes), formatBytes(qMax<qint64>(0, availableBytes))),
                "disk_space_insufficient");
        }

        const QString stagingDir = transaction.stagingDir();
        setProgress(
            30,
            (offlinePackagePath.isEmpty()
                 ? QCoreApplication::translate("Updater", "Downloading and validating %1 version files...")
                 : QCoreApplication::translate("Updater", "Extracting and validating %1 offline files..."))
                .arg(logic.getFileList().size()));
        if (!offlinePackagePath.isEmpty()) {
            if (!logic.loadOfflinePackage(offlinePackagePath, stagingDir))
                return fail(QCoreApplication::translate("Updater", "Offline Package Extraction Failed"), logic.offlineError(), "offline_package_invalid");
        } else {
            if (!logic.downloadAllFiles(stagingDir, targetDir)) {
                logic.reportDownloadResult(appId, channel, targetVersion, false);
                return fail(
                    QCoreApplication::translate("Updater", "Download Failed"),
                    QCoreApplication::translate("Updater", "Some files could not be downloaded or failed SHA-256 validation. Check the network and try again."),
                    "download_failed");
            }
            logic.reportDownloadResult(appId, channel, targetVersion, true);
        }

        setProgress(58, QCoreApplication::translate("Updater", "Validating the complete version files..."));
        if (!logic.validateLocalFiles(stagingDir, targetDir))
            return fail(
                QCoreApplication::translate("Updater", "File Validation Failed"),
                QCoreApplication::translate("Updater", "The staged files do not match the version manifest. The update has stopped."),
                "staging_verify_failed");

        QStringList changedPaths;
        QDir stagingRoot(stagingDir);
        QDirIterator stagingFiles(stagingDir, QDir::Files, QDirIterator::Subdirectories);
        while (stagingFiles.hasNext())
            changedPaths.append(QDir::fromNativeSeparators(stagingRoot.relativeFilePath(stagingFiles.next())));
        const QString runtimePrefix = config.runtimeRelativePath();
        const QString bootstrapManifestPath = QDir::fromNativeSeparators(
            runtimePrefix.isEmpty() ? bootstrapExecutable : runtimePrefix + "/" + bootstrapExecutable);
        for (const QString& path : changedPaths) {
            if (path.compare(bootstrapManifestPath, Qt::CaseInsensitive) == 0)
                return fail(
                    QCoreApplication::translate("Updater", "Bootstrap Cannot Update Itself"),
                    QCoreApplication::translate("Updater", "This version contains a new %1. Upgrade this component with the installer, then republish the application version.").arg(bootstrapManifestPath),
                    "bootstrap_self_update_blocked");
        }
        if (!transaction.recordVerifiedFiles(changedPaths, obsoletePaths))
            return fail(
                QCoreApplication::translate("Updater", "Transaction Recording Failed"),
                QCoreApplication::translate("Updater", "Cannot save the validated-file or pending-deletion lists. The update has stopped."),
                "transaction_record_failed");

        setProgress(66, QCoreApplication::translate("Updater", "Closing the main application..."));
        if (!FileHelper::killProcess(mainExecutable))
            return fail(
                QCoreApplication::translate("Updater", "Cannot Close Main Application"),
                QCoreApplication::translate("Updater", "%1 is still running. Close it manually and try again.").arg(mainExecutable),
                "mainapp_close_failed");

        setProgress(72, QCoreApplication::translate("Updater", "Backing up %1 files to be changed, including %2 deletions...")
            .arg(changedPaths.size() + obsoletePaths.size()).arg(obsoletePaths.size()));
        if (!transaction.backupCurrentFiles())
            return fail(
                QCoreApplication::translate("Updater", "Backup Failed"),
                QCoreApplication::translate("Updater", "Cannot back up the current version files. The new version has not been installed. Check disk space and directory permissions."),
                "backup_failed");

        const QString planFile = bootstrapPlanFile();
        QSaveFile plan(planFile);
        if (!plan.open(QIODevice::WriteOnly))
            return fail(
                QCoreApplication::translate("Updater", "Handoff Preparation Failed"),
                QCoreApplication::translate("Updater", "Cannot create the Bootstrap file plan."),
                "bootstrap_plan_failed");
        const auto writePlanItem = [&](char operation, const QString& path) {
            const QByteArray line = QByteArray(1, operation) + '\t' + path.toUtf8() + '\n';
            return plan.write(line) == line.size();
        };
        for (const QString& path : changedPaths) {
            if (!writePlanItem('C', path)) {
                plan.cancelWriting();
                return fail(
                    QCoreApplication::translate("Updater", "Handoff Preparation Failed"),
                    QCoreApplication::translate("Updater", "Cannot write the Bootstrap copy plan."),
                    "bootstrap_plan_failed");
            }
        }
        for (const QString& path : obsoletePaths) {
            if (!writePlanItem('D', path)) {
                plan.cancelWriting();
                return fail(
                    QCoreApplication::translate("Updater", "Handoff Preparation Failed"),
                    QCoreApplication::translate("Updater", "Cannot write the Bootstrap deletion plan."),
                    "bootstrap_plan_failed");
            }
        }
        if (!plan.commit() || !transaction.markAwaitingBootstrap())
            return fail(
                QCoreApplication::translate("Updater", "Handoff Preparation Failed"),
                QCoreApplication::translate("Updater", "Cannot commit the Bootstrap file plan or transaction state."),
                "bootstrap_plan_failed");

        setProgress(78, QCoreApplication::translate("Updater", "Handing installation over to Bootstrap..."));
        if (!launchBootstrap("install"))
            return fail(
                QCoreApplication::translate("Updater", "Bootstrap Startup Failed"),
                QCoreApplication::translate("Updater", "Cannot start the independent update handoff program: %1").arg(bootstrapPath),
                "bootstrap_start_failed");
        progress.close();
        return 0;
    }

    setProgress(82, QCoreApplication::translate("Updater", "Validating the Bootstrap installation result..."));
    if (!transaction.markPostVerify() || !logic.validateLocalFiles(targetDir))
        return delegateRollback(
            QCoreApplication::translate("Updater", "Installation Validation Failed"),
            QCoreApplication::translate("Updater", "The new version files failed validation after installation."));
    for (const QString& path : transaction.obsoletePaths()) {
        if (QFile::exists(QDir(targetDir).filePath(path)))
            return delegateRollback(
                QCoreApplication::translate("Updater", "Obsolete File Cleanup Failed"),
                QCoreApplication::translate("Updater", "An obsolete file still exists: %1").arg(path));
    }

    setProgress(89, QCoreApplication::translate("Updater", "Saving the new version state..."));
    if (!config.setValue("App", "current_version", targetVersion)
        || config.getValue("App", "current_version") != targetVersion)
        return delegateRollback(
            QCoreApplication::translate("Updater", "State Save Failed"),
            QCoreApplication::translate("Updater", "Cannot save the current version number."));

    const QString healthFile = transaction.healthFile();
    QFile::remove(healthFile);
    setProgress(94, QCoreApplication::translate("Updater", "Starting the new version and waiting for a health confirmation..."));
    if (!launchMainApp(healthFile))
        return delegateRollback(
            QCoreApplication::translate("Updater", "Startup Failed"),
            QCoreApplication::translate("Updater", "%1 could not be started.").arg(mainExecutable));

    QElapsedTimer healthTimer;
    healthTimer.start();
    while (healthTimer.elapsed() < healthCheckTimeoutMs && !QFile::exists(healthFile)) {
        QApplication::processEvents();
        QThread::msleep(100);
    }
    if (!QFile::exists(healthFile))
        return delegateRollback(
            QCoreApplication::translate("Updater", "Startup Confirmation Failed"),
            QCoreApplication::translate("Updater", "The new version did not complete its startup health confirmation within %1 milliseconds.").arg(healthCheckTimeoutMs));

    setProgress(99, QCoreApplication::translate("Updater", "Committing the update transaction..."));
    if (!transaction.commit())
        return delegateRollback(
            QCoreApplication::translate("Updater", "Transaction Commit Failed"),
            QCoreApplication::translate("Updater", "The new version started, but the update transaction could not be committed."));

    QFile::remove(healthFile);
    QFile::remove(bootstrapPlanFile());
    reportUpdateResult(true);
    progress.setValue(100);
    progress.close();
    QMessageBox::information(
        nullptr,
        QCoreApplication::translate("Updater", "Update Complete"),
        QCoreApplication::translate("Updater", "The software was updated successfully to %1 and passed the startup health check.").arg(targetVersion));
    return 0;
}