/*
    * MassStorage.cpp
    * USB Mass Storage Bulk-Only Transport driver
    * Copyright (c) 2026 Daniel Hammer
*/

#include "MassStorage.hpp"
#include "Xhci.hpp"
#include "UsbDevice.hpp"
#include <Drivers/Storage/BlockDevice.hpp>
#include <Drivers/Storage/Gpt.hpp>
#include <Fs/FsProbe.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Memory/HHDM.hpp>
#include <Memory/PageFrameAllocator.hpp>
#include <Libraries/Memory.hpp>
#include <Timekeeping/ApicTimer.hpp>
#include <atomic>

using namespace Kt;

namespace Drivers::USB::MassStorage {

    // -------------------------------------------------------------------------
    // USB Mass Storage Bulk-Only Transport constants
    // -------------------------------------------------------------------------

    constexpr uint32_t CBW_SIGNATURE = 0x43425355; // "USBC"
    constexpr uint32_t CSW_SIGNATURE = 0x53425355; // "USBS"

    constexpr uint8_t REQ_BULK_ONLY_RESET = 0xFF;
    constexpr uint8_t REQ_GET_MAX_LUN     = 0xFE;

    constexpr uint8_t CBW_FLAG_DATA_IN = 0x80;

    constexpr uint8_t CSW_STATUS_PASSED      = 0x00;
    constexpr uint8_t CSW_STATUS_FAILED      = 0x01;
    constexpr uint8_t CSW_STATUS_PHASE_ERROR = 0x02;

    constexpr uint8_t SCSI_TEST_UNIT_READY = 0x00;
    constexpr uint8_t SCSI_REQUEST_SENSE   = 0x03;
    constexpr uint8_t SCSI_INQUIRY         = 0x12;
    constexpr uint8_t SCSI_READ_CAPACITY10 = 0x25;
    constexpr uint8_t SCSI_READ10          = 0x28;
    constexpr uint8_t SCSI_WRITE10         = 0x2A;
    constexpr uint8_t SCSI_SERVICE_ACTION  = 0x9E;
    constexpr uint8_t SCSI_READ16          = 0x88;
    constexpr uint8_t SCSI_WRITE16         = 0x8A;

    constexpr uint8_t SERVICE_READ_CAPACITY16 = 0x10;

    constexpr uint32_t MAX_BULK_TRANSFER_BYTES = 0x10000;
    constexpr int MAX_MSC_DEVICES = 16;

    struct CommandBlockWrapper {
        uint32_t Signature;
        uint32_t Tag;
        uint32_t DataTransferLength;
        uint8_t  Flags;
        uint8_t  Lun;
        uint8_t  CommandBlockLength;
        uint8_t  CommandBlock[16];
    } __attribute__((packed));

    static_assert(sizeof(CommandBlockWrapper) == 31);

    struct CommandStatusWrapper {
        uint32_t Signature;
        uint32_t Tag;
        uint32_t DataResidue;
        uint8_t  Status;
    } __attribute__((packed));

    static_assert(sizeof(CommandStatusWrapper) == 13);

    struct Device {
        bool     Active;
        uint8_t  SlotId;
        uint8_t  Lun;
        uint8_t  InterfaceNumber;
        uint32_t Tag;

        uint64_t BlockCount;
        uint32_t BlockSize;
        char     Model[41];
        int      BlockDevIndex;

        uint8_t* CbwBuf;
        uint64_t CbwPhys;
        uint8_t* CswBuf;
        uint64_t CswPhys;
        uint8_t* ScratchBuf;
        uint64_t ScratchPhys;
    };

    static Device g_devices[MAX_MSC_DEVICES] = {};
    static std::atomic<bool> g_transportBusy{false};

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

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

    static uint64_t ReadBe64(const uint8_t* p) {
        return ((uint64_t)ReadBe32(p) << 32) | ReadBe32(p + 4);
    }

    static void WriteBe16(uint8_t* p, uint16_t v) {
        p[0] = (uint8_t)(v >> 8);
        p[1] = (uint8_t)v;
    }

    static void WriteBe32(uint8_t* p, uint32_t v) {
        p[0] = (uint8_t)(v >> 24);
        p[1] = (uint8_t)(v >> 16);
        p[2] = (uint8_t)(v >> 8);
        p[3] = (uint8_t)v;
    }

