// DragonX Wallet - ImGui Edition
// Copyright 2024-2026 The Hush Developers
// Released under the GPLv3

#include "bootstrap.h"
#include "../daemon/embedded_daemon.h"

#include <curl/curl.h>
#include <miniz.h>
#include <sodium.h>

#include <cstdio>
#include <cstring>
#include <filesystem>
#include <algorithm>
#include <sstream>
#include <iomanip>
#include "../util/logger.h"

namespace fs = std::filesystem;

namespace dragonx {
namespace util {

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

static bool endsWith(const std::string& s, const std::string& suffix) {
    if (suffix.size() > s.size()) return false;
    return s.compare(s.size() - suffix.size(), suffix.size(), suffix) == 0;
}

static size_t writeFileCallback(void* contents, size_t size, size_t nmemb, void* userp) {
    size_t total = size * nmemb;
    FILE* fp = static_cast<FILE*>(userp);
    return fwrite(contents, 1, total, fp);
}

std::string Bootstrap::formatSize(double bytes) {
    const char* units[] = { "B", "KB", "MB", "GB", "TB" };
    int idx = 0;
    double v = bytes;
    while (v >= 1024.0 && idx < 4) {
        v /= 1024.0;
        idx++;
    }
    char buf[64];
    if (idx == 0)
        snprintf(buf, sizeof(buf), "%.0f %s", v, units[idx]);
    else
        snprintf(buf, sizeof(buf), "%.2f %s", v, units[idx]);
    return buf;
}

// ---------------------------------------------------------------------------
// Lifecycle
// ---------------------------------------------------------------------------

Bootstrap::~Bootstrap() {
    cancel();
    if (worker_.joinable()) worker_.join();
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

void Bootstrap::start(const std::string& dataDir, const std::string& url) {
    if (worker_running_) return; // already running
    if (worker_.joinable()) worker_.join();

    cancel_requested_ = false;
    worker_running_ = true;

    // Ensure data dir exists
    fs::create_directories(dataDir);

    worker_ = std::thread([this, dataDir, url]() {
        std::string zipPath = dataDir + "/bootstrap_tmp.zip";

        // Step 1: Download (with resume support)
        // Check for a partial file from a previous interrupted download.
        {
            std::error_code ec;
            if (fs::exists(zipPath, ec) && fs::file_size(zipPath, ec) > 0 && !ec) {
                auto partial = fs::file_size(zipPath, ec);
                DEBUG_LOGF("[Bootstrap] Found partial download: %s (%s)\n",
                       zipPath.c_str(), formatSize((double)partial).c_str());
                setProgress(State::Downloading,
                            "Resuming download (" + formatSize((double)partial) + " already on disk)...");
            } else {
                setProgress(State::Downloading, "Connecting to dragonx.is...");
            }
        }
        if (!download(url, zipPath)) {
            if (cancel_requested_)
                setProgress(State::Failed, "Download cancelled");
            else
                setProgress(State::Failed, "Download failed — check your internet connection");
            // Keep the partial file so the next attempt can resume.
            // Only delete if the file is empty / doesn't exist.
            std::error_code ec;
            auto sz = fs::exists(zipPath, ec) ? fs::file_size(zipPath, ec) : 0;
            if (sz == 0) fs::remove(zipPath, ec);
            worker_running_ = false;
            return;
        }

        if (cancel_requested_) {
            setProgress(State::Failed, "Download cancelled");
            // Keep partial file for resume on retry.
            worker_running_ = false;
            return;
        }
        // Step 2: Verify checksums
        {
            // Derive base URL from the zip URL (strip filename)
            std::string baseUrl = url;
            auto lastSlash = baseUrl.rfind('/');
            if (lastSlash != std::string::npos)
                baseUrl = baseUrl.substr(0, lastSlash);

            if (!verifyChecksums(zipPath, baseUrl)) {
                if (!cancel_requested_) {
                    // Checksum failure — delete the corrupt file so next attempt re-downloads
                    std::error_code ec;
                    fs::remove(zipPath, ec);
                }
                worker_running_ = false;
                return;
            }
        }

        if (cancel_requested_) {
            setProgress(State::Failed, "Verification cancelled");
            worker_running_ = false;
            return;
        }
        // Step 3: Ensure daemon is fully stopped before touching chain data
        {
            setProgress(State::Extracting, "Waiting for daemon to stop...");
            int waited = 0;
            while (daemon::EmbeddedDaemon::isRpcPortInUse() && waited < 60 && !cancel_requested_) {
                std::this_thread::sleep_for(std::chrono::seconds(1));
                waited++;
                if (waited % 5 == 0)
                    DEBUG_LOGF("[Bootstrap] Still waiting for daemon to stop... (%ds)\n", waited);
            }
            if (cancel_requested_) {
                setProgress(State::Failed, "Cancelled while waiting for daemon");
                worker_running_ = false;
                return;
            }
            if (daemon::EmbeddedDaemon::isRpcPortInUse()) {
                setProgress(State::Failed, "Daemon is still running — stop it before using bootstrap");
                worker_running_ = false;
                return;
            }
        }

        // Step 4: Clean old chain data
        setProgress(State::Extracting, "Removing old chain data...");
        cleanChainData(dataDir);

        // Step 5: Extract (skipping wallet.dat)
        if (!extract(zipPath, dataDir)) {
            if (cancel_requested_)
                setProgress(State::Failed, "Extraction cancelled");
            else
                setProgress(State::Failed, "Extraction failed — zip file may be corrupted");
            worker_running_ = false;
            return;
        }

        setProgress(State::Completed, "Bootstrap complete!");
        worker_running_ = false;
    });
}

void Bootstrap::cancel() {
    cancel_requested_ = true;
}

Bootstrap::Progress Bootstrap::getProgress() const {
    std::lock_guard<std::mutex> lk(mutex_);
    return progress_;
}

bool Bootstrap::isDone() const {
    std::lock_guard<std::mutex> lk(mutex_);
    return progress_.state == State::Completed || progress_.state == State::Failed;
}

// ---------------------------------------------------------------------------
// Progress update (thread-safe)
// ---------------------------------------------------------------------------

void Bootstrap::setProgress(State state, const std::string& text, double downloaded, double total) {
    std::lock_guard<std::mutex> lk(mutex_);
    progress_.state = state;
    progress_.status_text = text;
    progress_.downloaded_bytes = downloaded;
    progress_.total_bytes = total;
    progress_.percent = (total > 0) ? (float)(100.0 * downloaded / total) : 0.0f;
    if (state == State::Failed) {
        progress_.error = text;
    }
}

// ---------------------------------------------------------------------------
// HEAD request for remote file size
// ---------------------------------------------------------------------------

long long Bootstrap::getRemoteFileSize(const std::string& url) {
    CURL* curl = curl_easy_init();
    if (!curl) return -1;

    curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
    curl_easy_setopt(curl, CURLOPT_NOBODY, 1L);           // HEAD request
    curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
    curl_easy_setopt(curl, CURLOPT_USERAGENT, "ObsidianDragon/1.0");
    curl_easy_setopt(curl, CURLOPT_TIMEOUT, 15L);
    curl_easy_setopt(curl, CURLOPT_CONNECTTIMEOUT, 10L);
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 1L);
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 2L);

