/*
    * IntelFirmware.cpp
    * Intel Bluetooth bootloader firmware download path
    *
    * Ports the btintel "TLV" bootloader sequence: read the version, derive the
    * ibt-<cnvi>-<cnvr>.sfi / .ddc file names, secure-send the signed RSA/ECDSA
    * header and command-buffer payload, boot the operational image and apply
    * the DDC parameters.  The .sfi/.ddc files are staged on the ramdisk at
    * 0:/os/firmware/intel/ by the userspace build.
    *
    * Copyright (c) 2026 Daniel Hammer
*/

#include "IntelFirmware.hpp"
#include <Arch/Cpu.hpp>
#include "Hci.hpp"
#include <Fs/Vfs.hpp>
#include <Memory/Heap.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Libraries/Memory.hpp>
#include <Timekeeping/Timer.hpp>

using namespace Kt;

namespace Drivers::USB::Bluetooth {

    // =========================================================================
    // .sfi / TLV format constants (from drivers/bluetooth/btintel.{c,h})
    // =========================================================================

    // CSS secure-boot header layout.
    static constexpr uint32_t RSA_HEADER_LEN   = 644;   // RSA header + key + sig
    static constexpr uint32_t ECDSA_HEADER_LEN = 320;   // follows the RSA header
    static constexpr uint32_t ECDSA_OFFSET     = 644;
    static constexpr uint32_t CSS_HEADER_OFFSET = 8;    // version field offset

    static constexpr uint32_t RSA_HEADER_VER    = 0x00010000;
    static constexpr uint32_t HYBRID_HEADER_VER = 0x00069700;
    static constexpr uint32_t ECDSA_HEADER_VER  = 0x00020000;

    // The operational image embeds its reset/boot address in a final
    // CMD_WRITE_BOOT_PARAMS HCI command.
    static constexpr uint16_t CMD_WRITE_BOOT_PARAMS = 0xFC0E;

    // Image type (INTEL_TLV_IMAGE_TYPE value): 0x01 bootloader, 0x02
    // intermediate loader, 0x03 operational.
    static constexpr uint8_t IMG_BOOTLOADER = 0x01;
    static constexpr uint8_t IMG_OPERATIONAL = 0x03;

    // Version TLV type ids (enum INTEL_TLV_* starts at 0x10).
    static constexpr uint8_t TLV_CNVI_TOP   = 0x10;
    static constexpr uint8_t TLV_CNVR_TOP   = 0x11;
    static constexpr uint8_t TLV_CNVI_BT    = 0x12;
    static constexpr uint8_t TLV_IMAGE_TYPE = 0x1C;
    static constexpr uint8_t TLV_MIN_FW     = 0x2D;
    static constexpr uint8_t TLV_SBE_TYPE   = 0x2F;

    // =========================================================================
    // Helpers
    // =========================================================================

    static uint32_t Rd32(const uint8_t* p) {
        return (uint32_t)p[0] | ((uint32_t)p[1] << 8)
             | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
    }

    struct IntelTlvVersion {
        uint32_t CnviTop = 0;
        uint32_t CnvrTop = 0;
        uint32_t CnviBt  = 0;
        uint8_t  ImgType = 0;
        uint8_t  SbeType = 0xFF;   // 0x00 RSA, 0x01 ECDSA, 0xFF unknown
        bool     SbeValid = false;
        bool     Valid = false;
    };

    static void ParseTlv(const uint8_t* buf, int len, IntelTlvVersion* v) {
        int i = 0;
        while (i + 2 <= len) {
            uint8_t type = buf[i];
            uint8_t l = buf[i + 1];
            if (i + 2 + l > len) break;
            const uint8_t* val = &buf[i + 2];

            switch (type) {
                case TLV_CNVI_TOP:   if (l >= 4) v->CnviTop = Rd32(val); break;
                case TLV_CNVR_TOP:   if (l >= 4) v->CnvrTop = Rd32(val); break;
                case TLV_CNVI_BT:    if (l >= 4) v->CnviBt  = Rd32(val); break;
                case TLV_IMAGE_TYPE: if (l >= 1) v->ImgType = val[0];    break;
                case TLV_SBE_TYPE:   if (l >= 1) { v->SbeType = val[0]; v->SbeValid = true; } break;
                default: break;
            }
            i += 2 + l;
        }
        // CNVI top is always present on a real TLV response.
        v->Valid = (v->CnviTop != 0);
    }