    static void WriteBe64(uint8_t* p, uint64_t v) {
        WriteBe32(p, (uint32_t)(v >> 32));
        WriteBe32(p + 4, (uint32_t)v);
    }

    static void PollWaitMs(uint32_t ms) {
        uint64_t flags;
        asm volatile("pushfq; pop %0" : "=r"(flags));
        if ((flags & (1 << 9)) == 0) {
            for (uint64_t i = 0; i < (uint64_t)ms * 1000; i++) {
                Xhci::PollEvents();
                asm volatile("outb %%al, $0x80" ::: "memory");
            }
            return;
        }

        uint64_t start = Timekeeping::GetMilliseconds();
        while (Timekeeping::GetMilliseconds() - start < ms) {
            Xhci::PollEvents();
            asm volatile("pause" ::: "memory");
        }
    }

    static void TrimTrailingSpaces(char* s, int len) {
        if (!s || len <= 0) return;
        s[len - 1] = '\0';
        for (int i = len - 2; i >= 0; i--) {
            if (s[i] == ' ' || s[i] == '\0') {
                s[i] = '\0';
            } else {
                break;
            }
        }
    }

    static void AppendField(char* dst, int dstMax, const uint8_t* src, int srcLen) {
        int pos = 0;
        while (pos < dstMax - 1 && dst[pos] != '\0') pos++;

        for (int i = 0; i < srcLen && pos < dstMax - 1; i++) {
            char c = (char)src[i];
            if (c < 32 || c > 126) c = ' ';
            dst[pos++] = c;
        }
        dst[pos] = '\0';
    }

    static void BuildInquiryModel(Device& dev, const uint8_t* inquiry) {
        memset(dev.Model, 0, sizeof(dev.Model));
        AppendField(dev.Model, sizeof(dev.Model), inquiry + 8, 8);

        int pos = 0;
        while (pos < 40 && dev.Model[pos] != '\0') pos++;
        if (pos < 40) dev.Model[pos++] = ' ';
        dev.Model[pos] = '\0';

        AppendField(dev.Model, sizeof(dev.Model), inquiry + 16, 16);
        TrimTrailingSpaces(dev.Model, sizeof(dev.Model));
        if (dev.Model[0] == '\0') {
            const char fallback[] = "USB Mass Storage";
            memcpy(dev.Model, fallback, sizeof(fallback));
        }
    }

    static int AllocateDevice() {
        for (int i = 0; i < MAX_MSC_DEVICES; i++) {
            if (g_devices[i].Active) continue;

            if (g_devices[i].CbwBuf == nullptr) {
                g_devices[i].CbwBuf = (uint8_t*)Memory::g_pfa->AllocateZeroed();
                g_devices[i].CbwPhys = Memory::SubHHDM(g_devices[i].CbwBuf);
                g_devices[i].CswBuf = (uint8_t*)Memory::g_pfa->AllocateZeroed();
                g_devices[i].CswPhys = Memory::SubHHDM(g_devices[i].CswBuf);
                g_devices[i].ScratchBuf = (uint8_t*)Memory::g_pfa->AllocateZeroed();
                g_devices[i].ScratchPhys = Memory::SubHHDM(g_devices[i].ScratchBuf);
            }

            return i;
        }

        return -1;
    }

    static bool IsGoodCompletion(uint32_t cc) {
        return cc == Xhci::CC_SUCCESS || cc == Xhci::CC_SHORT_PACKET;
    }

    bool IsTransportBusy() {
        return g_transportBusy.load(std::memory_order_acquire);
    }

    static void AcquireTransport() {
        bool expected = false;
        while (!g_transportBusy.compare_exchange_weak(expected, true,
                                                      std::memory_order_acquire,
                                                      std::memory_order_relaxed)) {
            expected = false;
            asm volatile("pause" ::: "memory");
        }
    }

    static void ReleaseTransport() {
        g_transportBusy.store(false, std::memory_order_release);
    }

    // -------------------------------------------------------------------------
    // BOT transport
    // -------------------------------------------------------------------------

