/*
    * UsbDevice.cpp
    * USB device enumeration and configuration
    * Copyright (c) 2025 Daniel Hammer
*/

#include "UsbDevice.hpp"
#include "Xhci.hpp"
#include "HidKeyboard.hpp"
#include "HidMouse.hpp"
#include "Bluetooth/Bluetooth.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>

using namespace Kt;

static void BusyWaitMs(uint64_t ms) {
    uint64_t flags;
    asm volatile("pushfq; pop %0" : "=r"(flags));
    if (flags & (1 << 9)) {
        // Interrupts enabled — use timer-based delay
        uint64_t start = Timekeeping::GetMilliseconds();
        while (Timekeeping::GetMilliseconds() - start < ms) {
            asm volatile("pause" ::: "memory");
        }
    } else {
        // Interrupts disabled (e.g. timer tick context) — use I/O port delay
        // Each outb to port 0x80 takes ~1µs on x86
        for (uint64_t i = 0; i < ms * 1000; i++) {
            asm volatile("outb %%al, $0x80" ::: "memory");
        }
    }
}

// Access xHCI internal state needed during enumeration
namespace Drivers::USB::Xhci {
    extern volatile uint32_t g_cmdCompletionSlotId;
    extern uint64_t* g_dcbaa;
}

namespace Drivers::USB::UsbDevice {

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

    static uint16_t MaxPacketSizeForSpeed(uint32_t speed) {
        switch (speed) {
            case Xhci::SPEED_LOW:   return 8;
            case Xhci::SPEED_FULL:  return 8;
            case Xhci::SPEED_HIGH:  return 64;
            case Xhci::SPEED_SUPER: return 512;
            default:                return 64;
        }
    }

    // Map xHCI port speed to the slot context speed field value.
    // Per the xHCI spec the slot context speed field uses the same encoding
    // as PORTSC (1=Full, 2=Low, 3=High, 4=Super).
    static uint32_t SpeedToSlotContextValue(uint32_t speed) {
        return speed;  // Same encoding
    }

    // Convert USB endpoint bInterval to xHCI Endpoint Context Interval value.
    // HS/SS: bInterval is already in 2^(n-1) * 125µs encoding — use directly.
    // FS/LS: bInterval is in milliseconds (frames) — convert via fls(bInterval * 8).
    static uint32_t ConvertInterval(uint32_t speed, uint8_t bInterval) {
        if (bInterval == 0) return 0;

        if (speed == Xhci::SPEED_HIGH || speed == Xhci::SPEED_SUPER) {
            return bInterval;
        }

        // FS/LS: bInterval ms → microframes, then find highest set bit position
        uint32_t microframes = (uint32_t)bInterval * 8;
        uint32_t interval = 0;
        while (microframes > 0) {
            interval++;
            microframes >>= 1;
        }
        if (interval > 15) interval = 15;
        return interval;
    }

    // Convert xHCI port speed to a human-readable string
    static const char* SpeedToString(uint32_t speed) {
        switch (speed) {
            case Xhci::SPEED_LOW:   return "Low";
            case Xhci::SPEED_FULL:  return "Full";
            case Xhci::SPEED_HIGH:  return "High";
            case Xhci::SPEED_SUPER: return "Super";
            default:                return "Unknown";
        }
    }

    // ---------------------------------------------------------------------------
    // EnumerateDevice
    // ---------------------------------------------------------------------------