    // ibt-<cnvi>-<cnvr> packing: __swab16(((top & 0xfff) << 4) | ((top >> 24) & 0xf))
    static uint16_t CnvxPack(uint32_t top) {
        uint16_t t = (uint16_t)(top & 0x0FFF);
        uint8_t  s = (uint8_t)((top >> 24) & 0x0F);
        uint16_t packed = (uint16_t)((t << 4) | s);
        return (uint16_t)((packed >> 8) | (packed << 8));   // __swab16
    }

    static uint8_t HwVariant(uint32_t cnviBt) {
        return (uint8_t)((cnviBt >> 16) & 0x3F);
    }

    static char HexDigit(uint8_t n) {
        return (char)(n < 10 ? ('0' + n) : ('a' + n - 10));
    }

    // Build "0:/os/firmware/intel/ibt-<cnvi>-<cnvr>.<ext>" into out (>= 48 bytes).
    static void BuildFwPath(char* out, uint16_t cnvi, uint16_t cnvr, const char* ext) {
        const char* prefix = "0:/os/firmware/intel/ibt-";
        char* p = out;
        while (*prefix) *p++ = *prefix++;
        for (int sh = 12; sh >= 0; sh -= 4) *p++ = HexDigit((cnvi >> sh) & 0xF);
        *p++ = '-';
        for (int sh = 12; sh >= 0; sh -= 4) *p++ = HexDigit((cnvr >> sh) & 0xF);
        *p++ = '.';
        while (*ext) *p++ = *ext++;
        *p = '\0';
    }

    // =========================================================================
    // Secure-boot header transmission
    // =========================================================================

    // Try one secure-boot header type.  Layout:
    //   RSA  : CSS 128 @0,   pkey 256 @128, sig 256 @388
    //   ECDSA: CSS 128 @644, pkey 96  @772, sig 96  @868   (after the RSA hdr)
    //
    // The CSS init fragment is sent and checked FIRST: a non-zero secure-send
    // result (0xFF/0x06) means the controller's secure-boot engine rejected
    // this header type, so we bail before blasting the ~200-512 byte key and
    // signature -- sending those after a rejected CSS wedges EP0 and (since
    // this runs on the boot path) freezes the kernel.  Returns true only if the
    // full header is accepted (or the controller stays silent, in which case
    // the payload + bootup event remain the final gate).
    static bool TryHeader(const uint8_t* fw, bool ecdsa) {
        const uint32_t cssOff  = ecdsa ? 644 : 0;
        const uint32_t pkeyOff = ecdsa ? 644 + 128 : 128;
        const uint32_t pkeyLen = ecdsa ? 96 : 256;
        const uint32_t sigOff  = ecdsa ? 644 + 224 : 388;
        const uint32_t sigLen  = ecdsa ? 96 : 256;
        const char* name = ecdsa ? "ECDSA" : "RSA";

        KernelLogStream(INFO, "BT-FW") << "Trying " << name << " secure-boot header";

        uint8_t result = 0, status = 0;

        // 1. CSS header init -- check the result before committing key/sig.
        // A rejection event arrives within milliseconds; on SUCCESS the
        // controller stays silent until the END of the whole download (traced
        // on the AX211), so a long timeout here only burns boot time.
        Hci::ClearSecureSendResult();
        if (!Hci::IntelSecureSend(0x00, fw + cssOff, 128)) return false;
        if (Hci::WaitSecureSendResult(250, &result, &status) && (result || status)) {
            KernelLogStream(WARNING, "BT-FW") << name << " CSS rejected (result="
                << base::hex << (uint64_t)result << " status=" << (uint64_t)status
                << base::dec << ")";
            return false;
        }

        // 2. Public key + signature.  Same rule: silence means accepted; a
        // rejected header also fails the payload phase immediately, so the
        // short window loses nothing.
        Hci::ClearSecureSendResult();
        if (!Hci::IntelSecureSend(0x03, fw + pkeyOff, pkeyLen)) return false;
        if (!Hci::IntelSecureSend(0x02, fw + sigOff, sigLen)) return false;
        if (Hci::WaitSecureSendResult(250, &result, &status) && (result || status)) {
            KernelLogStream(WARNING, "BT-FW") << name << " header rejected (result="
                << base::hex << (uint64_t)result << " status=" << (uint64_t)status
                << base::dec << ")";
            return false;
        }

        KernelLogStream(OK, "BT-FW") << name << " secure-boot header accepted";
        return true;
    }