    static bool BulkOnlyResetTransport(uint8_t slotId, uint8_t interfaceNumber) {
        uint32_t cc = Xhci::ControlTransfer(slotId,
            UsbDevice::REQTYPE_CLASS_IFACE,
            REQ_BULK_ONLY_RESET,
            0,
            interfaceNumber,
            0,
            nullptr,
            false);

        if (cc != Xhci::CC_SUCCESS) {
            KernelLogStream(WARNING, "USB-MSC") << "Bulk-only reset failed, cc="
                << (uint64_t)cc;
            return false;
        }

        PollWaitMs(50);
        return true;
    }

    static bool BulkOnlyReset(Device& dev) {
        return BulkOnlyResetTransport(dev.SlotId, dev.InterfaceNumber);
    }

    static uint8_t GetMaxLun(uint8_t slotId, uint8_t interfaceNumber) {
        uint8_t maxLun = 0;
        uint32_t cc = Xhci::ControlTransfer(slotId,
            0xA1, // Device-to-host, class, interface
            REQ_GET_MAX_LUN,
            0,
            interfaceNumber,
            1,
            &maxLun,
            true);

        if (cc != Xhci::CC_SUCCESS && cc != Xhci::CC_SHORT_PACKET) {
            return 0;
        }

        if (maxLun > 15) maxLun = 15;
        return maxLun;
    }

    static bool ScsiCommand(Device& dev, const uint8_t* cdb, uint8_t cdbLen,
                            void* data, uint64_t dataPhys, uint32_t dataLen,
                            bool dataIn) {
        if (cdbLen == 0 || cdbLen > 16) return false;
        if (dataLen > MAX_BULK_TRANSFER_BYTES) return false;
        if (dataLen > 0 && data == nullptr) return false;

        AcquireTransport();

        auto* cbw = (CommandBlockWrapper*)dev.CbwBuf;
        memset(cbw, 0, sizeof(CommandBlockWrapper));
        cbw->Signature = CBW_SIGNATURE;
        cbw->Tag = ++dev.Tag;
        cbw->DataTransferLength = dataLen;
        cbw->Flags = dataIn ? CBW_FLAG_DATA_IN : 0;
        cbw->Lun = dev.Lun;
        cbw->CommandBlockLength = cdbLen;
        memcpy(cbw->CommandBlock, cdb, cdbLen);

        uint32_t actual = 0;
        uint32_t cc = Xhci::BulkTransfer(dev.SlotId, false, dev.CbwBuf, dev.CbwPhys,
                                         sizeof(CommandBlockWrapper), &actual);
        if (!IsGoodCompletion(cc) || actual != sizeof(CommandBlockWrapper)) {
            ReleaseTransport();
            return false;
        }

        if (dataLen > 0) {
            cc = Xhci::BulkTransfer(dev.SlotId, dataIn, data, dataPhys, dataLen, &actual);
            if (!IsGoodCompletion(cc)) {
                ReleaseTransport();
                return false;
            }
        }

        memset(dev.CswBuf, 0, 0x1000);
        cc = Xhci::BulkTransfer(dev.SlotId, true, dev.CswBuf, dev.CswPhys,
                                sizeof(CommandStatusWrapper), &actual);
        if (!IsGoodCompletion(cc) || actual < sizeof(CommandStatusWrapper)) {
            ReleaseTransport();
            return false;
        }

        auto* csw = (CommandStatusWrapper*)dev.CswBuf;
        bool ok = csw->Signature == CSW_SIGNATURE &&
                  csw->Tag == cbw->Tag &&
                  csw->Status == CSW_STATUS_PASSED;

        if (csw->Status == CSW_STATUS_PHASE_ERROR) {
            BulkOnlyReset(dev);
        }

        ReleaseTransport();
        return ok;
    }

    // -------------------------------------------------------------------------
    // SCSI commands
    // -------------------------------------------------------------------------

    static bool RequestSense(Device& dev) {
        uint8_t cdb[6] = {};
        cdb[0] = SCSI_REQUEST_SENSE;
        cdb[4] = 18;

        memset(dev.ScratchBuf, 0, 18);
        bool ok = ScsiCommand(dev, cdb, 6, dev.ScratchBuf, dev.ScratchPhys, 18, true);
        if (ok) {
            uint8_t senseKey = dev.ScratchBuf[2] & 0x0F;
            uint8_t asc = dev.ScratchBuf[12];
            uint8_t ascq = dev.ScratchBuf[13];
            KernelLogStream(INFO, "USB-MSC") << "REQUEST SENSE: key="
                << (uint64_t)senseKey << " asc=" << base::hex << (uint64_t)asc
                << " ascq=" << (uint64_t)ascq << base::dec;
        }
        return ok;
    }