    CURLcode res = curl_easy_perform(curl);
    if (res != CURLE_OK) {
        DEBUG_LOGF("[Bootstrap] HEAD request failed: %s\n", curl_easy_strerror(res));
        curl_easy_cleanup(curl);
        return -1;
    }

    curl_off_t cl = -1;
    curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD_T, &cl);
    curl_easy_cleanup(curl);

    DEBUG_LOGF("[Bootstrap] Remote file size: %lld bytes\n", (long long)cl);
    return (long long)cl;
}

// ---------------------------------------------------------------------------
// Download (libcurl)
// ---------------------------------------------------------------------------

int Bootstrap::progressCallback(void* clientp, long long dltotal, long long dlnow,
                                long long /*ultotal*/, long long /*ulnow*/) {
    auto* self = static_cast<Bootstrap*>(clientp);
    if (self->cancel_requested_) return 1; // abort transfer

    // When resuming, dlnow/dltotal only reflect the *remaining* portion.
    // Add resume_offset_ so the user sees true total progress.
    double offset   = self->resume_offset_;
    double realNow  = offset + (double)dlnow;
    double realTotal = (dltotal > 0) ? (offset + (double)dltotal) : 0.0;

    std::string sizeText = formatSize(realNow) + " / " +
                           (realTotal > 0 ? formatSize(realTotal) : "unknown");
    self->setProgress(State::Downloading, sizeText, realNow, realTotal);
    return 0;
}