    // Push the command-buffer payload starting at `offset`.  Fragments are cut
    // on HCI command boundaries and only flushed when 4-byte aligned, matching
    // btintel_download_firmware_payload().
    static bool SendPayload(const uint8_t* fw, uint64_t size, uint64_t offset) {
        uint64_t fwPtr = offset;
        uint32_t fragLen = 0;
        uint32_t sends = 0;
        uint64_t start = Timekeeping::GetMilliseconds();

        while (fwPtr < size) {
            // Safety net: never let a stalled upload freeze the boot path.
            if (Timekeeping::GetMilliseconds() - start > 30000) {
                KernelLogStream(ERROR, "BT-FW") << "Payload upload stalled at " << fwPtr - offset
                    << "/" << (size - offset) << " bytes; aborting";
                return false;
            }
            uint64_t cmd = fwPtr + fragLen;
            if (cmd + 3 > size) break;              // truncated header
            uint8_t plen = fw[cmd + 2];
            fragLen += 3 + plen;
            if (fwPtr + fragLen > size) {
                KernelLogStream(ERROR, "BT-FW") << "Firmware payload command overruns file";
                return false;
            }

            if ((fragLen % 4) == 0) {
                // Payload goes over the bulk endpoint, pipelined (the btusb
                // bootloader path) -- EP0's synchronous 3-stage transfers made
                // the 720 KB image the dominant boot cost.
                if (!Hci::IntelSecureSendBulk(0x01, fw + fwPtr, fragLen)) {
                    KernelLogStream(ERROR, "BT-FW") << "Payload secure-send failed at offset "
                        << fwPtr;
                    return false;
                }
                fwPtr += fragLen;
                fragLen = 0;

                // A healthy bootloader only emits an 0xFF/0x06 result on the
                // final fragment.  One arriving mid-stream with a non-zero
                // result/status means it rejected what we just sent -- surface
                // the exact offset instead of blasting the rest of the image
                // into a wedged device (which then NAKs to a cc=4 on EP0).
                uint8_t r = 0, s = 0;
                if (Hci::PeekSecureSendResult(&r, &s) && (r || s)) {
                    KernelLogStream(ERROR, "BT-FW") << "Bootloader rejected payload at byte "
                        << (fwPtr - offset) << "/" << (size - offset) << " (result="
                        << base::hex << (uint64_t)r << " status=" << (uint64_t)s
                        << base::dec << ")";
                    return false;
                }

                // Coarse progress so a slow/stalled upload is distinguishable
                // from steady progress (one line per ~512 fragments).
                if ((++sends % 512) == 0) {
                    KernelLogStream(INFO, "BT-FW") << "  payload " << fwPtr - offset
                        << "/" << (size - offset) << " bytes (" << (uint64_t)sends << " frags)";
                }
            }
        }

        // The bulk pipeline may still hold up to 7 in-flight fragments; make
        // sure every byte reached the controller before the caller waits for
        // the download-complete result.
        if (!Hci::DrainBulkTx(3000)) {
            KernelLogStream(ERROR, "BT-FW") << "Bulk TX drain timed out at end of payload";
            return false;
        }
        return true;
    }

    // Walk the command stream from the start of the image to recover the boot
    // address embedded in the CMD_WRITE_BOOT_PARAMS command.
    static bool FindBootAddress(const uint8_t* fw, uint64_t size, uint32_t* outAddr) {
        uint64_t p = 0;
        while (p + 3 <= size) {
            uint16_t op = (uint16_t)fw[p] | ((uint16_t)fw[p + 1] << 8);
            uint8_t plen = fw[p + 2];
            if (op == CMD_WRITE_BOOT_PARAMS) {
                if (p + 3 + 4 > size) return false;
                *outAddr = Rd32(&fw[p + 3]);
                KernelLogStream(INFO, "BT-FW") << "Boot address: " << base::hex
                    << (uint64_t)*outAddr << base::dec
                    << " (fw build " << (uint64_t)fw[p + 7] << "-"
                    << (uint64_t)fw[p + 8] << "." << (uint64_t)fw[p + 9] << ")";
                return true;
            }
            p += 3 + plen;
        }
        return false;
    }

    // =========================================================================
    // DDC parameter application (best effort)
    // =========================================================================