    static bool TestUnitReady(Device& dev) {
        uint8_t cdb[6] = {};
        cdb[0] = SCSI_TEST_UNIT_READY;

        for (int i = 0; i < 5; i++) {
            if (ScsiCommand(dev, cdb, 6, nullptr, 0, 0, false)) return true;
            RequestSense(dev);
            PollWaitMs(50);
        }

        return false;
    }

    static bool Inquiry(Device& dev) {
        uint8_t cdb[6] = {};
        cdb[0] = SCSI_INQUIRY;
        cdb[4] = 36;

        memset(dev.ScratchBuf, 0, 36);
        if (!ScsiCommand(dev, cdb, 6, dev.ScratchBuf, dev.ScratchPhys, 36, true)) {
            return false;
        }

        BuildInquiryModel(dev, dev.ScratchBuf);
        return true;
    }

    static bool ReadCapacity16(Device& dev) {
        uint8_t cdb[16] = {};
        cdb[0] = SCSI_SERVICE_ACTION;
        cdb[1] = SERVICE_READ_CAPACITY16;
        cdb[13] = 32;

        memset(dev.ScratchBuf, 0, 32);
        if (!ScsiCommand(dev, cdb, 16, dev.ScratchBuf, dev.ScratchPhys, 32, true)) {
            return false;
        }

        uint64_t lastLba = ReadBe64(dev.ScratchBuf);
        uint32_t blockLen = ReadBe32(dev.ScratchBuf + 8);
        if (blockLen == 0) return false;

        dev.BlockCount = lastLba + 1;
        dev.BlockSize = blockLen;
        return dev.BlockCount != 0;
    }

    static bool ReadCapacity(Device& dev) {
        uint8_t cdb[10] = {};
        cdb[0] = SCSI_READ_CAPACITY10;

        memset(dev.ScratchBuf, 0, 8);
        if (!ScsiCommand(dev, cdb, 10, dev.ScratchBuf, dev.ScratchPhys, 8, true)) {
            return ReadCapacity16(dev);
        }

        uint32_t lastLba = ReadBe32(dev.ScratchBuf);
        uint32_t blockLen = ReadBe32(dev.ScratchBuf + 4);

        if (lastLba == 0xFFFFFFFF) {
            return ReadCapacity16(dev);
        }
        if (blockLen == 0) return false;

        dev.BlockCount = (uint64_t)lastLba + 1;
        dev.BlockSize = blockLen;
        return dev.BlockCount != 0;
    }

    static bool ReadBlocks(Device& dev, uint64_t lba, uint32_t blocks,
                           void* buffer, uint64_t bufferPhys) {
        uint8_t cdb[16] = {};

        if (lba > 0xFFFFFFFFULL) {
            cdb[0] = SCSI_READ16;
            WriteBe64(cdb + 2, lba);
            WriteBe32(cdb + 10, blocks);
            return ScsiCommand(dev, cdb, 16, buffer, bufferPhys,
                               blocks * dev.BlockSize, true);
        }

        cdb[0] = SCSI_READ10;
        WriteBe32(cdb + 2, (uint32_t)lba);
        WriteBe16(cdb + 7, (uint16_t)blocks);
        return ScsiCommand(dev, cdb, 10, buffer, bufferPhys,
                           blocks * dev.BlockSize, true);
    }

    static bool WriteBlocks(Device& dev, uint64_t lba, uint32_t blocks,
                            void* buffer, uint64_t bufferPhys) {
        uint8_t cdb[16] = {};

        if (lba > 0xFFFFFFFFULL) {
            cdb[0] = SCSI_WRITE16;
            WriteBe64(cdb + 2, lba);
            WriteBe32(cdb + 10, blocks);
            return ScsiCommand(dev, cdb, 16, buffer, bufferPhys,
                               blocks * dev.BlockSize, false);
        }

        cdb[0] = SCSI_WRITE10;
        WriteBe32(cdb + 2, (uint32_t)lba);
        WriteBe16(cdb + 7, (uint16_t)blocks);
        return ScsiCommand(dev, cdb, 10, buffer, bufferPhys,
                           blocks * dev.BlockSize, false);
    }