bool Bootstrap::download(const std::string& url, const std::string& destZip) {
    CURL* curl = curl_easy_init();
    if (!curl) {
        setProgress(State::Failed, "Failed to initialise libcurl");
        return false;
    }

    // --- Resume support with integrity check ---
    // Check if a partial file exists from a previous attempt.
    long long existing_bytes = 0;
    {
        std::error_code ec;
        if (fs::exists(destZip, ec) && !ec) {
            existing_bytes = (long long)fs::file_size(destZip, ec);
            if (ec) existing_bytes = 0;
        }
    }

    // If we have a partial file, check the remote file size to detect
    // server-side changes that would corrupt a resumed download.
    if (existing_bytes > 0) {
        long long remoteSize = getRemoteFileSize(url);
        if (remoteSize > 0) {
            if (existing_bytes >= remoteSize) {
                // File is already complete (or bigger — stale from an older zip)
                DEBUG_LOGF("[Bootstrap] Local file (%lld bytes) >= remote (%lld bytes), re-downloading\n",
                       existing_bytes, remoteSize);
                std::error_code ec;
                fs::remove(destZip, ec);
                existing_bytes = 0;
            }
        } else {
            // Could not determine remote size — play it safe and re-download
            DEBUG_LOGF("[Bootstrap] Could not determine remote file size, re-downloading to be safe\n");
            std::error_code ec;
            fs::remove(destZip, ec);
            existing_bytes = 0;
        }
    }

    resume_offset_ = (double)existing_bytes;

    FILE* fp = nullptr;
    if (existing_bytes > 0) {
        // Open in append-binary mode and tell curl to resume
        fp = fopen(destZip.c_str(), "ab");
        if (!fp) {
            curl_easy_cleanup(curl);
            setProgress(State::Failed, "Failed to open output file for resume: " + destZip);
            return false;
        }
        curl_easy_setopt(curl, CURLOPT_RESUME_FROM_LARGE, (curl_off_t)existing_bytes);
        DEBUG_LOGF("[Bootstrap] Resuming download from byte %lld\n", existing_bytes);
    } else {
        fp = fopen(destZip.c_str(), "wb");
        if (!fp) {
            curl_easy_cleanup(curl);
            setProgress(State::Failed, "Failed to open output file: " + destZip);
            return false;
        }
    }

    curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, writeFileCallback);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, fp);
    curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L);
    curl_easy_setopt(curl, CURLOPT_XFERINFOFUNCTION, progressCallback);
    curl_easy_setopt(curl, CURLOPT_XFERINFODATA, this);
    curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
    curl_easy_setopt(curl, CURLOPT_USERAGENT, "ObsidianDragon/1.0");
    curl_easy_setopt(curl, CURLOPT_TIMEOUT, 0L);           // no timeout for large file
    curl_easy_setopt(curl, CURLOPT_CONNECTTIMEOUT, 30L);
    curl_easy_setopt(curl, CURLOPT_LOW_SPEED_LIMIT, 1024L); // abort if < 1 KB/s
    curl_easy_setopt(curl, CURLOPT_LOW_SPEED_TIME, 60L);    // ... for 60 seconds

    // HTTPS certificate verification
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 1L);
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 2L);

    CURLcode res = curl_easy_perform(curl);
    fclose(fp);
    curl_easy_cleanup(curl);

    if (res == CURLE_RANGE_ERROR) {
        // Server does not support range requests — restart from scratch
        DEBUG_LOGF("[Bootstrap] Server does not support resume, restarting download\n");
        resume_offset_ = 0;
        std::error_code ec;
        fs::remove(destZip, ec);
        return download(url, destZip); // recursive retry without resume
    }

    if (res != CURLE_OK) {
        DEBUG_LOGF("[Bootstrap] curl error: %s\n", curl_easy_strerror(res));
        return false;
    }
    return true;
}

