/*
 * SPDX-License-Identifier: Apache-2.0
 * Copyright (C) 2025 Raspberry Pi Ltd
 */

#include "rpiboot_protocol.h"

#include <thread>

namespace rpiboot {

bool RpibootProtocol::uploadBootcode(IUsbTransport& transport,
                                      ChipGeneration gen,
                                      const std::filesystem::path& firmwareDir,
                                      ProgressCallback progress,
                                      std::atomic<bool>& cancelled)
{
    if (progress)
        progress(0, 3, "Uploading bootloader...");

    if (!_bootcodeLoader.send(transport, gen, firmwareDir, cancelled)) {
        _lastError = "Bootcode upload failed: " + _bootcodeLoader.lastError();
        return false;
    }

    return true;
}

bool RpibootProtocol::serveFiles(IUsbTransport& transport,
                                  ChipGeneration gen,
                                  SideloadMode mode,
                                  const std::filesystem::path& firmwareDir,
                                  ProgressCallback progress,
                                  std::atomic<bool>& cancelled)
{
    if (progress)
        progress(2, 3, "Loading firmware...");

    auto sideloadDir = resolveSideloadDir(firmwareDir, gen, mode);

    // If the sideload directory contains a bootfiles.bin tar archive,
    // extract it and provide files from memory.
    auto bootfilesTar = sideloadDir / "bootfiles.bin";
    bool haveBootfiles = false;

    if (std::filesystem::exists(bootfilesTar)) {
        if (!_bootfiles.extractFromFile(bootfilesTar.string())) {
            _lastError = "Failed to extract bootfiles.bin: " + _bootfiles.lastError();
            return false;
        }
        haveBootfiles = true;
    }

    // Custom file resolver that first checks the extracted tar, then disk
    FileResolver resolver;
    if (haveBootfiles) {
        auto prefix = chipDirectoryPrefix(gen);
        resolver = [this, &sideloadDir, prefix](const std::string& filename) -> std::vector<uint8_t> {
            // Try the tar archive first (with chip-specific prefix fallback)
            const auto* data = _bootfiles.find(filename, prefix);
            if (data)
                return *data;

            // Fall back to on-disk files
            return FileServer::readFileFromDisk(sideloadDir, filename);
        };
    }

    // Both sideload modes have an external confirmation mechanism: the
    // caller (RpibootThread) runs a scanner that watches for the next-stage
    // device on the same port path and flips the cancellation flag once it
    // appears.  Requiring that confirmation on disconnect is what lets the
    // file server tell a real cable yank apart from a legitimate reboot.
    //   Fastboot → scanner looks for the fastboot gadget VID/PID.
    //   SBR      → scanner looks for the original rpiboot VID/PID returning
    //              on a fresh device address (FirmwareManager rewrites the
    //              recovery config.txt to ensure set_boot_order=0x3 +
    //              recovery_reboot=1 so the device actually does reboot
    //              back to rpiboot).
    const bool requireReEnumConfirmation =
        (mode == SideloadMode::Fastboot || mode == SideloadMode::SecureBootRecovery);
    if (!_fileServer.run(transport, sideloadDir, progress, cancelled, resolver,
                          requireReEnumConfirmation)) {
        _lastError = "File server failed: " + _fileServer.lastError();
        return false;
    }

    if (progress)
        progress(3, 3, "Device boot complete");

    return true;
}

bool RpibootProtocol::execute(IUsbTransport& transport,
                               ChipGeneration gen,
                               SideloadMode mode,
                               const std::filesystem::path& firmwareDir,
                               ProgressCallback progress,
                               std::atomic<bool>& cancelled)
{
    if (!uploadBootcode(transport, gen, firmwareDir, progress, cancelled))
        return false;

    if (cancelled.load())
        return false;

    // Wait for device to re-enumerate.  Mock transports used in tests
    // don't re-enumerate, so isOpen() returns true immediately.
    if (progress)
        progress(1, 3, "Waiting for device to restart...");

    for (int i = 0; i < 10 && !cancelled.load(); ++i) {
        std::this_thread::sleep_for(std::chrono::milliseconds(200));
        if (transport.isOpen())
            break;
    }

    if (cancelled.load())
        return false;

    if (!transport.isOpen()) {
        _lastError = "Device did not re-enumerate after bootcode upload";
        return false;
    }

    return serveFiles(transport, gen, mode, firmwareDir, progress, cancelled);
}

std::filesystem::path RpibootProtocol::resolveSideloadDir(
    const std::filesystem::path& firmwareDir,
    ChipGeneration gen,
    SideloadMode mode) const
{
    switch (mode) {
    case SideloadMode::Fastboot:
        {
            auto dir = firmwareDir / "fastboot";
            if (std::filesystem::exists(dir))
                return dir;
        }
        break;

    case SideloadMode::SecureBootRecovery:
        if (gen == ChipGeneration::BCM2712) {
            auto dir = firmwareDir / "secure-boot-recovery5";
            if (std::filesystem::exists(dir))
                return dir;
        }
        {
            auto dir = firmwareDir / "secure-boot-recovery";
            if (std::filesystem::exists(dir))
                return dir;
        }
        break;
    }

    // Fall back to the base firmware directory
    return firmwareDir;
}

} // namespace rpiboot