    // -------------------------------------------------------------------------
    // Block device callbacks
    // -------------------------------------------------------------------------

    static bool ReadSectors(void* ctx, uint64_t lba, uint32_t count, void* buffer) {
        auto* dev = (Device*)ctx;
        if (!dev || !dev->Active || !buffer || count == 0) return false;
        if (lba >= dev->BlockCount || count > dev->BlockCount - lba) return false;

        uint32_t maxBlocks = MAX_BULK_TRANSFER_BYTES / dev->BlockSize;
        if (maxBlocks == 0) maxBlocks = 1;

        uint8_t* dst = (uint8_t*)buffer;
        uint32_t remaining = count;
        uint64_t curLba = lba;

        while (remaining > 0) {
            uint32_t chunk = remaining > maxBlocks ? maxBlocks : remaining;
            uint32_t bytes = chunk * dev->BlockSize;
            int pages = (bytes + 0xFFF) / 0x1000;

            uint64_t dmaPhys;
            uint8_t* dma = (uint8_t*)Memory::g_pfa->ReallocConsecutive(nullptr, pages);
            memset(dma, 0, pages * 0x1000);
            dmaPhys = Memory::SubHHDM(dma);

            bool ok = ReadBlocks(*dev, curLba, chunk, dma, dmaPhys);
            if (ok) {
                memcpy(dst, dma, bytes);
            }

            Memory::g_pfa->Free(dma, pages);
            if (!ok) return false;

            dst += bytes;
            curLba += chunk;
            remaining -= chunk;
        }

        return true;
    }

    static bool WriteSectors(void* ctx, uint64_t lba, uint32_t count, const void* buffer) {
        auto* dev = (Device*)ctx;
        if (!dev || !dev->Active || !buffer || count == 0) return false;
        if (lba >= dev->BlockCount || count > dev->BlockCount - lba) return false;

        uint32_t maxBlocks = MAX_BULK_TRANSFER_BYTES / dev->BlockSize;
        if (maxBlocks == 0) maxBlocks = 1;

        const uint8_t* src = (const uint8_t*)buffer;
        uint32_t remaining = count;
        uint64_t curLba = lba;

        while (remaining > 0) {
            uint32_t chunk = remaining > maxBlocks ? maxBlocks : remaining;
            uint32_t bytes = chunk * dev->BlockSize;
            int pages = (bytes + 0xFFF) / 0x1000;

            uint64_t dmaPhys;
            uint8_t* dma = (uint8_t*)Memory::g_pfa->ReallocConsecutive(nullptr, pages);
            memset(dma, 0, pages * 0x1000);
            memcpy(dma, src, bytes);
            dmaPhys = Memory::SubHHDM(dma);

            bool ok = WriteBlocks(*dev, curLba, chunk, dma, dmaPhys);
            Memory::g_pfa->Free(dma, pages);
            if (!ok) return false;

            src += bytes;
            curLba += chunk;
            remaining -= chunk;
        }

        return true;
    }

    // -------------------------------------------------------------------------
    // Registration
    // -------------------------------------------------------------------------

    static void RegisterBlockDevice(Device& dev) {
        Storage::BlockDevice bdev = {};
        bdev.ReadSectors = ReadSectors;
        bdev.WriteSectors = WriteSectors;
        bdev.Ctx = &dev;
        bdev.SectorCount = dev.BlockCount;
        bdev.SectorSize = (uint16_t)dev.BlockSize;
        bdev.Kind = Storage::BLOCK_KIND_USB_MSC;
        memcpy(bdev.Model, dev.Model, 41);

        int blockIndex = Storage::RegisterBlockDevice(bdev);
        if (blockIndex < 0) {
            KernelLogStream(ERROR, "USB-MSC") << "Failed to register block device";
            return;
        }
        dev.BlockDevIndex = blockIndex;

        uint64_t sizeMB = (dev.BlockCount * dev.BlockSize) / (1024 * 1024);
        uint64_t sizeGB = sizeMB / 1024;

        if (sizeGB > 0) {
            KernelLogStream(OK, "USB-MSC") << "LUN " << (uint64_t)dev.Lun
                << ": " << dev.Model << " (" << sizeGB << " GiB, drive "
                << blockIndex << ")";
        } else {
            KernelLogStream(OK, "USB-MSC") << "LUN " << (uint64_t)dev.Lun
                << ": " << dev.Model << " (" << sizeMB << " MiB, drive "
                << blockIndex << ")";
        }

        Storage::Gpt::ProbeDevice(blockIndex);
        Fs::FsProbe::MountPartitionsForBlockDevice(blockIndex);
    }