// ---------------------------------------------------------------------------
// Zip extraction (miniz)
// ---------------------------------------------------------------------------

bool Bootstrap::extract(const std::string& zipPath, const std::string& dataDir) {
    mz_zip_archive zip = {};
    if (!mz_zip_reader_init_file(&zip, zipPath.c_str(), 0)) {
        setProgress(State::Failed, "Failed to open zip file");
        return false;
    }

    int numFiles = (int)mz_zip_reader_get_num_files(&zip);
    for (int i = 0; i < numFiles; i++) {
        if (cancel_requested_) {
            mz_zip_reader_end(&zip);
            return false;
        }

        mz_zip_archive_file_stat stat;
        if (!mz_zip_reader_file_stat(&zip, i, &stat)) continue;

        std::string filename = stat.m_filename;

        // *** CRITICAL: Skip wallet.dat ***
        if (filename == "wallet.dat" || endsWith(filename, "/wallet.dat")) {
            DEBUG_LOGF("[Bootstrap] Skipping wallet.dat (protected)\n");
            continue;
        }

        std::string destPath = dataDir;
        // Ensure trailing separator
        if (!destPath.empty() && destPath.back() != '/' && destPath.back() != '\\')
            destPath += '/';
        destPath += filename;

        if (mz_zip_reader_is_file_a_directory(&zip, i)) {
            std::error_code ec;
            fs::create_directories(destPath, ec);
        } else {
            // Ensure parent directory exists
            std::error_code ec;
            fs::create_directories(fs::path(destPath).parent_path(), ec);

            if (!mz_zip_reader_extract_to_file(&zip, i, destPath.c_str(), 0)) {
                DEBUG_LOGF("[Bootstrap] Failed to extract: %s\n", filename.c_str());
                // Critical chain data must extract successfully
                if (filename.rfind("blocks/", 0) == 0 || filename.rfind("chainstate/", 0) == 0) {
                    setProgress(State::Failed, "Failed to extract critical file: " + filename);
                    mz_zip_reader_end(&zip);
                    return false;
                }
                // Non-fatal for other files: continue with remaining
            }
        }

        // Update progress
        float pct = (numFiles > 0) ? (100.0f * (i + 1) / numFiles) : 0.0f;
        setProgress(State::Extracting,
                    "Extracting: " + filename,
                    (double)(i + 1), (double)numFiles);
        progress_.percent = pct; // override to use file count ratio
    }

    mz_zip_reader_end(&zip);

    // Clean up zip file
    std::error_code ec;
    fs::remove(zipPath, ec);

    return true;
}

// ---------------------------------------------------------------------------
// Pre-extraction cleanup
// ---------------------------------------------------------------------------