    static void ApplyDdc(uint16_t cnvi, uint16_t cnvr) {
        char path[48];
        BuildFwPath(path, cnvi, cnvr, "ddc");

        Fs::Vfs::BackendFile file;
        if (Fs::Vfs::OpenBackendFile(path, file) < 0) {
            KernelLogStream(INFO, "BT-FW") << "No DDC parameters (" << path << ")";
            return;
        }

        uint64_t size = Fs::Vfs::GetBackendFileSize(file);
        if (size == 0 || size > 4096) {
            Fs::Vfs::CloseBackendFile(file);
            return;
        }

        uint8_t* ddc = (uint8_t*)Memory::g_heap->Request(size);
        if (!ddc) {
            Fs::Vfs::CloseBackendFile(file);
            return;
        }
        Fs::Vfs::ReadBackendFile(file, ddc, 0, size);
        Fs::Vfs::CloseBackendFile(file);

        // Each record is [len][len bytes of parameters]; the whole record is
        // sent as the 0xFC8B payload.
        uint64_t p = 0;
        int applied = 0;
        while (p < size) {
            uint8_t total = (uint8_t)(ddc[p] + 1);
            if (p + total > size) break;
            if (Hci::IntelWriteDdcRecord(&ddc[p], total)) applied++;
            p += total;
        }

        Memory::g_heap->Free(ddc);
        KernelLogStream(OK, "BT-FW") << "Applied " << (uint64_t)applied << " DDC record(s)";
    }

    // =========================================================================
    // Download orchestration
    // =========================================================================

    static bool DownloadImage(const uint8_t* fw, uint64_t size,
                              const IntelTlvVersion& ver,
                              uint16_t cnvi, uint16_t cnvr) {
        if (size < RSA_HEADER_LEN + 16) {
            KernelLogStream(ERROR, "BT-FW") << "Firmware image too small (" << size << " bytes)";
            return false;
        }

        uint32_t cssVer = Rd32(&fw[CSS_HEADER_OFFSET]);
        if (cssVer != RSA_HEADER_VER && cssVer != HYBRID_HEADER_VER) {
            KernelLogStream(ERROR, "BT-FW") << "Invalid CSS header version: "
                << base::hex << (uint64_t)cssVer << base::dec;
            return false;
        }

        uint32_t bootAddr = 0;
        if (!FindBootAddress(fw, size, &bootAddr)) {
            KernelLogStream(ERROR, "BT-FW") << "Could not locate boot address in firmware";
            return false;
        }

        uint8_t hw = HwVariant(ver.CnviBt);
        uint64_t payloadOffset;
        bool headerOk = false;

        if (hw <= 0x14) {
            // RSA-only parts: RSA header then command buffer.
            headerOk = TryHeader(fw, false);
            payloadOffset = RSA_HEADER_LEN;
        } else {
            // 0x17+ parts ship RSA(644) + ECDSA(320) headers; the command
            // buffer follows both.  sbe_type selects which header the part's
            // secure-boot engine expects.
            if (size < RSA_HEADER_LEN + ECDSA_HEADER_LEN || fw[ECDSA_OFFSET] != 0x06) {
                KernelLogStream(ERROR, "BT-FW") << "Missing ECDSA header for hw variant "
                    << base::hex << (uint64_t)hw << base::dec;
                return false;
            }
            payloadOffset = RSA_HEADER_LEN + ECDSA_HEADER_LEN;

            // Honor sbe_type if the version reply included it; otherwise (the
            // common case here -- the reply is truncated to one USB packet so
            // sbe_type is missing) try ECDSA first.  These newer parts
            // (AX210/AX211, hw_variant >= 0x17) use the ECDSA secure-boot
            // engine, and TryHeader checks the CSS result before committing the
            // key/signature, so a wrong guess fails cleanly instead of wedging.
            bool order[2];
            int n;
            if (ver.SbeValid && ver.SbeType == 0x00) { order[0] = false; n = 1; }      // RSA
            else if (ver.SbeValid && ver.SbeType == 0x01) { order[0] = true; n = 1; }  // ECDSA
            else { order[0] = true; order[1] = false; n = 2; }                          // ECDSA, then RSA

            for (int k = 0; k < n && !headerOk; k++) {
                headerOk = TryHeader(fw, order[k]);
            }
        }

        if (!headerOk) {
            KernelLogStream(ERROR, "BT-FW") << "Secure-boot header rejected";
            return false;
        }

        KernelLogStream(INFO, "BT-FW") << "Downloading firmware payload ("
            << (size - payloadOffset) << " bytes)";
        Hci::ClearSecureSendResult();   // discard header-era result
        if (!SendPayload(fw, size, payloadOffset)) {
            KernelLogStream(ERROR, "BT-FW") << "Firmware payload download failed";
            return false;
        }

        // The bootloader signals download completion with a final 0xFF/0x06
        // secure-send result.  Wait for it (and check it) before booting.
        uint8_t r = 0, s = 0;
        if (Hci::WaitSecureSendResult(5000, &r, &s)) {
            if (r != 0 || s != 0) {
                KernelLogStream(ERROR, "BT-FW") << "Firmware download failed (result="
                    << base::hex << (uint64_t)r << " status=" << (uint64_t)s
                    << base::dec << ")";
                return false;
            }
            KernelLogStream(OK, "BT-FW") << "Firmware download complete";
        } else {
            KernelLogStream(INFO, "BT-FW") << "No download-complete event; booting anyway";
        }

        KernelLogStream(INFO, "BT-FW") << "Booting operational firmware";
        if (!Hci::IntelBootFirmware(bootAddr)) return false;

        KernelLogStream(OK, "BT-FW") << "Operational firmware booted";

        // DDC parameters are applied against the now-running firmware.
        ApplyDdc(cnvi, cnvr);
        return true;
    }