    void RegisterDevice(uint8_t slotId) {
        auto* usbDev = Xhci::GetDevice(slotId);
        if (!usbDev || !usbDev->BulkInEpNum || !usbDev->BulkOutEpNum) {
            KernelLogStream(ERROR, "USB-MSC") << "Missing bulk endpoints on slot "
                << (uint64_t)slotId;
            return;
        }

        uint8_t maxLun = GetMaxLun(slotId, usbDev->InterfaceNumber);
        KernelLogStream(INFO, "USB-MSC") << "Slot " << (uint64_t)slotId
            << ": max LUN " << (uint64_t)maxLun;

        BulkOnlyResetTransport(slotId, usbDev->InterfaceNumber);

        for (uint8_t lun = 0; lun <= maxLun; lun++) {
            int idx = AllocateDevice();
            if (idx < 0) {
                KernelLogStream(ERROR, "USB-MSC") << "No free device slots";
                return;
            }

            Device& dev = g_devices[idx];
            dev.Active = true;
            dev.SlotId = slotId;
            dev.Lun = lun;
            dev.InterfaceNumber = usbDev->InterfaceNumber;
            dev.Tag = 0x4D544B00u + idx;
            dev.BlockCount = 0;
            dev.BlockSize = 0;
            dev.BlockDevIndex = -1;
            memset(dev.Model, 0, sizeof(dev.Model));

            if (!Inquiry(dev)) {
                KernelLogStream(WARNING, "USB-MSC") << "INQUIRY failed for LUN "
                    << (uint64_t)lun;
                dev.Active = false;
                continue;
            }

            if (!TestUnitReady(dev)) {
                KernelLogStream(WARNING, "USB-MSC") << "LUN " << (uint64_t)lun
                    << " not ready";
                dev.Active = false;
                continue;
            }

            if (!ReadCapacity(dev)) {
                KernelLogStream(WARNING, "USB-MSC") << "READ CAPACITY failed for LUN "
                    << (uint64_t)lun;
                dev.Active = false;
                continue;
            }

            if (dev.BlockSize == 0 || dev.BlockSize > 0xFFFF ||
                dev.BlockSize > MAX_BULK_TRANSFER_BYTES) {
                KernelLogStream(WARNING, "USB-MSC") << "Unsupported block size "
                    << (uint64_t)dev.BlockSize;
                dev.Active = false;
                continue;
            }

            RegisterBlockDevice(dev);
        }
    }

    void UnregisterDevice(uint8_t slotId) {
        for (int i = 0; i < MAX_MSC_DEVICES; i++) {
            Device& dev = g_devices[i];
            if (!dev.Active || dev.SlotId != slotId) continue;

            int blockIndex = dev.BlockDevIndex;
            if (blockIndex >= 0) {
                Fs::FsProbe::UnmountPartitionsForBlockDevice(blockIndex);
                Storage::Gpt::RemovePartitionsForBlockDevice(blockIndex);
                Storage::UnregisterBlockDevice(blockIndex);
                KernelLogStream(INFO, "USB-MSC") << "Unregistered block device "
                    << blockIndex << " from slot " << (uint64_t)slotId;
            }

            dev.Active = false;
            dev.SlotId = 0;
            dev.Lun = 0;
            dev.InterfaceNumber = 0;
            dev.Tag = 0;
            dev.BlockCount = 0;
            dev.BlockSize = 0;
            dev.BlockDevIndex = -1;
            dev.Model[0] = '\0';
        }
    }

}