void Bootstrap::cleanChainData(const std::string& dataDir) {
    // Directories to remove completely.
    // NOTE: "database" is intentionally NOT removed here. It contains BDB
    // environment/log files for wallet.dat. Deleting it while keeping wallet.dat
    // creates a BDB LSN mismatch that causes the daemon to "salvage" the wallet
    // (renaming it to wallet.{timestamp}.bak and creating an empty replacement).
    // The daemon's DB_RECOVER flag in CDBEnv::Open() handles stale environment
    // files gracefully when the environment dir still exists.
    for (const char* subdir : {"blocks", "chainstate", "notarizations"}) {
        fs::path p = fs::path(dataDir) / subdir;
        std::error_code ec;
        if (fs::exists(p, ec)) {
            fs::remove_all(p, ec);
            if (ec) DEBUG_LOGF("[Bootstrap] Warning: could not remove %s: %s\n", subdir, ec.message().c_str());
        }
    }
    // Individual files to remove (will be replaced by bootstrap)
    for (const char* file : {"fee_estimates.dat", "peers.dat", "banlist.dat", "db.log", ".lock"}) {
        fs::path p = fs::path(dataDir) / file;
        std::error_code ec;
        if (fs::exists(p, ec)) {
            fs::remove(p, ec);
        }
    }
    // NEVER remove: wallet.dat, debug.log, *.conf
}

// ---------------------------------------------------------------------------
// Small-file download (checksums)
// ---------------------------------------------------------------------------

static size_t writeStringCallback(void* contents, size_t size, size_t nmemb, void* userp) {
    size_t total = size * nmemb;
    auto* str = static_cast<std::string*>(userp);
    str->append(static_cast<const char*>(contents), total);
    return total;
}

std::string Bootstrap::downloadSmallFile(const std::string& url) {
    CURL* curl = curl_easy_init();
    if (!curl) return {};

    std::string result;
    curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, writeStringCallback);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, &result);
    curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
    curl_easy_setopt(curl, CURLOPT_USERAGENT, "ObsidianDragon/1.0");
    curl_easy_setopt(curl, CURLOPT_TIMEOUT, 30L);
    curl_easy_setopt(curl, CURLOPT_CONNECTTIMEOUT, 15L);
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 1L);
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 2L);

    CURLcode res = curl_easy_perform(curl);
    curl_easy_cleanup(curl);

    if (res != CURLE_OK) {
        DEBUG_LOGF("[Bootstrap] Failed to download %s: %s\n", url.c_str(), curl_easy_strerror(res));
        return {};
    }
    return result;
}

// ---------------------------------------------------------------------------
// Checksum helpers
// ---------------------------------------------------------------------------

std::string Bootstrap::parseChecksumFile(const std::string& content) {
    // Typical format: "<hex_hash>  <filename>" or just "<hex_hash>"
    // Extract the first whitespace-delimited token.
    std::istringstream iss(content);
    std::string token;
    if (iss >> token) {
        // Normalise to lowercase
        std::transform(token.begin(), token.end(), token.begin(),
                       [](unsigned char c) { return std::tolower(c); });
        return token;
    }
    return {};
}

std::string Bootstrap::computeSHA256(const std::string& filePath, float pctBase, float pctRange) {
    FILE* fp = fopen(filePath.c_str(), "rb");
    if (!fp) return {};

    // Get file size for progress reporting
    fseek(fp, 0, SEEK_END);
    long long fileSize = ftell(fp);
    fseek(fp, 0, SEEK_SET);

    crypto_hash_sha256_state state;
    crypto_hash_sha256_init(&state);

    unsigned char buf[65536];
    size_t n;
    long long processed = 0;
    while ((n = fread(buf, 1, sizeof(buf), fp)) > 0) {
        if (cancel_requested_) {
            fclose(fp);
            return {};
        }
        crypto_hash_sha256_update(&state, buf, n);
        processed += (long long)n;

        // Update progress every ~4MB
        if (fileSize > 0 && (processed % (4 * 1024 * 1024)) < (long long)sizeof(buf)) {
            float filePct = (float)(100.0 * processed / fileSize);
            float overallPct = pctBase + pctRange * (float)processed / (float)fileSize;
            char msg[128];
            snprintf(msg, sizeof(msg), "Verifying SHA-256... %.0f%% (%s / %s)",
                     filePct, formatSize((double)processed).c_str(),
                     formatSize((double)fileSize).c_str());
            {
                std::lock_guard<std::mutex> lk(mutex_);
                progress_.state = State::Verifying;
                progress_.status_text = msg;
                progress_.percent = overallPct;
            }
        }
    }
    fclose(fp);

    unsigned char hash[crypto_hash_sha256_BYTES];
    crypto_hash_sha256_final(&state, hash);

    std::ostringstream oss;
    oss << std::hex << std::setfill('0');
    for (int i = 0; i < crypto_hash_sha256_BYTES; i++)
        oss << std::setw(2) << (int)hash[i];
    return oss.str();
}