    bool DownloadIntelFirmware() {
        // 1. Read the TLV version to learn the SKU and current image type.
        uint8_t raw[256];
        int rawLen = Hci::ReadIntelVersionTlv(raw, sizeof(raw));
        if (rawLen < 1 || raw[0] != 0) {
            KernelLogStream(WARNING, "BT-FW") << "Could not read Intel TLV version";
            return false;
        }

        IntelTlvVersion ver;
        ParseTlv(raw + 1, rawLen - 1, &ver);
        if (!ver.Valid) {
            KernelLogStream(WARNING, "BT-FW") << "Unrecognized Intel version response";
            return false;
        }

        uint16_t cnvi = CnvxPack(ver.CnviTop);
        uint16_t cnvr = CnvxPack(ver.CnvrTop);

        KernelLogStream(INFO, "BT-FW") << "Intel CNVi=" << base::hex << (uint64_t)cnvi
            << " CNVr=" << (uint64_t)cnvr << " hw_variant=" << (uint64_t)HwVariant(ver.CnviBt)
            << " img_type=" << (uint64_t)ver.ImgType << base::dec;

        // Diagnostic: how much of the TLV version response we actually received
        // (a value near 60 means the controller's reply was truncated to one
        // interrupt packet), and whether sbe_type made it through.
        KernelLogStream(INFO, "BT-FW") << "TLV bytes=" << (uint64_t)rawLen
            << " sbe_valid=" << (uint64_t)(ver.SbeValid ? 1 : 0)
            << " sbe_type=" << base::hex << (uint64_t)ver.SbeType << base::dec;

        // 2. If the operational image is already running, nothing to do.
        if (ver.ImgType == IMG_OPERATIONAL) {
            KernelLogStream(OK, "BT-FW") << "Operational firmware already loaded";
            return true;
        }
        if (ver.ImgType != IMG_BOOTLOADER) {
            KernelLogStream(WARNING, "BT-FW") << "Unexpected image type, attempting download anyway";
        }

        // 3. Load the .sfi image from the ramdisk.
        char sfiPath[48];
        BuildFwPath(sfiPath, cnvi, cnvr, "sfi");
        KernelLogStream(INFO, "BT-FW") << "Loading firmware: " << sfiPath;

        Fs::Vfs::BackendFile file;
        if (Fs::Vfs::OpenBackendFile(sfiPath, file) < 0) {
            KernelLogStream(ERROR, "BT-FW") << "Firmware not found: " << sfiPath;
            KernelLogStream(WARNING, "BT-FW") << "Bluetooth stays in bootloader mode (no firmware)";
            return false;
        }

        uint64_t size = Fs::Vfs::GetBackendFileSize(file);
        uint8_t* fw = (uint8_t*)Memory::g_heap->Request(size);
        if (!fw) {
            KernelLogStream(ERROR, "BT-FW") << "Failed to allocate " << size << " bytes for firmware";
            Fs::Vfs::CloseBackendFile(file);
            return false;
        }
        Fs::Vfs::ReadBackendFile(file, fw, 0, size);
        Fs::Vfs::CloseBackendFile(file);
        Arch::CompilerBarrier();

        // 4. Run the secure download + boot sequence.
        bool ok = DownloadImage(fw, size, ver, cnvi, cnvr);
        Memory::g_heap->Free(fw);
        return ok;
    }

}