    uint8_t EnumerateDevice(uint8_t portId, uint32_t speed) {
        KernelLogStream(INFO, "USB") << "Enumerating device on port "
            << (uint64_t)portId << " speed=" << SpeedToString(speed);

        // -----------------------------------------------------------------
        // Step 1: Enable Slot
        // -----------------------------------------------------------------
        Xhci::TRB enableSlotTrb = {};
        enableSlotTrb.Control = (Xhci::TRB_ENABLE_SLOT << Xhci::TRB_TYPE_SHIFT);

        uint32_t cc = Xhci::SendCommand(enableSlotTrb);
        if (cc != Xhci::CC_SUCCESS) {
            KernelLogStream(ERROR, "USB") << "Enable Slot failed, cc=" << (uint64_t)cc;
            return 0;
        }

        uint8_t slotId = (uint8_t)Xhci::g_cmdCompletionSlotId;
        if (slotId == 0 || slotId > Xhci::MAX_SLOTS) {
            KernelLogStream(ERROR, "USB") << "Invalid slot ID: " << (uint64_t)slotId;
            return 0;
        }

        KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId << " enabled";

        // -----------------------------------------------------------------
        // Step 2: Allocate device output context and set DCBAA entry
        // -----------------------------------------------------------------
        auto* dev = Xhci::GetDevice(slotId);
        dev->Active = true;
        dev->PortId = portId;
        dev->Speed  = speed;

        // Allocate a zeroed page for the output DeviceContext
        auto* outputCtx = (Xhci::DeviceContext*)Memory::g_pfa->AllocateZeroed();
        dev->OutputContext     = outputCtx;
        dev->OutputContextPhys = Memory::SubHHDM(outputCtx);

        // Point DCBAA[slotId] to the output context physical address
        Xhci::g_dcbaa[slotId] = dev->OutputContextPhys;

        // -----------------------------------------------------------------
        // Step 3: Build Input Context for Address Device command
        // -----------------------------------------------------------------
        auto* inputCtx = (Xhci::InputContext*)Memory::g_pfa->AllocateZeroed();

        // Input Control Context: add Slot Context (bit 0) and EP0 (bit 1)
        inputCtx->ICC.AddFlags = 0x3;

        // Slot Context
        uint32_t speedVal    = SpeedToSlotContextValue(speed);
        uint32_t ctxEntries  = 1;  // Context Entries = 1 (Slot + EP0 only)
        inputCtx->Slot.Field0 = (ctxEntries << 27) | (speedVal << 20);
        inputCtx->Slot.Field1 = ((uint32_t)portId << 16);  // Root Hub Port Number

        // EP0 Context
        // Allocate EP0 transfer ring
        auto* ep0Ring = (Xhci::TRB*)Memory::g_pfa->AllocateZeroed();
        dev->EP0Ring        = ep0Ring;
        dev->EP0RingPhys    = Memory::SubHHDM(ep0Ring);
        dev->EP0RingEnqueue = 0;
        dev->EP0RingCCS     = true;

        // Set up Link TRB at last position to wrap back to start
        // (bit 1 = Toggle Cycle on Link TRBs)
        Xhci::TRB& ep0Link = ep0Ring[Xhci::XFER_RING_SIZE - 1];
        ep0Link.Parameter0 = (uint32_t)(dev->EP0RingPhys & 0xFFFFFFFF);
        ep0Link.Parameter1 = (uint32_t)(dev->EP0RingPhys >> 32);
        ep0Link.Status = 0;
        ep0Link.Control = (Xhci::TRB_LINK << Xhci::TRB_TYPE_SHIFT) | Xhci::TRB_ENT;

        uint16_t maxPacket = MaxPacketSizeForSpeed(speed);

        // Field1: CErr=3 (bits 2:1), EP Type=Control=4 (bits 5:3), Max Packet Size (bits 31:16)
        inputCtx->EP[0].Field1 = (3 << 1)
                                | (Xhci::EP_TYPE_CONTROL << 3)
                                | ((uint32_t)maxPacket << 16);
        // TR Dequeue Pointer with DCS=1
        inputCtx->EP[0].TRDequeuePtr = dev->EP0RingPhys | 1;
        // Average TRB Length = 8
        inputCtx->EP[0].Field2 = 8;

        // -----------------------------------------------------------------
        // Step 4a: Address Device (BSR=1) — initialize slot without SET_ADDRESS
        // -----------------------------------------------------------------
        Xhci::TRB addrTrb = {};
        uint64_t inputCtxPhys = Memory::SubHHDM(inputCtx);
        addrTrb.Parameter0 = (uint32_t)(inputCtxPhys & 0xFFFFFFFF);
        addrTrb.Parameter1 = (uint32_t)(inputCtxPhys >> 32);
        addrTrb.Control    = (Xhci::TRB_ADDRESS_DEVICE << Xhci::TRB_TYPE_SHIFT)
                           | Xhci::TRB_BSR
                           | ((uint32_t)slotId << 24);

        cc = Xhci::SendCommand(addrTrb);
        if (cc != Xhci::CC_SUCCESS) {
            KernelLogStream(ERROR, "USB") << "Address Device (BSR=1) failed, slot="
                << (uint64_t)slotId << " cc=" << (uint64_t)cc;
            dev->Active = false;
            return 0;
        }

        KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId << " initialized (BSR=1)";

        // -----------------------------------------------------------------
        // Step 4b: GET_DESCRIPTOR (Device, 8 bytes) — read bMaxPacketSize0
        // -----------------------------------------------------------------
        uint8_t partialDesc[8] = {};
        cc = Xhci::ControlTransfer(slotId, REQTYPE_DEV_TO_HOST, REQ_GET_DESCRIPTOR,
                                   (DESC_DEVICE << 8), 0, 8,
                                   partialDesc, true);
        if (cc != Xhci::CC_SUCCESS && cc != Xhci::CC_SHORT_PACKET) {
            KernelLogStream(ERROR, "USB") << "GET_DESCRIPTOR(8-byte) failed, cc=" << (uint64_t)cc;
            dev->Active = false;
            return 0;
        }

        uint8_t bMaxPacketSize0 = partialDesc[7];
        if (bMaxPacketSize0 == 0) bMaxPacketSize0 = maxPacket; // fallback

        KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId
            << ": bMaxPacketSize0=" << (uint64_t)bMaxPacketSize0;

        // -----------------------------------------------------------------
        // Step 4c: Evaluate Context — update EP0 max packet size if needed
        // -----------------------------------------------------------------
        if (bMaxPacketSize0 != maxPacket) {
            auto* evalCtx = (Xhci::InputContext*)Memory::g_pfa->AllocateZeroed();

            // Only updating EP0 — set AddFlags bit 1 (EP0), no slot context needed
            evalCtx->ICC.AddFlags = (1 << 1);

            // Copy current EP0 context and update max packet size
            evalCtx->EP[0] = dev->OutputContext->EP[0];
            evalCtx->EP[0].Field1 = (evalCtx->EP[0].Field1 & 0x0000FFFF)
                                  | ((uint32_t)bMaxPacketSize0 << 16);

            Xhci::TRB evalTrb = {};
            uint64_t evalCtxPhys = Memory::SubHHDM(evalCtx);
            evalTrb.Parameter0 = (uint32_t)(evalCtxPhys & 0xFFFFFFFF);
            evalTrb.Parameter1 = (uint32_t)(evalCtxPhys >> 32);
            evalTrb.Control    = (Xhci::TRB_EVALUATE_CONTEXT << Xhci::TRB_TYPE_SHIFT)
                               | ((uint32_t)slotId << 24);

            cc = Xhci::SendCommand(evalTrb);
            if (cc != Xhci::CC_SUCCESS) {
                KernelLogStream(WARNING, "USB") << "Evaluate Context failed, slot="
                    << (uint64_t)slotId << " cc=" << (uint64_t)cc;
                // Non-fatal: continue with original max packet size
            } else {
                KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId
                    << ": EP0 max packet updated to " << (uint64_t)bMaxPacketSize0;
            }
        }

        // -----------------------------------------------------------------
        // Step 4d: Address Device (BSR=0) — actually send SET_ADDRESS
        // -----------------------------------------------------------------
        // Update input context EP0 to current ring position and actual max
        // packet size.  BSR=0 re-initializes the output EP0 context from the
        // input context, so both fields must reflect reality.
        uint64_t curDeq = dev->EP0RingPhys + (uint64_t)dev->EP0RingEnqueue * sizeof(Xhci::TRB);
        if (dev->EP0RingCCS) {
            curDeq |= 1; // DCS bit
        }
        inputCtx->EP[0].TRDequeuePtr = curDeq;
        inputCtx->EP[0].Field1 = (3 << 1)
                                | (Xhci::EP_TYPE_CONTROL << 3)
                                | ((uint32_t)bMaxPacketSize0 << 16);

        Xhci::TRB addrTrb2 = {};
        addrTrb2.Parameter0 = (uint32_t)(inputCtxPhys & 0xFFFFFFFF);
        addrTrb2.Parameter1 = (uint32_t)(inputCtxPhys >> 32);
        addrTrb2.Control    = (Xhci::TRB_ADDRESS_DEVICE << Xhci::TRB_TYPE_SHIFT)
                            | ((uint32_t)slotId << 24);

        cc = Xhci::SendCommand(addrTrb2);
        if (cc != Xhci::CC_SUCCESS) {
            KernelLogStream(ERROR, "USB") << "Address Device failed, slot="
                << (uint64_t)slotId << " cc=" << (uint64_t)cc;
            dev->Active = false;
            return 0;
        }

        // Set-address recovery time (USB spec requires >= 2ms, use 10ms for safety)
        BusyWaitMs(10);

        KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId << " addressed";

        // -----------------------------------------------------------------
        // Step 5: GET_DESCRIPTOR (Device, full 18 bytes)
        // -----------------------------------------------------------------
        DeviceDescriptor devDesc = {};
        cc = Xhci::ControlTransfer(slotId, REQTYPE_DEV_TO_HOST, REQ_GET_DESCRIPTOR,
                                   (DESC_DEVICE << 8), 0, sizeof(DeviceDescriptor),
                                   &devDesc, true);
        if (cc != Xhci::CC_SUCCESS && cc != Xhci::CC_SHORT_PACKET) {
            KernelLogStream(ERROR, "USB") << "GET_DESCRIPTOR(Device) failed, cc=" << (uint64_t)cc;
            dev->Active = false;
            return 0;
        }

        dev->VendorId    = devDesc.idVendor;
        dev->ProductId   = devDesc.idProduct;
        dev->DeviceClass = devDesc.bDeviceClass;

        KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId
            << ": VID:PID = " << base::hex << (uint64_t)devDesc.idVendor
            << ":" << (uint64_t)devDesc.idProduct << base::dec;

        // -----------------------------------------------------------------
        // Step 6: GET_DESCRIPTOR (Configuration) -- header first, then full
        // -----------------------------------------------------------------
        ConfigDescriptor cfgHdr = {};
        cc = Xhci::ControlTransfer(slotId, REQTYPE_DEV_TO_HOST, REQ_GET_DESCRIPTOR,
                                   (DESC_CONFIGURATION << 8), 0, sizeof(ConfigDescriptor),
                                   &cfgHdr, true);
        if (cc != Xhci::CC_SUCCESS && cc != Xhci::CC_SHORT_PACKET) {
            KernelLogStream(ERROR, "USB") << "GET_DESCRIPTOR(Config header) failed, cc=" << (uint64_t)cc;
            dev->Active = false;
            return 0;
        }

        uint8_t cfgBuf[256] = {};
        uint16_t totalLen = cfgHdr.wTotalLength;
        if (totalLen > 256) totalLen = 256;

        cc = Xhci::ControlTransfer(slotId, REQTYPE_DEV_TO_HOST, REQ_GET_DESCRIPTOR,
                                   (DESC_CONFIGURATION << 8), 0, totalLen,
                                   cfgBuf, true);
        if (cc != Xhci::CC_SUCCESS && cc != Xhci::CC_SHORT_PACKET) {
            KernelLogStream(ERROR, "USB") << "GET_DESCRIPTOR(Config full) failed, cc=" << (uint64_t)cc;
            dev->Active = false;
            return 0;
        }

        // -----------------------------------------------------------------
        // Step 7: Parse configuration descriptor blob
        // -----------------------------------------------------------------
        uint16_t offset = 0;
        bool foundHid = false;
        bool foundBt = false;
        bool foundEp = false;
        bool foundBulkIn = false;
        bool foundBulkOut = false;
        uint16_t hidReportDescLen = 0;

        while (offset + 2 <= totalLen) {
            uint8_t len  = cfgBuf[offset];
            uint8_t type = cfgBuf[offset + 1];
            if (len == 0) break;

            if (type == DESC_INTERFACE && offset + sizeof(InterfaceDescriptor) <= totalLen) {
                auto* iface = (InterfaceDescriptor*)&cfgBuf[offset];
                // Reset at each new interface boundary
                foundHid = false;
                foundBt = false;

                if (!foundEp &&
                    iface->bInterfaceClass == CLASS_HID &&
                    iface->bInterfaceSubClass == SUBCLASS_BOOT) {
                    dev->InterfaceClass    = iface->bInterfaceClass;
                    dev->InterfaceSubClass = iface->bInterfaceSubClass;
                    dev->InterfaceProtocol = iface->bInterfaceProtocol;
                    foundHid = true;
                }

                // Bluetooth HCI interface (class 0xE0, subclass 0x01, protocol 0x01)
                if (iface->bInterfaceClass == CLASS_WIRELESS &&
                    iface->bInterfaceSubClass == SUBCLASS_RF &&
                    iface->bInterfaceProtocol == PROTOCOL_BLUETOOTH) {
                    if (dev->InterfaceClass == 0 || dev->InterfaceClass == CLASS_WIRELESS) {
                        dev->InterfaceClass    = iface->bInterfaceClass;
                        dev->InterfaceSubClass = iface->bInterfaceSubClass;
                        dev->InterfaceProtocol = iface->bInterfaceProtocol;
                    }
                    foundBt = true;
                }
            }

            // HID descriptor (0x21): extract report descriptor length
            if (type == DESC_HID && foundHid && !foundEp && len >= 9 && offset + 8 < totalLen) {
                hidReportDescLen = (uint16_t)cfgBuf[offset + 7]
                                 | ((uint16_t)cfgBuf[offset + 8] << 8);
            }

            if (type == DESC_ENDPOINT && offset + sizeof(EndpointDescriptor) <= totalLen) {
                auto* ep = (EndpointDescriptor*)&cfgBuf[offset];
                uint8_t xferType = ep->bmAttributes & EP_XFER_TYPE_MASK;
                bool isIn = (ep->bEndpointAddress & EP_DIR_IN) != 0;

                // HID interrupt IN endpoint
                if (foundHid && !foundEp && isIn && xferType == EP_XFER_INTERRUPT) {
                    dev->InterruptEpNum     = ep->bEndpointAddress & 0x0F;
                    dev->InterruptMaxPacket = ep->wMaxPacketSize & 0x7FF;
                    dev->InterruptInterval  = ep->bInterval;
                    foundEp = true;
                }

                // Bluetooth endpoints
                if (foundBt) {
                    if (isIn && xferType == EP_XFER_INTERRUPT && !foundEp) {
                        // HCI event pipe (interrupt IN)
                        dev->InterruptEpNum     = ep->bEndpointAddress & 0x0F;
                        dev->InterruptMaxPacket = ep->wMaxPacketSize & 0x7FF;
                        dev->InterruptInterval  = ep->bInterval;
                        foundEp = true;
                    } else if (isIn && xferType == EP_XFER_BULK && !foundBulkIn) {
                        // ACL data IN pipe (bulk IN)
                        dev->BulkInEpNum     = ep->bEndpointAddress & 0x0F;
                        dev->BulkInMaxPacket = ep->wMaxPacketSize & 0x7FF;
                        foundBulkIn = true;
                    } else if (!isIn && xferType == EP_XFER_BULK && !foundBulkOut) {
                        // ACL data OUT pipe (bulk OUT)
                        dev->BulkOutEpNum     = ep->bEndpointAddress & 0x0F;
                        dev->BulkOutMaxPacket = ep->wMaxPacketSize & 0x7FF;
                        foundBulkOut = true;
                    }
                }
            }

            offset += len;
        }

        // For Bluetooth devices, also check device class for correct identification
        // Some BT adapters use bDeviceClass=0xE0 at device level
        if (!foundBt && devDesc.bDeviceClass == CLASS_WIRELESS &&
            devDesc.bDeviceSubClass == SUBCLASS_RF &&
            devDesc.bDeviceProtocol == PROTOCOL_BLUETOOTH) {
            dev->InterfaceClass    = CLASS_WIRELESS;
            dev->InterfaceSubClass = SUBCLASS_RF;
            dev->InterfaceProtocol = PROTOCOL_BLUETOOTH;
            foundBt = true;
        }

        // -----------------------------------------------------------------
        // Step 8: SET_CONFIGURATION
        // -----------------------------------------------------------------
        cc = Xhci::ControlTransfer(slotId, REQTYPE_HOST_TO_DEV, REQ_SET_CONFIGURATION,
                                   cfgHdr.bConfigurationValue, 0, 0, nullptr, false);
        if (cc != Xhci::CC_SUCCESS) {
            KernelLogStream(ERROR, "USB") << "SET_CONFIGURATION failed, cc=" << (uint64_t)cc;
            dev->Active = false;
            return 0;
        }

        // -----------------------------------------------------------------
        // Step 9: Configure Endpoints
        // -----------------------------------------------------------------
        if (foundEp || foundBulkIn || foundBulkOut) {
            auto* inputCtx2 = (Xhci::InputContext*)Memory::g_pfa->AllocateZeroed();

            // Start with slot context
            inputCtx2->ICC.AddFlags = (1 << 0);

            // Copy the current slot context from the output context
            inputCtx2->Slot = dev->OutputContext->Slot;

            // Track the highest DCI to set Context Entries
            uint32_t maxDci = 0;

            // --- Configure Interrupt IN endpoint ---
            if (foundEp) {
                uint8_t dci = dev->InterruptEpNum * 2 + 1;
                inputCtx2->ICC.AddFlags |= (1 << dci);
                if (dci > maxDci) maxDci = dci;

                // Allocate interrupt transfer ring
                auto* intRing = (Xhci::TRB*)Memory::g_pfa->AllocateZeroed();
                dev->InterruptRing        = intRing;
                dev->InterruptRingPhys    = Memory::SubHHDM(intRing);
                dev->InterruptRingEnqueue = 0;
                dev->InterruptRingCCS     = true;

                Xhci::TRB& intLink = intRing[Xhci::XFER_RING_SIZE - 1];
                intLink.Parameter0 = (uint32_t)(dev->InterruptRingPhys & 0xFFFFFFFF);
                intLink.Parameter1 = (uint32_t)(dev->InterruptRingPhys >> 32);
                intLink.Status = 0;
                intLink.Control = (Xhci::TRB_LINK << Xhci::TRB_TYPE_SHIFT) | Xhci::TRB_ENT;

                auto& epCtx = inputCtx2->EP[dci - 1];
                uint32_t xhciInterval = ConvertInterval(speed, dev->InterruptInterval);
                epCtx.Field0 = (xhciInterval << 16);
                epCtx.Field1 = (3 << 1)
                             | (Xhci::EP_TYPE_INTERRUPT_IN << 3)
                             | ((uint32_t)dev->InterruptMaxPacket << 16);
                epCtx.TRDequeuePtr = dev->InterruptRingPhys | 1;
                epCtx.Field2 = dev->InterruptMaxPacket;
            }

            // --- Configure Bulk IN endpoint ---
            if (foundBulkIn) {
                uint8_t dci = dev->BulkInEpNum * 2 + 1;
                inputCtx2->ICC.AddFlags |= (1 << dci);
                if (dci > maxDci) maxDci = dci;

                auto* bulkInRing = (Xhci::TRB*)Memory::g_pfa->AllocateZeroed();
                dev->BulkInRing        = bulkInRing;
                dev->BulkInRingPhys    = Memory::SubHHDM(bulkInRing);
                dev->BulkInRingEnqueue = 0;
                dev->BulkInRingCCS     = true;

                Xhci::TRB& biLink = bulkInRing[Xhci::XFER_RING_SIZE - 1];
                biLink.Parameter0 = (uint32_t)(dev->BulkInRingPhys & 0xFFFFFFFF);
                biLink.Parameter1 = (uint32_t)(dev->BulkInRingPhys >> 32);
                biLink.Status = 0;
                biLink.Control = (Xhci::TRB_LINK << Xhci::TRB_TYPE_SHIFT) | Xhci::TRB_ENT;

                auto& epCtx = inputCtx2->EP[dci - 1];
                epCtx.Field0 = 0;
                epCtx.Field1 = (3 << 1)
                             | (Xhci::EP_TYPE_BULK_IN << 3)
                             | ((uint32_t)dev->BulkInMaxPacket << 16);
                epCtx.TRDequeuePtr = dev->BulkInRingPhys | 1;
                epCtx.Field2 = dev->BulkInMaxPacket;
            }

            // --- Configure Bulk OUT endpoint ---
            if (foundBulkOut) {
                uint8_t dci = dev->BulkOutEpNum * 2;
                inputCtx2->ICC.AddFlags |= (1 << dci);
                if (dci > maxDci) maxDci = dci;

                auto* bulkOutRing = (Xhci::TRB*)Memory::g_pfa->AllocateZeroed();
                dev->BulkOutRing        = bulkOutRing;
                dev->BulkOutRingPhys    = Memory::SubHHDM(bulkOutRing);
                dev->BulkOutRingEnqueue = 0;
                dev->BulkOutRingCCS     = true;

                Xhci::TRB& boLink = bulkOutRing[Xhci::XFER_RING_SIZE - 1];
                boLink.Parameter0 = (uint32_t)(dev->BulkOutRingPhys & 0xFFFFFFFF);
                boLink.Parameter1 = (uint32_t)(dev->BulkOutRingPhys >> 32);
                boLink.Status = 0;
                boLink.Control = (Xhci::TRB_LINK << Xhci::TRB_TYPE_SHIFT) | Xhci::TRB_ENT;

                auto& epCtx = inputCtx2->EP[dci - 1];
                epCtx.Field0 = 0;
                epCtx.Field1 = (3 << 1)
                             | (Xhci::EP_TYPE_BULK_OUT << 3)
                             | ((uint32_t)dev->BulkOutMaxPacket << 16);
                epCtx.TRDequeuePtr = dev->BulkOutRingPhys | 1;
                epCtx.Field2 = dev->BulkOutMaxPacket;
            }

            // Update Context Entries to cover the highest DCI
            inputCtx2->Slot.Field0 = (inputCtx2->Slot.Field0 & ~(0x1Fu << 27))
                                   | (maxDci << 27);

            // Send Configure Endpoint command
            Xhci::TRB cfgTrb = {};
            uint64_t inputCtx2Phys = Memory::SubHHDM(inputCtx2);
            cfgTrb.Parameter0 = (uint32_t)(inputCtx2Phys & 0xFFFFFFFF);
            cfgTrb.Parameter1 = (uint32_t)(inputCtx2Phys >> 32);
            cfgTrb.Control    = (Xhci::TRB_CONFIGURE_ENDPOINT << Xhci::TRB_TYPE_SHIFT)
                              | ((uint32_t)slotId << 24);

            cc = Xhci::SendCommand(cfgTrb);
            if (cc != Xhci::CC_SUCCESS) {
                KernelLogStream(ERROR, "USB") << "Configure Endpoint failed, slot="
                    << (uint64_t)slotId << " cc=" << (uint64_t)cc;
                dev->Active = false;
                return 0;
            }

            if (foundEp) {
                KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId
                    << ": Interrupt EP " << (uint64_t)dev->InterruptEpNum << " configured";
            }
            if (foundBulkIn) {
                KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId
                    << ": Bulk IN EP " << (uint64_t)dev->BulkInEpNum << " configured";
            }
            if (foundBulkOut) {
                KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId
                    << ": Bulk OUT EP " << (uint64_t)dev->BulkOutEpNum << " configured";
            }
        }

        // -----------------------------------------------------------------
        // Step 10: SET_PROTOCOL -- Boot Protocol for keyboards only
        // -----------------------------------------------------------------
        // Set Boot Protocol for keyboards only.
        // Mice stay in Report Protocol (the default) for scroll wheel support;
        // HidMouse parses the HID Report Descriptor to handle variable formats.
        if (foundEp && dev->InterfaceClass == CLASS_HID && dev->InterfaceProtocol == PROTOCOL_KEYBOARD) {
            cc = Xhci::ControlTransfer(slotId, REQTYPE_CLASS_IFACE, REQ_SET_PROTOCOL,
                                       0, 0, 0, nullptr, false);
            if (cc != Xhci::CC_SUCCESS) {
                KernelLogStream(WARNING, "USB") << "SET_PROTOCOL(Boot) failed, cc=" << (uint64_t)cc;
                // Non-fatal: some devices only support boot protocol anyway
            }
        }

        // -----------------------------------------------------------------
        // Step 10b: Fetch HID Report Descriptor for mice
        // -----------------------------------------------------------------
        if (foundEp && dev->InterfaceProtocol == PROTOCOL_MOUSE && hidReportDescLen > 0) {
            uint8_t rdBuf[256] = {};
            uint16_t rdLen = hidReportDescLen;
            if (rdLen > 256) rdLen = 256;

            cc = Xhci::ControlTransfer(slotId, REQTYPE_STD_IFACE_IN, REQ_GET_DESCRIPTOR,
                                       (DESC_HID_REPORT << 8), 0, rdLen,
                                       rdBuf, true);
            if (cc == Xhci::CC_SUCCESS || cc == Xhci::CC_SHORT_PACKET) {
                HidMouse::ParseReportDescriptor(rdBuf, rdLen);
            } else {
                KernelLogStream(WARNING, "USB") << "GET_DESCRIPTOR(HID Report) failed, cc=" << (uint64_t)cc;
            }
        }

        // -----------------------------------------------------------------
        // Step 11: SET_IDLE(0) -- only report on changes (no idle reports)
        // -----------------------------------------------------------------
        if (foundEp && dev->InterfaceClass == CLASS_HID && dev->InterfaceProtocol == PROTOCOL_KEYBOARD) {
            // wValue upper byte = duration (0 = indefinite), lower byte = report ID
            cc = Xhci::ControlTransfer(slotId, REQTYPE_CLASS_IFACE, REQ_SET_IDLE,
                                       (0 << 8), 0, 0, nullptr, false);
            if (cc != Xhci::CC_SUCCESS) {
                KernelLogStream(WARNING, "USB") << "SET_IDLE(0) failed, cc=" << (uint64_t)cc;
                // Non-fatal: not all devices support SET_IDLE
            }
        }

        // -----------------------------------------------------------------
        // Step 12: Queue first interrupt transfer (HID only)
        // Bluetooth manages its own interrupt/bulk transfers via StartEventPipe()
        // -----------------------------------------------------------------
        if (foundEp && !foundBt) {
            Xhci::QueueInterruptTransfer(slotId);
        }

        // -----------------------------------------------------------------
        // Step 13: Register with the appropriate class driver
        // -----------------------------------------------------------------
        if (dev->InterfaceClass == CLASS_HID && dev->InterfaceProtocol == PROTOCOL_KEYBOARD) {
            HidKeyboard::RegisterDevice(slotId);
            KernelLogStream(OK, "USB") << "Slot " << (uint64_t)slotId << ": HID Boot Keyboard";
        } else if (dev->InterfaceClass == CLASS_HID && dev->InterfaceProtocol == PROTOCOL_MOUSE) {
            HidMouse::RegisterDevice(slotId);
            KernelLogStream(OK, "USB") << "Slot " << (uint64_t)slotId << ": HID Boot Mouse";
        } else if (dev->InterfaceClass == CLASS_WIRELESS &&
                   dev->InterfaceSubClass == SUBCLASS_RF &&
                   dev->InterfaceProtocol == PROTOCOL_BLUETOOTH) {
            Bluetooth::RegisterAdapter(slotId);
            KernelLogStream(OK, "USB") << "Slot " << (uint64_t)slotId << ": Bluetooth Adapter"
                << " VID:" << base::hex << (uint64_t)dev->VendorId
                << " PID:" << (uint64_t)dev->ProductId << base::dec;
        } else if (foundEp) {
            KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId
                << ": USB device, class=" << (uint64_t)dev->InterfaceClass
                << " protocol=" << (uint64_t)dev->InterfaceProtocol;
        } else {
            KernelLogStream(INFO, "USB") << "Slot " << (uint64_t)slotId
                << ": Non-HID device, class=" << (uint64_t)devDesc.bDeviceClass;
        }

        return slotId;
    }

}