// ---------------------------------------------------------------------------
// MD5 — minimal embedded implementation (RFC 1321)
// ---------------------------------------------------------------------------

namespace {

struct MD5Context {
    uint32_t state[4];
    uint64_t count;
    uint8_t  buffer[64];
};

static const uint32_t md5_T[64] = {
    0xd76aa478,0xe8c7b756,0x242070db,0xc1bdceee,0xf57c0faf,0x4787c62a,0xa8304613,0xfd469501,
    0x698098d8,0x8b44f7af,0xffff5bb1,0x895cd7be,0x6b901122,0xfd987193,0xa679438e,0x49b40821,
    0xf61e2562,0xc040b340,0x265e5a51,0xe9b6c7aa,0xd62f105d,0x02441453,0xd8a1e681,0xe7d3fbc8,
    0x21e1cde6,0xc33707d6,0xf4d50d87,0x455a14ed,0xa9e3e905,0xfcefa3f8,0x676f02d9,0x8d2a4c8a,
    0xfffa3942,0x8771f681,0x6d9d6122,0xfde5380c,0xa4beea44,0x4bdecfa9,0xf6bb4b60,0xbebfbc70,
    0x289b7ec6,0xeaa127fa,0xd4ef3085,0x04881d05,0xd9d4d039,0xe6db99e5,0x1fa27cf8,0xc4ac5665,
    0xf4292244,0x432aff97,0xab9423a7,0xfc93a039,0x655b59c3,0x8f0ccc92,0xffeff47d,0x85845dd1,
    0x6fa87e4f,0xfe2ce6e0,0xa3014314,0x4e0811a1,0xf7537e82,0xbd3af235,0x2ad7d2bb,0xeb86d391
};
static const int md5_S[64] = {
    7,12,17,22,7,12,17,22,7,12,17,22,7,12,17,22,
    5, 9,14,20,5, 9,14,20,5, 9,14,20,5, 9,14,20,
    4,11,16,23,4,11,16,23,4,11,16,23,4,11,16,23,
    6,10,15,21,6,10,15,21,6,10,15,21,6,10,15,21
};

static inline uint32_t md5_rotl(uint32_t x, int n) { return (x << n) | (x >> (32 - n)); }

static void md5_transform(uint32_t state[4], const uint8_t block[64]) {
    uint32_t M[16];
    for (int i = 0; i < 16; i++)
        M[i] = (uint32_t)block[i*4] | ((uint32_t)block[i*4+1]<<8) |
               ((uint32_t)block[i*4+2]<<16) | ((uint32_t)block[i*4+3]<<24);

    uint32_t a=state[0], b=state[1], c=state[2], d=state[3];
    for (int i = 0; i < 64; i++) {
        uint32_t f, g;
        if (i < 16)      { f = (b & c) | (~b & d);         g = i; }
        else if (i < 32) { f = (d & b) | (~d & c);         g = (5*i+1) % 16; }
        else if (i < 48) { f = b ^ c ^ d;                  g = (3*i+5) % 16; }
        else              { f = c ^ (b | ~d);               g = (7*i) % 16; }
        uint32_t tmp = d; d = c; c = b;
        b = b + md5_rotl(a + f + md5_T[i] + M[g], md5_S[i]);
        a = tmp;
    }
    state[0]+=a; state[1]+=b; state[2]+=c; state[3]+=d;
}

static void md5_init(MD5Context* ctx) {
    ctx->state[0]=0x67452301; ctx->state[1]=0xefcdab89;
    ctx->state[2]=0x98badcfe; ctx->state[3]=0x10325476;
    ctx->count = 0;
    memset(ctx->buffer, 0, 64);
}

static void md5_update(MD5Context* ctx, const uint8_t* data, size_t len) {
    size_t idx = (size_t)(ctx->count % 64);
    ctx->count += len;
    for (size_t i = 0; i < len; i++) {
        ctx->buffer[idx++] = data[i];
        if (idx == 64) { md5_transform(ctx->state, ctx->buffer); idx = 0; }
    }
}

static void md5_final(MD5Context* ctx, uint8_t digest[16]) {
    uint64_t bits = ctx->count * 8;
    uint8_t pad = 0x80;
    md5_update(ctx, &pad, 1);
    pad = 0;
    while (ctx->count % 64 != 56) md5_update(ctx, &pad, 1);
    uint8_t bitbuf[8];
    for (int i = 0; i < 8; i++) bitbuf[i] = (uint8_t)(bits >> (i*8));
    md5_update(ctx, bitbuf, 8);
    for (int i = 0; i < 4; i++) {
        digest[i*4+0] = (uint8_t)(ctx->state[i]);
        digest[i*4+1] = (uint8_t)(ctx->state[i]>>8);
        digest[i*4+2] = (uint8_t)(ctx->state[i]>>16);
        digest[i*4+3] = (uint8_t)(ctx->state[i]>>24);
    }
}

} // anon namespace

std::string Bootstrap::computeMD5(const std::string& filePath, float pctBase, float pctRange) {
    FILE* fp = fopen(filePath.c_str(), "rb");
    if (!fp) return {};

    // Get file size for progress reporting
    fseek(fp, 0, SEEK_END);
    long long fileSize = ftell(fp);
    fseek(fp, 0, SEEK_SET);

    MD5Context ctx;
    md5_init(&ctx);

    uint8_t buf[65536];
    size_t n;
    long long processed = 0;
    while ((n = fread(buf, 1, sizeof(buf), fp)) > 0) {
        if (cancel_requested_) {
            fclose(fp);
            return {};
        }
        md5_update(&ctx, buf, n);
        processed += (long long)n;

        // Update progress every ~4MB
        if (fileSize > 0 && (processed % (4 * 1024 * 1024)) < (long long)sizeof(buf)) {
            float filePct = (float)(100.0 * processed / fileSize);
            float overallPct = pctBase + pctRange * (float)processed / (float)fileSize;
            char msg[128];
            snprintf(msg, sizeof(msg), "Verifying MD5... %.0f%% (%s / %s)",
                     filePct, formatSize((double)processed).c_str(),
                     formatSize((double)fileSize).c_str());
            {
                std::lock_guard<std::mutex> lk(mutex_);
                progress_.state = State::Verifying;
                progress_.status_text = msg;
                progress_.percent = overallPct;
            }
        }
    }
    fclose(fp);

    uint8_t digest[16];
    md5_final(&ctx, digest);

    std::ostringstream oss;
    oss << std::hex << std::setfill('0');
    for (int i = 0; i < 16; i++)
        oss << std::setw(2) << (int)digest[i];
    return oss.str();
}

// ---------------------------------------------------------------------------
// Checksum verification pipeline
// ---------------------------------------------------------------------------

bool Bootstrap::verifyChecksums(const std::string& zipPath, const std::string& baseUrl) {
    setProgress(State::Verifying, "Downloading checksums...");

    std::string sha256Url = baseUrl + "/" + kZipName + ".sha256";
    std::string md5Url    = baseUrl + "/" + kZipName + ".md5";

    std::string sha256Content = downloadSmallFile(sha256Url);
    std::string md5Content    = downloadSmallFile(md5Url);

    if (cancel_requested_) return false;

    bool haveSHA256 = !sha256Content.empty();
    bool haveMD5    = !md5Content.empty();

    if (!haveSHA256 && !haveMD5) {
        DEBUG_LOGF("[Bootstrap] Warning: no checksum files available — skipping verification\n");
        // Allow the process to continue (server may not have checksum files yet)
        return true;
    }

    // Determine progress ranges: if both checksums exist, split 0-50% / 50-100%
    float sha256Base = 0.0f, sha256Range = 0.0f;
    float md5Base    = 0.0f, md5Range    = 0.0f;
    if (haveSHA256 && haveMD5) {
        sha256Base = 0.0f;  sha256Range = 50.0f;
        md5Base    = 50.0f; md5Range    = 50.0f;
    } else if (haveSHA256) {
        sha256Base = 0.0f;  sha256Range = 100.0f;
    } else {
        md5Base = 0.0f;     md5Range = 100.0f;
    }

    // --- SHA-256 ---
    if (haveSHA256) {
        setProgress(State::Verifying, "Verifying SHA-256...", 0, 1); // percent = 0
        std::string expected = parseChecksumFile(sha256Content);
        std::string actual   = computeSHA256(zipPath, sha256Base, sha256Range);

        if (cancel_requested_) return false;

        if (expected.empty() || actual.empty()) {
            DEBUG_LOGF("[Bootstrap] SHA-256: could not compute/parse (expected=%s, actual=%s)\n",
                   expected.c_str(), actual.c_str());
            setProgress(State::Failed, "SHA-256 verification error");
            return false;
        }

        if (expected != actual) {
            DEBUG_LOGF("[Bootstrap] SHA-256 MISMATCH!\n  expected: %s\n  actual:   %s\n",
                   expected.c_str(), actual.c_str());
            setProgress(State::Failed,
                "SHA-256 mismatch — the server's checksum file may be out of date.\n"
                "Expected: " + expected.substr(0, 16) + "...\n"
                "Got:      " + actual.substr(0, 16) + "...\n"
                "Try again or use a fresh download.");
            return false;
        }
        DEBUG_LOGF("[Bootstrap] SHA-256 verified: %s\n", actual.c_str());
    }

    // --- MD5 ---
    if (haveMD5) {
        setProgress(State::Verifying, "Verifying MD5...", (double)md5Base, 100.0);
        std::string expected = parseChecksumFile(md5Content);
        std::string actual   = computeMD5(zipPath, md5Base, md5Range);

        if (cancel_requested_) return false;

        if (expected.empty() || actual.empty()) {
            DEBUG_LOGF("[Bootstrap] MD5: could not compute/parse (expected=%s, actual=%s)\n",
                   expected.c_str(), actual.c_str());
            setProgress(State::Failed, "MD5 verification error");
            return false;
        }

        if (expected != actual) {
            DEBUG_LOGF("[Bootstrap] MD5 MISMATCH!\n  expected: %s\n  actual:   %s\n",
                   expected.c_str(), actual.c_str());
            setProgress(State::Failed,
                "MD5 mismatch — the server's checksum file may be out of date.\n"
                "Expected: " + expected + "\n"
                "Got:      " + actual + "\n"
                "Try again or use a fresh download.");
            return false;
        }
        DEBUG_LOGF("[Bootstrap] MD5 verified: %s\n", actual.c_str());
    }

    setProgress(State::Verifying, "Checksums verified \xe2\x9c\x93", 100.0, 100.0);
    return true;
}

} // namespace util
} // namespace dragonx