nt-lite/code/drivers/bus/usb/usb_xhci.c

// xHCI USB host-controller driver. This file owns MMIO setup, rings, slots, endpoints, and transfer completion.

#include "usb.h"

// Implements the align Up32 operation.
uint32_t UsbAlignUp32(uint32_t value, uint32_t align) {
    return (value + align - 1u) & ~(align - 1u);
}

// Implements the barrier operation.
void XhciBarrier(void) {
    _ReadWriteBarrier();
    _mm_mfence();
    _ReadWriteBarrier();
}

// Implements the read8 operation.
uint8_t XhciRead8(volatile uint8_t* base, uint32_t reg) {
    return *(volatile uint8_t*)(base + reg);
}

// Implements the read16 operation.
uint16_t XhciRead16(volatile uint8_t* base, uint32_t reg) {
    return *(volatile uint16_t*)(base + reg);
}

// Implements the read32 operation.
uint32_t XhciRead32(volatile uint8_t* base, uint32_t reg) {
    return *(volatile uint32_t*)(base + reg);
}

// Implements the write32 operation.
void XhciWrite32(volatile uint8_t* base, uint32_t reg, uint32_t value) {
    *(volatile uint32_t*)(base + reg) = value;
}

// Implements the read64 operation.
uint64_t XhciRead64(volatile uint8_t* base, uint32_t reg) {
    uint32_t lo = XhciRead32(base, reg);
    uint32_t hi = XhciRead32(base, reg + 4u);
    return ((uint64_t)hi << 32) | lo;
}

// Implements the write64 operation.
void XhciWrite64(volatile uint8_t* base, uint32_t reg, uint64_t value) {
    XhciWrite32(base, reg, (uint32_t)value);
    XhciWrite32(base, reg + 4u, (uint32_t)(value >> 32));
}

// Implements the delay Us operation.
void UsbDelayUs(uint32_t us) {
    KeStallExecution(us);
}

bool UsbWaitSet32(volatile uint8_t* base, uint32_t reg, uint32_t mask, uint32_t timeoutUs) {
    for (uint32_t i = 0; i < timeoutUs; i++) {
        if ((XhciRead32(base, reg) & mask) == mask) return true;
        UsbDelayUs(1);
    }
    return false;
}

bool UsbWaitClear32(volatile uint8_t* base, uint32_t reg, uint32_t mask, uint32_t timeoutUs) {
    for (uint32_t i = 0; i < timeoutUs; i++) {
        if ((XhciRead32(base, reg) & mask) == 0) return true;
        UsbDelayUs(1);
    }
    return false;
}

// Allocates page-aligned common-buffer memory from the controller adapter.
void* UsbDmaAllocPages(XHCI_CONTROLLER* xhci,
                       size_t pages,
                       uint64_t* outDeviceAddress,
                       IO_DMA_BUFFER** outBuffer) {
    if (!xhci || !xhci->DmaAdapter || !outDeviceAddress || !outBuffer || pages == 0 ||
        pages > ((size_t)-1) / PAGE_SIZE)
        return (void*)0;

    IO_DMA_BUFFER* buffer = (IO_DMA_BUFFER*)0;
    KSTATUS status = IoDmaAllocateCommonBuffer(xhci->DmaAdapter,
                                                           pages * (size_t)PAGE_SIZE,
                                                           PAGE_SIZE,
                                                           &buffer);
    if (!KSUCCESS(status)) {
        DbgPrint("usb: DMA allocation failed pages=%llu st=0x%x\n",
                 (unsigned long long)pages, status);
        return (void*)0;
    }

    *outDeviceAddress = IoDmaGetCommonBufferDeviceAddress(buffer);
    *outBuffer = buffer;
    return IoDmaGetCommonBufferAddress(buffer);
}

// Releases all DMA allocations owned by one xHCI slot.
void XhciReleaseDeviceDma(XHCI_DEVICE* dev) {
    if (!dev)
        return;

    for (uint32_t endpoint = 1; endpoint < 32; endpoint++) {
        USB_PIPE* pipe = dev->EndpointPipes[endpoint];
        dev->EndpointPipes[endpoint] = (USB_PIPE*)0;
        if (pipe)
            UsbDereferencePipe(pipe);
    }

    IoDmaFreeCommonBuffer(dev->MassBufferDma);
    IoDmaFreeCommonBuffer(dev->HidBufferDma);
    IoDmaFreeCommonBuffer(dev->DescriptorDma);
    IoDmaFreeCommonBuffer(dev->MassOutRingDma);
    IoDmaFreeCommonBuffer(dev->MassInRingDma);
    IoDmaFreeCommonBuffer(dev->HidRingDma);
    IoDmaFreeCommonBuffer(dev->Ep0RingDma);
    IoDmaFreeCommonBuffer(dev->DeviceContextDma);
    IoDmaFreeCommonBuffer(dev->InputContextDma);

    dev->MassBufferDma = (IO_DMA_BUFFER*)0;
    dev->HidBufferDma = (IO_DMA_BUFFER*)0;
    dev->DescriptorDma = (IO_DMA_BUFFER*)0;
    dev->MassOutRingDma = (IO_DMA_BUFFER*)0;
    dev->MassInRingDma = (IO_DMA_BUFFER*)0;
    dev->HidRingDma = (IO_DMA_BUFFER*)0;
    dev->Ep0RingDma = (IO_DMA_BUFFER*)0;
    dev->DeviceContextDma = (IO_DMA_BUFFER*)0;
    dev->InputContextDma = (IO_DMA_BUFFER*)0;

    dev->MassBuffer = (uint8_t*)0;
    dev->HidBuffer = (uint8_t*)0;
    dev->DescriptorBuffer = (uint8_t*)0;
    dev->MassOutRing = (XHCI_TRB*)0;
    dev->MassInRing = (XHCI_TRB*)0;
    dev->HidRing = (XHCI_TRB*)0;
    dev->Ep0Ring = (XHCI_TRB*)0;
    dev->DeviceContext = (uint8_t*)0;
    dev->InputContext = (uint8_t*)0;
}

// Releases controller-wide xHCI DMA allocations.
static void XhciReleaseControllerDma(XHCI_CONTROLLER* xhci) {
    if (!xhci)
        return;

    for (uint32_t slot = 1; slot <= XHCI_MAX_TRACKED_SLOTS; slot++)
        XhciReleaseDeviceDma(&xhci->Devices[slot]);
    for (uint16_t i = 0; i < XHCI_MAX_SCRATCHPADS; i++) {
        IoDmaFreeCommonBuffer(xhci->ScratchpadDma[i]);
        xhci->ScratchpadDma[i] = (IO_DMA_BUFFER*)0;
        xhci->ScratchpadDeviceAddress[i] = 0;
    }

    IoDmaFreeCommonBuffer(xhci->ScratchpadArrayDma);
    IoDmaFreeCommonBuffer(xhci->ErstDma);
    IoDmaFreeCommonBuffer(xhci->EventRingDma);
    IoDmaFreeCommonBuffer(xhci->CommandRingDma);
    IoDmaFreeCommonBuffer(xhci->DcbaaDma);
    xhci->ScratchpadArrayDma = (IO_DMA_BUFFER*)0;
    xhci->ErstDma = (IO_DMA_BUFFER*)0;
    xhci->EventRingDma = (IO_DMA_BUFFER*)0;
    xhci->CommandRingDma = (IO_DMA_BUFFER*)0;
    xhci->DcbaaDma = (IO_DMA_BUFFER*)0;
    xhci->ScratchpadArray = (uint64_t*)0;
    xhci->Erst = (XHCI_ERST_ENTRY*)0;
    xhci->EventRing = (XHCI_TRB*)0;
    xhci->CommandRing = (XHCI_TRB*)0;
    xhci->Dcbaa = (uint64_t*)0;

    IoDmaPutAdapter(xhci->DmaAdapter);
    xhci->DmaAdapter = (IO_DMA_ADAPTER*)0;
}

// Implements the scratchpad Count operation.
uint16_t XhciScratchpadCount(uint32_t hcs2) {
    return (uint16_t)((((hcs2 >> 16) & 0x3E0u) | ((hcs2 >> 27) & 0x1Fu)) & 0x3FFu);
}

// Implements the trb Type operation.
uint32_t XhciTrbType(const XHCI_TRB* trb) {
    return (trb->Control >> XHCI_TRB_TYPE_SHIFT) & XHCI_TRB_TYPE_MASK;
}

// Implements the completion Code operation.
uint32_t XhciCompletionCode(const XHCI_TRB* trb) {
    return (trb->Status >> 24) & 0xFFu;
}

// Implements the event Slot Id operation.
uint8_t XhciEventSlotId(const XHCI_TRB* trb) {
    return (uint8_t)(trb->Control >> 24);
}

// Implements the event Endpoint Id operation.
uint8_t XhciEventEndpointId(const XHCI_TRB* trb) {
    return (uint8_t)((trb->Control >> 16) & 0x1Fu);
}

// Implements the transfer Residue operation.
uint32_t XhciTransferResidue(const XHCI_TRB* trb) {
    return trb->Status & 0xFFFFFFu;
}

// Implements the input Control Context operation.
uint32_t* XhciInputControlContext(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    (void)xhci;
    return (uint32_t*)dev->InputContext;
}

// Implements the input Slot Context operation.
uint32_t* XhciInputSlotContext(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    return (uint32_t*)(dev->InputContext + xhci->ContextBytes);
}

uint32_t* XhciInputEndpointContext(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                          uint8_t endpointId) {
    return (uint32_t*)(dev->InputContext + (uint32_t)(endpointId + 1u) * xhci->ContextBytes);
}

// Reconstructs the addressed device topology when configuring an endpoint or hub.
// Configure Endpoint must not discard route, TT, or hub attributes on the slot.
void XhciBuildConfigureSlotContext(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                   uint32_t* slot, uint8_t endpointId) {
    if (!xhci || !dev || !slot)
        return;
    memset(slot, 0, xhci->ContextBytes);

    uint8_t lastContext = endpointId;
    if (dev->DeviceContext) {
        uint32_t previous = ((const uint32_t*)dev->DeviceContext)[0];
        uint8_t previousLast = (uint8_t)((previous >> 27) & 31u);
        if (previousLast > lastContext)
            lastContext = previousLast;
    }
    uint8_t rootPort = dev->Path.RootPort ? dev->Path.RootPort : dev->Port;
    slot[0] = XHCI_CONTEXT_ENTRIES(lastContext) | XHCI_SLOT_SPEED(dev->Speed) |
              (dev->Path.RouteString & 0xfffffu);
    slot[1] = XHCI_SLOT_ROOT_PORT(rootPort);
    if (dev->Path.BehindHub && dev->Speed <= 2u && dev->Path.ParentHubSlot) {
        slot[2] = (uint32_t)dev->Path.ParentHubSlot |
                  ((uint32_t)dev->Path.ParentHubPort << 8);
    }
    if (dev->HubPortCount) {
        slot[0] |= XHCI_SLOT_HUB;
        slot[1] |= XHCI_SLOT_PORT_COUNT(dev->HubPortCount);
        if (dev->Speed == 3u)
            slot[2] |= XHCI_SLOT_TT_THINK(dev->HubCharacteristics >> 5);
    }
}

// Publishes an external hub's topology to xHCI before its child ports are used.
KSTATUS XhciConfigureHubContext(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                uint8_t portCount, uint16_t characteristics) {
    if (!xhci || !dev || !dev->Addressed || !dev->InputContext ||
        !dev->DeviceContext || !portCount)
        return NTLSTATUS_INVALID_PARAMETER;

    uint8_t oldPortCount = dev->HubPortCount;
    uint16_t oldCharacteristics = dev->HubCharacteristics;
    dev->HubPortCount = portCount;
    dev->HubCharacteristics = characteristics;

    memset(dev->InputContext, 0, PAGE_SIZE);
    uint32_t* icc = XhciInputControlContext(xhci, dev);
    uint32_t* slot = XhciInputSlotContext(xhci, dev);
    // Update only the slot context; existing endpoint contexts remain active.
    icc[1] = 1u;
    XhciBuildConfigureSlotContext(xhci, dev, slot, 1u);

    XHCI_COMMAND_RESULT result;
    memset(&result, 0, sizeof(result));
    DbgPrint("usb-stage: hcd=%02x:%02x.%u hub-slot=%u configure-context begin ports=%u tt=%u speed=%u\n",
             xhci->Bus, xhci->Dev, xhci->Fn, dev->SlotId,
             portCount, (characteristics >> 5) & 3u, dev->Speed);
    KSTATUS st = XhciSubmitCommand(xhci, dev->InputContextDeviceAddress, 0,
                                   (XHCI_TRB_TYPE_CONFIGURE_ENDPOINT << XHCI_TRB_TYPE_SHIFT) |
                                       ((uint32_t)dev->SlotId << 24),
                                   &result);
    if (!KSUCCESS(st) || result.CompletionCode != XHCI_CC_SUCCESS) {
        DbgPrint("usb-stage: hub-slot=%u configure-context failed st=0x%x cc=%u\n",
                 dev->SlotId, st, result.CompletionCode);
        dev->HubPortCount = oldPortCount;
        dev->HubCharacteristics = oldCharacteristics;
        return KSUCCESS(st) ? NTLSTATUS_DEVICE_ERROR : st;
    }
    DbgPrint("usb-stage: hub-slot=%u configure-context done\n", dev->SlotId);
    return NTLSTATUS_SUCCESS;
}

// Implements the ep Max Packet operation.
uint32_t XhciEpMaxPacket(uint8_t speed) {
    switch (speed) {
        case 2: return 8;    // low-speed
        case 3: return 64;   // high-speed
        case 4: return 512;  // SuperSpeed
        case 5: return 512;  // SuperSpeed+
        case 1:
        default: return 8;   // full-speed before descriptor refinement
    }
}

// Implements the setup Packet To U64 operation.
uint64_t UsbSetupPacketToU64(const USB_SETUP_PACKET* setup) {
    uint64_t v = 0;
    v |= (uint64_t)setup->bmRequestType;
    v |= (uint64_t)setup->bRequest << 8;
    v |= (uint64_t)setup->wValue << 16;
    v |= (uint64_t)setup->wIndex << 32;
    v |= (uint64_t)setup->wLength << 48;
    return v;
}

const char* XhciPortSpeedName(uint32_t speed) {
    switch (speed) {
        case 1: return "full";
        case 2: return "low";
        case 3: return "high";
        case 4: return "super";
        case 5: return "super+";
        default: return "unknown";
    }
}

// Advances an xHCI extended-capability offset using the relative DWORD link.
static uint32_t XhciNextExtendedCapabilityDword(uint32_t current, uint8_t next,
                                               uint32_t mappedBytes) {
    if (!next) return 0;
    uint64_t following = (uint64_t)current + next;
    if (following * 4u + sizeof(uint32_t) > mappedBytes) return 0;
    return (uint32_t)following;
}

// Parses supported Protocols.
void XhciParseSupportedProtocols(XHCI_CONTROLLER* xhci) {
    memset(xhci->PortProtocolMajor, 0, sizeof(xhci->PortProtocolMajor));
    memset(xhci->PortProtocolMinor, 0, sizeof(xhci->PortProtocolMinor));

    uint32_t xecp = (xhci->Hcc1 >> 16) & 0xFFFFu;
    for (uint32_t guard = 0; xecp && guard < 64; guard++) {
        uint32_t off = xecp << 2;
        if (off + 12u > xhci->MmioBytes)
            break;

        uint32_t cap = XhciRead32(xhci->Mmio, off);
        uint8_t id = (uint8_t)(cap & 0xFFu);
        uint8_t next = (uint8_t)((cap >> 8) & 0xFFu);

        if (id == XHCI_EXT_CAP_PROTOCOL) {
            uint8_t minor = (uint8_t)((cap >> 16) & 0xFFu);
            uint8_t major = (uint8_t)((cap >> 24) & 0xFFu);
            uint32_t name = XhciRead32(xhci->Mmio, off + 4u);
            uint32_t ports = XhciRead32(xhci->Mmio, off + 8u);
            uint8_t portOffset = (uint8_t)(ports & 0xFFu);
            uint8_t portCount = (uint8_t)((ports >> 8) & 0xFFu);

            UsbTrace("usb: xHCI protocol cap +0x%x name='%c%c%c%c' usb=%u.%u ports=%u..%u\n",
                     off,
                     (char)(name & 0xFFu),
                     (char)((name >> 8) & 0xFFu),
                     (char)((name >> 16) & 0xFFu),
                     (char)((name >> 24) & 0xFFu),
                     major, minor, portOffset,
                     portCount ? (uint8_t)(portOffset + portCount - 1u) : portOffset);

            if (portOffset != 0 && portCount != 0) {
                for (uint32_t i = 0; i < portCount; i++) {
                    uint32_t port = (uint32_t)portOffset + i;
                    if (port <= XHCI_MAX_ROOT_PORTS) {
                        xhci->PortProtocolMajor[port] = major;
                        xhci->PortProtocolMinor[port] = minor;
                    }
                }
            }
        }

        xecp = XhciNextExtendedCapabilityDword(xecp, next, xhci->MmioBytes);
    }

    uint32_t usb2Ports = 0;
    uint32_t usb3Ports = 0;
    for (uint32_t port = 1; port <= xhci->MaxPorts; port++) {
        if (xhci->PortProtocolMajor[port] == 2u) usb2Ports++;
        else if (xhci->PortProtocolMajor[port] >= 3u) usb3Ports++;
    }
    DbgPrint("usb: xHCI hcd=%02x:%02x.%u protocol ports usb2=%u usb3=%u unknown=%u\n",
             xhci->Bus, xhci->Dev, xhci->Fn, usb2Ports, usb3Ports,
             (uint32_t)xhci->MaxPorts - usb2Ports - usb3Ports);
}

// Implements the port Protocol Major operation.
uint8_t XhciPortProtocolMajor(XHCI_CONTROLLER* xhci, uint8_t port, uint32_t portsc) {
    if (port <= XHCI_MAX_ROOT_PORTS && xhci->PortProtocolMajor[port])
        return xhci->PortProtocolMajor[port];

    uint32_t speed = (portsc & XHCI_PORTSC_SPEED_MASK) >> XHCI_PORTSC_SPEED_SHIFT;
    if (speed >= 4u)
        return 3u;
    if (speed >= 1u && speed <= 3u)
        return 2u;
    return 0u;
}

const char* UsbEndpointTypeName(uint8_t type) {
    switch (type & USB_ENDPOINT_ATTR_TYPE_MASK) {
        case USB_ENDPOINT_TYPE_CONTROL: return "control";
        case USB_ENDPOINT_TYPE_ISOCH: return "isochronous";
        case USB_ENDPOINT_TYPE_BULK: return "bulk";
        case USB_ENDPOINT_TYPE_INTERRUPT: return "interrupt";
        default: return "unknown";
    }
}

const char* UsbDirectionName(uint8_t endpointAddress) {
    return (endpointAddress & USB_ENDPOINT_DIR_IN) ? "in" : "out";
}

const char* UsbInterfaceClassName(uint8_t klass, uint8_t subclass, uint8_t proto) {
    if (klass == USB_CLASS_HID) {
        if (subclass == 1 && proto == 1) return "hid-boot-keyboard";
        if (subclass == 1 && proto == 2) return "hid-boot-mouse";
        return "hid";
    }
    if (klass == USB_CLASS_MASS_STORAGE) {
        if (subclass == USB_MASS_SUBCLASS_SCSI && proto == USB_MASS_PROTOCOL_BOT)
            return "mass-storage-bot";
        return "mass-storage";
    }
    return "class";
}

// Implements the log Port Status operation.
void XhciLogPortStatus(XHCI_CONTROLLER* xhci) {
    for (uint8_t port = 1; port <= xhci->MaxPorts; port++) {
        uint32_t off = XHCI_OP_PORT_BASE + ((uint32_t)port - 1u) * XHCI_OP_PORT_STRIDE;
        uint32_t portsc = XhciRead32(xhci->Op, off);
        uint32_t speed = (portsc & XHCI_PORTSC_SPEED_MASK) >> XHCI_PORTSC_SPEED_SHIFT;

        if ((portsc & (XHCI_PORTSC_CCS | XHCI_PORTSC_CHANGE_MASK)) || port <= 8u) {
            uint8_t proto = XhciPortProtocolMajor(xhci, port, portsc);
            UsbTrace("usb: xHCI port %u: portsc=%08x ccs=%u ped=%u pp=%u proto=usb%u speed=%s(%u)%s%s\n",
                     port, portsc,
                     (portsc & XHCI_PORTSC_CCS) ? 1u : 0u,
                     (portsc & XHCI_PORTSC_PED) ? 1u : 0u,
                     (portsc & XHCI_PORTSC_PP) ? 1u : 0u,
                     proto,
                     XhciPortSpeedName(speed), speed,
                     (portsc & XHCI_PORTSC_OCA) ? " over-current" : "",
                     (portsc & XHCI_PORTSC_PR) ? " resetting" : "");
        }
    }
}

// Implements the power Root Ports operation.
void XhciPowerRootPorts(XHCI_CONTROLLER* xhci) {
    if (!xhci || !xhci->Op)
        return;

    uint32_t powered = 0;
    for (uint8_t port = 1; port <= xhci->MaxPorts; port++) {
        uint32_t off = XhciPortReg(port);
        uint32_t portsc = XhciRead32(xhci->Op, off);
        uint32_t changes = portsc & XHCI_PORTSC_CHANGE_MASK;

        if (!(portsc & XHCI_PORTSC_PP)) {
            XhciWrite32(xhci->Op, off, XHCI_PORTSC_PP | changes);
            powered++;
            UsbTrace("usb: xHCI powered root port %u old-portsc=%08x\n", port, portsc);
        } else if (changes) {
            XhciWrite32(xhci->Op, off, XHCI_PORTSC_PP | changes);
        }
    }

    if (powered) {
        UsbDelayUs(20000u);
        DbgPrint("usb: xHCI powered %u previously-off root port(s)\n", powered);
    }
}

// Implements the ack Interrupter0 operation.
void XhciAckInterrupter0(XHCI_CONTROLLER* xhci) {
    if (!xhci->Runtime) return;
    volatile uint8_t* ir0 = xhci->Runtime + 0x20u;
    uint32_t iman = XhciRead32(ir0, 0x00u);
    XhciWrite32(ir0, 0x00u, (iman & XHCI_IMAN_IE) | XHCI_IMAN_IP);
}

// Updates erdp.
void XhciUpdateErdp(XHCI_CONTROLLER* xhci) {
    volatile uint8_t* ir0 = xhci->Runtime + 0x20u;
    uint64_t erdp = xhci->EventRingDeviceAddress + (uint64_t)xhci->EventDequeue * sizeof(XHCI_TRB);
    XhciWrite64(ir0, 0x18u, erdp | (1ull << 3)); // EHB clear
}

// Returns command Waiter Locked.
static XHCI_WAIT_TARGET* XhciFindCommandWaiterLocked(XHCI_CONTROLLER* xhci, uint64_t commandDeviceAddress) {
    for (XHCI_WAIT_TARGET* w = xhci->WaiterList; w; w = w->Next) {
        if (w->CommandDeviceAddress && w->CommandDeviceAddress == commandDeviceAddress)
            return w;
    }
    return (XHCI_WAIT_TARGET*)0;
}

static XHCI_WAIT_TARGET* XhciFindTransferWaiterLocked(XHCI_CONTROLLER* xhci,
                                                      uint64_t transferDeviceAddress,
                                                      uint8_t slotId,
                                                      uint8_t endpointId) {
    for (XHCI_WAIT_TARGET* w = xhci->WaiterList; w; w = w->Next) {
        if (w->TransferDeviceAddress &&
            w->TransferDeviceAddress == transferDeviceAddress &&
            w->TransferSlotId == slotId &&
            w->TransferEndpointId == endpointId) {
            return w;
        }
    }
    return (XHCI_WAIT_TARGET*)0;
}

static XHCI_WAIT_TARGET* XhciClaimCommandWaiter(XHCI_CONTROLLER* xhci,
                                                uint64_t commandDeviceAddress,
                                                uint32_t completionCode,
                                                uint8_t slotId) {
    uintptr_t flags = KeAcquireSpinlock(&xhci->WaitLock);
    XHCI_WAIT_TARGET* w = XhciFindCommandWaiterLocked(xhci, commandDeviceAddress);
    if (w) {
        w->Command.CompletionCode = completionCode;
        w->Command.SlotId = slotId;
        XhciBarrier();
        w->CommandMatched = true;
    }
    KeReleaseSpinlock(&xhci->WaitLock, flags);
    return w;
}

static XHCI_WAIT_TARGET* XhciClaimTransferWaiter(XHCI_CONTROLLER* xhci,
                                                 uint64_t transferDeviceAddress,
                                                 uint8_t slotId,
                                                 uint8_t endpointId,
                                                 uint32_t completionCode,
                                                 uint32_t residue) {
    uintptr_t flags = KeAcquireSpinlock(&xhci->WaitLock);
    XHCI_WAIT_TARGET* w = XhciFindTransferWaiterLocked(xhci, transferDeviceAddress, slotId, endpointId);
    if (w) {
        w->Transfer.CompletionCode = completionCode;
        w->Transfer.SlotId = slotId;
        w->Transfer.EndpointId = endpointId;
        w->Transfer.Residue = residue;
        XhciBarrier();
        w->TransferMatched = true;
    }
    KeReleaseSpinlock(&xhci->WaitLock, flags);
    return w;
}

// A completion leaves EventLock with its class-device rundown reference held.
// Client callbacks can rearm transfers or acquire unrelated PnP locks.
#define XHCI_COMPLETION_BATCH 16u

typedef struct XHCI_ASYNC_COMPLETION {
    USB_URB* Urb;
    IRP* Irp;
    KSTATUS Status;
    uint32_t ActualLength;
} XHCI_ASYNC_COMPLETION;

// Claims an asynchronous TD exactly once while the event lock is owned.
static bool XhciClaimAsyncCompletion(XHCI_CONTROLLER* xhci, uint64_t trbAddress,
                                      uint32_t completionCode, uint8_t slotId,
                                      uint8_t endpointId, uint32_t residue,
                                      XHCI_ASYNC_COMPLETION* completion);

// Executes callbacks outside the event lock; class remove locks own their lifetime.
static void XhciDispatchAsyncCompletion(const XHCI_ASYNC_COMPLETION* completed) {
    USB_URB* urb = completed->Urb;
    if (urb->Complete)
        urb->Complete(urb, urb->CompleteContext);
    if (completed->Irp)
        IoCompleteRequest(completed->Irp, completed->Status,
                          KSUCCESS(completed->Status) ? completed->ActualLength : 0);
}

// Drains a bounded batch while the event lock is held.
static bool XhciDrainEventRing(XHCI_CONTROLLER* xhci,
                               XHCI_ASYNC_COMPLETION* completions, uint32_t* completionCount,
                               uint32_t* changedPorts) {
    bool processed = false;
    *completionCount = 0;

    for (uint32_t drained = 0; drained < XHCI_EVENT_TRBS; drained++) {
        if (*completionCount == XHCI_COMPLETION_BATCH)
            break;
        XHCI_TRB* ev = &xhci->EventRing[xhci->EventDequeue];
        uint32_t control = ev->Control;
        uint8_t cycle = (uint8_t)(control & XHCI_TRB_CYCLE);
        if (cycle != xhci->EventCycle)
            break;

        processed = true;
        uint32_t type = XhciTrbType(ev);
        uint32_t cc = XhciCompletionCode(ev);
        uint64_t param = ((uint64_t)ev->ParameterHi << 32) | ev->ParameterLo;

        if (type == XHCI_TRB_TYPE_CMD_DONE) {
            uint8_t slotId = XhciEventSlotId(ev);
            UsbTrace("usb: xHCI command completion ptr=0x%016llX cc=%u slot=%u\n",
                     (unsigned long long)param, cc, slotId);
            (void)XhciClaimCommandWaiter(xhci, param, cc, slotId);
        } else if (type == XHCI_TRB_TYPE_TRANSFER) {
            uint8_t slotId = XhciEventSlotId(ev);
            uint8_t epId = XhciEventEndpointId(ev);
            uint32_t residue = XhciTransferResidue(ev);
            if (cc != XHCI_CC_SUCCESS) {
                UsbTrace("usb: xHCI transfer event trb=0x%016llX cc=%u slot=%u ep=%u residue=%u\n",
                         (unsigned long long)param, cc, slotId, epId, residue);
            }
            XHCI_WAIT_TARGET* waiter = XhciClaimTransferWaiter(xhci, param, slotId, epId, cc, residue);
            if (!waiter && XhciClaimAsyncCompletion(xhci, param, cc, slotId, epId, residue,
                                                     &completions[*completionCount]))
                (*completionCount)++;
        } else if (type == XHCI_TRB_TYPE_PORT_SC) {
            uint32_t portId = ev->ParameterLo >> 24;
            UsbTrace("usb: xHCI port-status-change event port=%u cc=%u\n", portId, cc);
            if (portId && portId <= xhci->MaxPorts && portId <= USB_ROOT_HUB_MAX_PORTS)
                changedPorts[portId / 32u] |= 1u << (portId & 31u);
        } else {
            DbgPrint("usb: xHCI event type=%u cc=%u p=0x%016llX status=%08x ctrl=%08x\n",
                     type, cc, (unsigned long long)param, ev->Status, ev->Control);
        }

        xhci->EventDequeue++;
        if (xhci->EventDequeue == XHCI_EVENT_TRBS) {
            xhci->EventDequeue = 0;
            xhci->EventCycle ^= 1u;
        }
    }

    if (processed)
        XhciUpdateErdp(xhci);
    return processed;
}

// Serializes event-ring ownership without calling into class drivers under the lock.
bool XhciDrainEventRingLocked(XHCI_CONTROLLER* xhci, XHCI_WAIT_TARGET* target) {
    (void)target;
    if (!xhci) return false;
    XHCI_ASYNC_COMPLETION completions[XHCI_COMPLETION_BATCH];
    uint32_t changedPorts[8] = { 0 };
    uint32_t count = 0;
    uintptr_t flags = KeAcquireSpinlock(&xhci->EventLock);
    bool processed = XhciDrainEventRing(xhci, completions, &count, changedPorts);
    KeReleaseSpinlock(&xhci->EventLock, flags);

    // Hub wakeups can take scheduler locks; keep them outside EventLock too.
    for (uint32_t word = 0; word < 8; word++) {
        uint32_t bits = changedPorts[word];
        for (uint32_t bit = 0; bit < 32; bit++) {
            if (bits & (1u << bit))
                UsbHubQueueRootPortChange(&xhci->Hcd, (uint8_t)(word * 32u + bit));
        }
    }
    for (uint32_t i = 0; i < count; i++)
        XhciDispatchAsyncCompletion(&completions[i]);
    return processed;
}

// Waits for target Complete.
bool XhciWaitTargetComplete(const XHCI_WAIT_TARGET* target) {
    if (!target) return false;
    if (target->CommandDeviceAddress && !target->CommandMatched) return false;
    if (target->TransferDeviceAddress && !target->TransferMatched) return false;
    return true;
}

// Registers wait Target.
void XhciRegisterWaitTarget(XHCI_CONTROLLER* xhci, XHCI_WAIT_TARGET* target) {
    if (!xhci || !target || target->Registered)
        return;

    uintptr_t flags = KeAcquireSpinlock(&xhci->WaitLock);
    target->Next = xhci->WaiterList;
    target->Registered = true;
    xhci->WaiterList = target;
    KeReleaseSpinlock(&xhci->WaitLock, flags);
}

// Implements the unregister Wait Target operation.
void XhciUnregisterWaitTarget(XHCI_CONTROLLER* xhci, XHCI_WAIT_TARGET* target) {
    if (!xhci || !target || !target->Registered)
        return;

    uintptr_t flags = KeAcquireSpinlock(&xhci->WaitLock);
    XHCI_WAIT_TARGET** link = &xhci->WaiterList;
    while (*link) {
        if (*link == target) {
            *link = target->Next;
            target->Next = (XHCI_WAIT_TARGET*)0;
            target->Registered = false;
            break;
        }
        link = &(*link)->Next;
    }
    KeReleaseSpinlock(&xhci->WaitLock, flags);
}

KSTATUS XhciWaitForArmedTarget(XHCI_CONTROLLER* xhci, XHCI_WAIT_TARGET* target,
                               uint32_t timeoutUs) {
    if (!xhci || !target)
        return NTLSTATUS_INVALID_PARAMETER;

    for (uint32_t i = 0; i < timeoutUs; i++) {
        if (XhciWaitTargetComplete(target)) {
            XhciUnregisterWaitTarget(xhci, target);
            return NTLSTATUS_SUCCESS;
        }

        // The IRQ/DPC path normally owns the event ring after enumeration, but
        // synchronous storage/control waits must not depend exclusively on an
        // interrupt being delivered. Opportunistically drain the ring too;
        // EventLock serializes this with the DPC while WaitLock protects waiter
        // registration/matching.
        if (!xhci->InterruptsEnabled || (i & 0x7u) == 0) {
            XhciDrainEventRingLocked(xhci, target);
        } else if ((XhciRead32(xhci->Op, XHCI_OP_USBSTS) & XHCI_USBSTS_EINT) && (i & 0x3Fu) == 0) {
            KeInsertQueueDpc(&xhci->EventDpc, (void*)0, (void*)0);
        }

        UsbDelayUs(1);
    }

    XhciUnregisterWaitTarget(xhci, target);
    return NTLSTATUS_TIMEOUT;
}

KSTATUS XhciWaitForTarget(XHCI_CONTROLLER* xhci, XHCI_WAIT_TARGET* target,
                          uint32_t timeoutUs) {
    if (!xhci || !target)
        return NTLSTATUS_INVALID_PARAMETER;
    XhciRegisterWaitTarget(xhci, target);
    return XhciWaitForArmedTarget(xhci, target, timeoutUs);
}

void XhciEventDpcRoutine(KDPC* dpc, void* context,
                                void* systemArgument1, void* systemArgument2) {
    (void)dpc;
    (void)systemArgument1;
    (void)systemArgument2;

    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)context;
    if (!xhci || !xhci->Running || !xhci->Runtime)
        return;

    xhci->EventDpcCount++;
    // Processing is bounded even if connected hubs continuously report changes.
    // Updating ERDP clears EHB; remaining events cause the interrupter to
    // assert again, allowing other DPCs and scheduled threads to make progress.
    for (uint32_t pass = 0; pass < 4u; pass++) {
        if (!XhciDrainEventRingLocked(xhci, NULL))
            return;
    }
    uint32_t hits = ++xhci->EventDpcBudgetHits;
    if ((hits & (hits - 1u)) == 0u)
        DbgPrint("usb: xHCI hcd=%02x:%02x.%u event DPC budget reached count=%u\n",
                 xhci->Bus, xhci->Dev, xhci->Fn, hits);
}

#define XHCI_CONTROLLER_MAGIC 0x58484349u  // XHCI
#define XHCI_MAX_CONTROLLERS 8u

static XHCI_CONTROLLER* gXhciControllers[XHCI_MAX_CONTROLLERS];
static uint32_t gXhciControllerCount;

// Registers controller.
static KSTATUS XhciRegisterController(XHCI_CONTROLLER* xhci) {
    if (!xhci) return NTLSTATUS_INVALID_PARAMETER;
    for (uint32_t i = 0; i < gXhciControllerCount; i++) {
        if (gXhciControllers[i] == xhci)
            return NTLSTATUS_SUCCESS;
    }
    if (gXhciControllerCount >= XHCI_MAX_CONTROLLERS)
        return NTLSTATUS_NO_MEMORY;
    gXhciControllers[gXhciControllerCount++] = xhci;
    return NTLSTATUS_SUCCESS;
}

// Removes a controller from the interrupt-dispatch registry.
static void XhciUnregisterController(XHCI_CONTROLLER* xhci) {
    if (!xhci)
        return;
    for (uint32_t i = 0; i < gXhciControllerCount; i++) {
        if (gXhciControllers[i] != xhci)
            continue;
        for (uint32_t next = i + 1; next < gXhciControllerCount; next++)
            gXhciControllers[next - 1] = gXhciControllers[next];
        gXhciControllerCount--;
        gXhciControllers[gXhciControllerCount] = (XHCI_CONTROLLER*)0;
        return;
    }
}

// Implements the isr operation.
void XhciIsr(void) {
    for (uint32_t i = 0; i < gXhciControllerCount; i++) {
        XHCI_CONTROLLER* xhci = gXhciControllers[i];
        if (!xhci || !xhci->Running || !xhci->Op || !xhci->Runtime)
            continue;

        uint32_t status = XhciRead32(xhci->Op, XHCI_OP_USBSTS);
        if (!(status & (XHCI_USBSTS_EINT | XHCI_USBSTS_PCD | XHCI_USBSTS_HSE | XHCI_USBSTS_HCE)))
            continue;

        xhci->InterruptCount++;
        XhciWrite32(xhci->Op, XHCI_OP_USBSTS, status & XHCI_USBSTS_CLEAR_MASK);
        XhciAckInterrupter0(xhci);
        KeInsertQueueDpc(&xhci->EventDpc, (void*)0, (void*)0);
    }
}

// Implements the bios Handoff operation.
KSTATUS XhciBiosHandoff(XHCI_CONTROLLER* xhci) {
    uint32_t xecp = (xhci->Hcc1 >> 16) & 0xFFFFu;
    if (!xecp) return NTLSTATUS_SUCCESS;

    for (uint32_t guard = 0; xecp && guard < 64; guard++) {
        uint32_t off = xecp << 2;
        if (off + 8u > xhci->MmioBytes)
            break;

        uint32_t cap = XhciRead32(xhci->Mmio, off);
        uint8_t id = (uint8_t)(cap & 0xFFu);
        uint8_t next = (uint8_t)((cap >> 8) & 0xFFu);

        if (id == XHCI_EXT_CAP_LEGACY) {
            UsbTrace("usb: xHCI legacy-support cap at +0x%x legsup=%08x\n", off, cap);
            if (cap & XHCI_LEGSUP_BIOS_OWNED) {
                XhciWrite32(xhci->Mmio, off, cap | XHCI_LEGSUP_OS_OWNED);
                bool ok = false;
                for (uint32_t i = 0; i < 1000000u; i++) {
                    uint32_t cur = XhciRead32(xhci->Mmio, off);
                    if ((cur & XHCI_LEGSUP_BIOS_OWNED) == 0 &&
                        (cur & XHCI_LEGSUP_OS_OWNED)) {
                        ok = true;
                        break;
                    }
                    UsbDelayUs(1);
                }
                if (!ok) {
                    DbgPrint("usb: xHCI BIOS handoff timed out; continuing anyway\n");
                }
            }

            // Clear legacy SMI enables/status. RW1C status bits tolerate this.
            XhciWrite32(xhci->Mmio, off + 4u, 0xFFFFFFFFu);
            return NTLSTATUS_SUCCESS;
        }

        xecp = XhciNextExtendedCapabilityDword(xecp, next, xhci->MmioBytes);
    }

    return NTLSTATUS_SUCCESS;
}

// Implements the stop And Reset operation.
KSTATUS XhciStopAndReset(XHCI_CONTROLLER* xhci) {
    uint32_t cmd = XhciRead32(xhci->Op, XHCI_OP_USBCMD);
    if (cmd & XHCI_USBCMD_RS) {
        XhciWrite32(xhci->Op, XHCI_OP_USBCMD, cmd & ~XHCI_USBCMD_RS);
        if (!UsbWaitSet32(xhci->Op, XHCI_OP_USBSTS,
                          XHCI_USBSTS_HCH, 1000000u)) {
            DbgPrint("usb: xHCI stop timeout usbsts=%08x\n",
                     XhciRead32(xhci->Op, XHCI_OP_USBSTS));
            return NTLSTATUS_TIMEOUT;
        }
    }

    XhciWrite32(xhci->Op, XHCI_OP_USBCMD, XHCI_USBCMD_HCRST);
    if (!UsbWaitClear32(xhci->Op, XHCI_OP_USBCMD,
                        XHCI_USBCMD_HCRST, 1000000u)) {
        DbgPrint("usb: xHCI reset timeout usbcmd=%08x\n",
                 XhciRead32(xhci->Op, XHCI_OP_USBCMD));
        return NTLSTATUS_TIMEOUT;
    }

    if (!UsbWaitClear32(xhci->Op, XHCI_OP_USBSTS,
                        XHCI_USBSTS_CNR, 1000000u)) {
        DbgPrint("usb: xHCI controller-not-ready timeout usbsts=%08x\n",
                 XhciRead32(xhci->Op, XHCI_OP_USBSTS));
        return NTLSTATUS_TIMEOUT;
    }

    XhciWrite32(xhci->Op, XHCI_OP_USBSTS,
                XhciRead32(xhci->Op, XHCI_OP_USBSTS) & XHCI_USBSTS_CLEAR_MASK);
    return NTLSTATUS_SUCCESS;
}

// Halts the controller before any device-visible memory can be released.
static void XhciHaltForDmaRelease(XHCI_CONTROLLER* xhci) {
    if (!xhci || !xhci->DmaAdapter)
        return;

    xhci->InterruptsEnabled = false;
    xhci->Running = false;
    if (xhci->SurpriseRemoved || !xhci->Op)
        return;

    if (xhci->Runtime) {
        volatile uint8_t* interrupter = xhci->Runtime + 0x20u;
        XhciWrite32(interrupter, 0x00u, XHCI_IMAN_IP);
    }

    uint32_t command = XhciRead32(xhci->Op, XHCI_OP_USBCMD);
    command &= ~(XHCI_USBCMD_RS | XHCI_USBCMD_INTE);
    XhciWrite32(xhci->Op, XHCI_OP_USBCMD, command);
    if (!UsbWaitSet32(xhci->Op,
                      XHCI_OP_USBSTS,
                      XHCI_USBSTS_HCH,
                      1000000u))
        KePanic(KE_PANIC_CODE_DRIVER,
                "usb: xHCI failed to halt before DMA release usbsts=%08x\n",
                XhciRead32(xhci->Op, XHCI_OP_USBSTS));
}

// Creates scratchpads.
KSTATUS XhciAllocateScratchpads(XHCI_CONTROLLER* xhci) {
    if (xhci->ScratchpadCount == 0) return NTLSTATUS_SUCCESS;
    if (xhci->ScratchpadCount > XHCI_MAX_SCRATCHPADS) {
        DbgPrint("usb: xHCI wants %u scratchpads, capped driver supports %u\n",
                 xhci->ScratchpadCount, XHCI_MAX_SCRATCHPADS);
        return NTLSTATUS_NOT_SUPPORTED;
    }

    xhci->ScratchpadArray = (uint64_t*)UsbDmaAllocPages(xhci, 1, &xhci->ScratchpadArrayDeviceAddress, &xhci->ScratchpadArrayDma);
    if (!xhci->ScratchpadArray) {
        DbgPrint("usb: xHCI scratchpad pointer-array allocation failed count=%u\n",
                 xhci->ScratchpadCount);
        return NTLSTATUS_NO_MEMORY;
    }

    for (uint16_t i = 0; i < xhci->ScratchpadCount; i++) {
        void* page = UsbDmaAllocPages(xhci, 1, &xhci->ScratchpadDeviceAddress[i], &xhci->ScratchpadDma[i]);
        if (!page) {
            DbgPrint("usb: xHCI scratchpad buffer allocation failed index=%u/%u\n",
                     i, xhci->ScratchpadCount);
            return NTLSTATUS_NO_MEMORY;
        }
        (void)page;
        xhci->ScratchpadArray[i] = xhci->ScratchpadDeviceAddress[i];
    }

    xhci->Dcbaa[0] = xhci->ScratchpadArrayDeviceAddress;
    UsbTrace("usb: xHCI allocated %u scratchpad buffer(s) array=0x%016llX\n",
             xhci->ScratchpadCount,
             (unsigned long long)xhci->ScratchpadArrayDeviceAddress);
    return NTLSTATUS_SUCCESS;
}

// Creates runtime Memory.
KSTATUS XhciAllocateRuntimeMemory(XHCI_CONTROLLER* xhci) {
    xhci->Dcbaa = (uint64_t*)UsbDmaAllocPages(xhci, 1, &xhci->DcbaaDeviceAddress, &xhci->DcbaaDma);
    xhci->CommandRing = (XHCI_TRB*)UsbDmaAllocPages(xhci, 1, &xhci->CommandRingDeviceAddress, &xhci->CommandRingDma);
    xhci->EventRing = (XHCI_TRB*)UsbDmaAllocPages(xhci, 1, &xhci->EventRingDeviceAddress, &xhci->EventRingDma);
    xhci->Erst = (XHCI_ERST_ENTRY*)UsbDmaAllocPages(xhci, 1, &xhci->ErstDeviceAddress, &xhci->ErstDma);

    if (!xhci->Dcbaa || !xhci->CommandRing || !xhci->EventRing || !xhci->Erst) {
        DbgPrint("usb: xHCI runtime DMA allocation failed dcbaa=0x%016llX cmd=0x%016llX event=0x%016llX erst=0x%016llX\n",
                 xhci->Dcbaa, xhci->CommandRing, xhci->EventRing, xhci->Erst);
        return NTLSTATUS_NO_MEMORY;
    }

    xhci->CommandEnqueue = 0;
    xhci->CommandCycle = 1;
    xhci->EventDequeue = 0;
    xhci->EventCycle = 1;

    XHCI_TRB* link = &xhci->CommandRing[XHCI_RING_TRBS - 1u];
    link->ParameterLo = (uint32_t)xhci->CommandRingDeviceAddress;
    link->ParameterHi = (uint32_t)(xhci->CommandRingDeviceAddress >> 32);
    link->Status = 0;
    link->Control = (XHCI_TRB_TYPE_LINK << XHCI_TRB_TYPE_SHIFT) |
                    XHCI_TRB_LINK_TC |
                    xhci->CommandCycle;

    xhci->Erst[0].RingSegmentBase = xhci->EventRingDeviceAddress;
    xhci->Erst[0].RingSegmentSize = XHCI_EVENT_TRBS;
    xhci->Erst[0].Reserved = 0;

    KSTATUS st = XhciAllocateScratchpads(xhci);
    if (!KSUCCESS(st)) return st;

    UsbTrace("usb: xHCI dma dcbaa=0x%016llX cmd=0x%016llX event=0x%016llX erst=0x%016llX\n",
             (unsigned long long)xhci->DcbaaDeviceAddress,
             (unsigned long long)xhci->CommandRingDeviceAddress,
             (unsigned long long)xhci->EventRingDeviceAddress,
             (unsigned long long)xhci->ErstDeviceAddress);
    return NTLSTATUS_SUCCESS;
}

// Implements the program Runtime operation.
KSTATUS XhciProgramRuntime(XHCI_CONTROLLER* xhci) {
    uint8_t slots = xhci->MaxSlots;
    if (slots == 0) return NTLSTATUS_DEVICE_ERROR;

    XhciWrite64(xhci->Op, XHCI_OP_DCBAAP, xhci->DcbaaDeviceAddress);
    XhciWrite64(xhci->Op, XHCI_OP_CRCR, xhci->CommandRingDeviceAddress | xhci->CommandCycle);
    XhciWrite32(xhci->Op, XHCI_OP_DNCTRL, 0);
    XhciWrite32(xhci->Op, XHCI_OP_CONFIG, slots);

    volatile uint8_t* ir0 = xhci->Runtime + 0x20u;
    XhciWrite32(ir0, 0x00u, XHCI_IMAN_IP); // clear pending, leave disabled until smoke test passes
    XhciWrite32(ir0, 0x04u, 0);            // no interrupt moderation yet
    XhciWrite32(ir0, 0x08u, 1);            // one ERST entry
    XhciWrite64(ir0, 0x10u, xhci->ErstDeviceAddress);
    XhciWrite64(ir0, 0x18u, xhci->EventRingDeviceAddress);

    XhciBarrier();
    return NTLSTATUS_SUCCESS;
}

// Returns the device address of one command TRB.
uint64_t XhciCommandTrbDeviceAddress(XHCI_CONTROLLER* xhci, uint32_t index) {
    return xhci->CommandRingDeviceAddress + (uint64_t)index * sizeof(XHCI_TRB);
}

static KSTATUS XhciCommandRingAppend(XHCI_CONTROLLER* xhci,
                                      uint64_t parameter,
                                      uint32_t status,
                                      uint32_t controlWithoutCycle,
                                      uint64_t* outDeviceAddress) {
    if (!xhci || !xhci->CommandRing)
        return NTLSTATUS_INVALID_PARAMETER;

    if (xhci->CommandEnqueue >= XHCI_RING_TRBS - 1u) {
        XHCI_TRB* link = &xhci->CommandRing[XHCI_RING_TRBS - 1u];
        link->ParameterLo = (uint32_t)xhci->CommandRingDeviceAddress;
        link->ParameterHi = (uint32_t)(xhci->CommandRingDeviceAddress >> 32);
        link->Status = 0;
        XhciBarrier();
        link->Control = (XHCI_TRB_TYPE_LINK << XHCI_TRB_TYPE_SHIFT) |
                        XHCI_TRB_LINK_TC |
                        xhci->CommandCycle;
        xhci->CommandEnqueue = 0;
        xhci->CommandCycle ^= 1u;
    }

    uint32_t index = xhci->CommandEnqueue++;
    XHCI_TRB* trb = &xhci->CommandRing[index];
    uint64_t trbDeviceAddress = XhciCommandTrbDeviceAddress(xhci, index);

    trb->ParameterLo = (uint32_t)parameter;
    trb->ParameterHi = (uint32_t)(parameter >> 32);
    trb->Status = status;
    XhciBarrier();
    trb->Control = controlWithoutCycle | xhci->CommandCycle;

    if (outDeviceAddress) *outDeviceAddress = trbDeviceAddress;
    return NTLSTATUS_SUCCESS;
}

KSTATUS XhciSubmitCommand(XHCI_CONTROLLER* xhci,
                                  uint64_t parameter,
                                  uint32_t status,
                                  uint32_t controlWithoutCycle,
                                  XHCI_COMMAND_RESULT* outResult) {
    if (!xhci)
        return NTLSTATUS_INVALID_PARAMETER;

    XHCI_WAIT_TARGET target;
    memset(&target, 0, sizeof(target));

    uintptr_t cmdFlags = KeAcquireSpinlock(&xhci->CommandLock);
    uint64_t trbDeviceAddress = 0;
    KSTATUS st = XhciCommandRingAppend(xhci, parameter, status, controlWithoutCycle, &trbDeviceAddress);
    if (!KSUCCESS(st)) {
        KeReleaseSpinlock(&xhci->CommandLock, cmdFlags);
        return st;
    }

    target.CommandDeviceAddress = trbDeviceAddress;
    XhciRegisterWaitTarget(xhci, &target);

    XhciBarrier();
    XhciWrite32(xhci->Doorbells, 0, 0); // host-controller doorbell, command ring
    KeReleaseSpinlock(&xhci->CommandLock, cmdFlags);

    KSTATUS waitStatus = XhciWaitForArmedTarget(xhci, &target, 500000u);
    if (KSUCCESS(waitStatus) && target.CommandMatched) {
        if (outResult) *outResult = target.Command;
        return (target.Command.CompletionCode == XHCI_CC_SUCCESS)
            ? NTLSTATUS_SUCCESS
            : NTLSTATUS_DEVICE_ERROR;
    }

    DbgPrint("usb: xHCI command timeout type=%u usbsts=%08x crcr=0x%016llX\n",
             (controlWithoutCycle >> XHCI_TRB_TYPE_SHIFT) & XHCI_TRB_TYPE_MASK,
             XhciRead32(xhci->Op, XHCI_OP_USBSTS),
             (unsigned long long)XhciRead64(xhci->Op, XHCI_OP_CRCR));
    return waitStatus;
}

// Submits noop Command.
KSTATUS XhciSubmitNoopCommand(XHCI_CONTROLLER* xhci) {
    XHCI_COMMAND_RESULT result;
    KSTATUS st = XhciSubmitCommand(
        xhci,
        0,
        0,
        (XHCI_TRB_TYPE_NOOP_CMD << XHCI_TRB_TYPE_SHIFT),
        &result);

    if (KSUCCESS(st)) {
        UsbTrace("usb: xHCI No-Op command completed successfully\n");
        return NTLSTATUS_SUCCESS;
    }

    DbgPrint("usb: xHCI No-Op command failed st=0x%x cc=%u\n",
             st, result.CompletionCode);
    return st;
}

// Enables slot.
KSTATUS XhciEnableSlot(XHCI_CONTROLLER* xhci, uint8_t* outSlotId) {
    XHCI_COMMAND_RESULT result;
    memset(&result, 0, sizeof(result));
    KSTATUS st = XhciSubmitCommand(
        xhci,
        0,
        0,
        (XHCI_TRB_TYPE_ENABLE_SLOT << XHCI_TRB_TYPE_SHIFT),
        &result);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI Enable Slot failed st=0x%x cc=%u\n",
                 st, result.CompletionCode);
        return st;
    }
    if (result.SlotId == 0 || result.SlotId > XHCI_MAX_TRACKED_SLOTS) {
        DbgPrint("usb: xHCI Enable Slot returned unsupported slot=%u\n", result.SlotId);
        return NTLSTATUS_DEVICE_ERROR;
    }
    *outSlotId = result.SlotId;
    UsbTrace("usb: xHCI Enable Slot -> slot %u\n", result.SlotId);
    return NTLSTATUS_SUCCESS;
}

// Disables slot.
KSTATUS XhciDisableSlot(XHCI_CONTROLLER* xhci, uint8_t slotId) {
    if (!xhci || slotId == 0)
        return NTLSTATUS_INVALID_PARAMETER;

    XHCI_COMMAND_RESULT res;
    memset(&res, 0, sizeof(res));
    KSTATUS st = XhciSubmitCommand(xhci,
                                   0,
                                   0,
                                   (XHCI_TRB_TYPE_DISABLE_SLOT << XHCI_TRB_TYPE_SHIFT) |
                                       ((uint32_t)slotId << 24),
                                   &res);
    if (!KSUCCESS(st))
        return st;
    if (res.CompletionCode != XHCI_CC_SUCCESS) {
        DbgPrint("usb: Disable Slot slot=%u failed cc=%u\n", slotId, res.CompletionCode);
        return NTLSTATUS_DEVICE_ERROR;
    }
    return NTLSTATUS_SUCCESS;
}

// Implements the address Device operation.
KSTATUS XhciAddressDevice(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    XHCI_COMMAND_RESULT result;
    memset(&result, 0, sizeof(result));
    KSTATUS st = XhciSubmitCommand(
        xhci,
        dev->InputContextDeviceAddress,
        0,
        (XHCI_TRB_TYPE_ADDRESS_DEVICE << XHCI_TRB_TYPE_SHIFT) |
            ((uint32_t)dev->SlotId << 24),
        &result);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI Address Device slot=%u failed st=0x%x cc=%u\n",
                 dev->SlotId, st, result.CompletionCode);
        return st;
    }
    UsbTrace("usb: xHCI Address Device slot=%u complete\n", dev->SlotId);
    return NTLSTATUS_SUCCESS;
}

// Starts the requested operation.
KSTATUS XhciStart(XHCI_CONTROLLER* xhci) {
    uint32_t cmd = XhciRead32(xhci->Op, XHCI_OP_USBCMD);
    cmd |= XHCI_USBCMD_RS;
    XhciWrite32(xhci->Op, XHCI_OP_USBCMD, cmd);

    if (!UsbWaitClear32(xhci->Op, XHCI_OP_USBSTS,
                        XHCI_USBSTS_HCH, 1000000u)) {
        DbgPrint("usb: xHCI start timeout usbsts=%08x\n",
                 XhciRead32(xhci->Op, XHCI_OP_USBSTS));
        return NTLSTATUS_TIMEOUT;
    }

    xhci->Running = true;
    return NTLSTATUS_SUCCESS;
}

// Enables interrupts.
void XhciEnableInterrupts(XHCI_CONTROLLER* xhci) {
    volatile uint8_t* ir0 = xhci->Runtime + 0x20u;
    XhciWrite32(xhci->Op, XHCI_OP_USBSTS,
                XhciRead32(xhci->Op, XHCI_OP_USBSTS) & XHCI_USBSTS_CLEAR_MASK);
    XhciWrite32(ir0, 0x00u, XHCI_IMAN_IE | XHCI_IMAN_IP);
    uint32_t cmd = XhciRead32(xhci->Op, XHCI_OP_USBCMD);
    XhciWrite32(xhci->Op, XHCI_OP_USBCMD, cmd | XHCI_USBCMD_INTE | XHCI_USBCMD_HSEE);
}

// Implements the port Reg operation.
uint32_t XhciPortReg(uint8_t port) {
    return XHCI_OP_PORT_BASE + ((uint32_t)port - 1u) * XHCI_OP_PORT_STRIDE;
}

// Clears port Changes.
void XhciClearPortChanges(XHCI_CONTROLLER* xhci, uint8_t port) {
    uint32_t off = XhciPortReg(port);
    uint32_t portsc = XhciRead32(xhci->Op, off);
    uint32_t changes = portsc & XHCI_PORTSC_CHANGE_MASK;
    if (changes) {

        // PORTSC is a mixed register. The change bits are RW1C, but many
        // other bits are either RO or have side effects when written as one
        // (PED disables the port; PR starts a reset). Preserve only Port
        // Power and write one to the RW1C change bits.

        XhciWrite32(xhci->Op, off, (portsc & XHCI_PORTSC_PP) | changes);
    }
}

// Implements the reset Port operation.
KSTATUS XhciResetPort(XHCI_CONTROLLER* xhci, uint8_t port, uint8_t* outSpeed) {
    uint32_t off = XhciPortReg(port);
    uint32_t portsc = XhciRead32(xhci->Op, off);
    if (!(portsc & XHCI_PORTSC_CCS))
        return NTLSTATUS_NOT_FOUND;

    uint32_t speed = (portsc & XHCI_PORTSC_SPEED_MASK) >> XHCI_PORTSC_SPEED_SHIFT;
    uint8_t protocol = XhciPortProtocolMajor(xhci, port, portsc);
    DbgPrint("usb-stage: hcd=%02x:%02x.%u root-port=%u reset begin proto=%u speed=%u portsc=%08x\n",
             xhci->Bus, xhci->Dev, xhci->Fn, port, protocol, speed, portsc);

    XhciClearPortChanges(xhci, port);
    portsc = XhciRead32(xhci->Op, off);

    if (protocol >= 3u || speed >= 4u) {

        // SuperSpeed ports are enabled by link training. Do not issue a
        // USB2-style PORTSC.PR reset; on real Intel controllers that can make
        // the link drop during early enumeration. If firmware/controller
        // already has PR set, just wait for it to settle.

        bool enabled = false;
        for (uint32_t i = 0; i < 1000000u; i++) {
            portsc = XhciRead32(xhci->Op, off);
            if (!(portsc & XHCI_PORTSC_CCS))
                break;
            if (!(portsc & XHCI_PORTSC_PR) && (portsc & XHCI_PORTSC_PED)) {
                enabled = true;
                break;
            }
            XhciDrainEventRingLocked(xhci, NULL);
            UsbDelayUs(1);
        }

        XhciClearPortChanges(xhci, port);
        portsc = XhciRead32(xhci->Op, off);
        speed = (portsc & XHCI_PORTSC_SPEED_MASK) >> XHCI_PORTSC_SPEED_SHIFT;
        if (outSpeed) *outSpeed = (uint8_t)speed;

        UsbTrace("usb: xHCI port %u link ready portsc=%08x ped=%u proto=usb%u speed=%s(%u)\n",
                 port, portsc,
                 (portsc & XHCI_PORTSC_PED) ? 1u : 0u,
                 protocol, XhciPortSpeedName(speed), speed);

        DbgPrint("usb-stage: hcd=%02x:%02x.%u root-port=%u superspeed ready=%u portsc=%08x\n",
                 xhci->Bus, xhci->Dev, xhci->Fn, port, enabled ? 1u : 0u, portsc);
        if (!enabled || !(portsc & XHCI_PORTSC_CCS) || !(portsc & XHCI_PORTSC_PED))
            return NTLSTATUS_DEVICE_ERROR;
        return NTLSTATUS_SUCCESS;
    }

    uint32_t reset = (portsc & XHCI_PORTSC_PP) | XHCI_PORTSC_PR;
    XhciWrite32(xhci->Op, off, reset);

    bool resetDone = false;
    for (uint32_t i = 0; i < 500000u; i++) {
        portsc = XhciRead32(xhci->Op, off);
        if (!(portsc & XHCI_PORTSC_PR) && (portsc & XHCI_PORTSC_PRC)) {
            resetDone = true;
            break;
        }
        XhciDrainEventRingLocked(xhci, NULL);
        UsbDelayUs(1);
    }

    if (!resetDone) {
        DbgPrint("usb: xHCI hcd=%02x:%02x.%u port %u reset timeout portsc=%08x\n",
                 xhci->Bus, xhci->Dev, xhci->Fn, port, portsc);
        return NTLSTATUS_TIMEOUT;
    }

    XhciClearPortChanges(xhci, port);
    portsc = XhciRead32(xhci->Op, off);
    speed = (portsc & XHCI_PORTSC_SPEED_MASK) >> XHCI_PORTSC_SPEED_SHIFT;
    if (outSpeed) *outSpeed = (uint8_t)speed;

    UsbTrace("usb: xHCI port %u reset complete portsc=%08x ped=%u proto=usb%u speed=%s(%u)\n",
             port, portsc,
             (portsc & XHCI_PORTSC_PED) ? 1u : 0u,
             protocol, XhciPortSpeedName(speed), speed);

    DbgPrint("usb-stage: hcd=%02x:%02x.%u root-port=%u reset done portsc=%08x speed=%u\n",
             xhci->Bus, xhci->Dev, xhci->Fn, port, portsc, speed);
    if (!(portsc & XHCI_PORTSC_CCS) || !(portsc & XHCI_PORTSC_PED))
        return NTLSTATUS_DEVICE_ERROR;

    return NTLSTATUS_SUCCESS;
}

// Creates device Memory.
KSTATUS XhciAllocateDeviceMemory(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    dev->InputContext = (uint8_t*)UsbDmaAllocPages(xhci, 1, &dev->InputContextDeviceAddress, &dev->InputContextDma);
    dev->DeviceContext = (uint8_t*)UsbDmaAllocPages(xhci, 1, &dev->DeviceContextDeviceAddress, &dev->DeviceContextDma);
    dev->Ep0Ring = (XHCI_TRB*)UsbDmaAllocPages(xhci, 1, &dev->Ep0RingDeviceAddress, &dev->Ep0RingDma);
    dev->DescriptorBuffer = (uint8_t*)UsbDmaAllocPages(xhci, 1, &dev->DescriptorDeviceAddress, &dev->DescriptorDma);

    if (!dev->InputContext || !dev->DeviceContext || !dev->Ep0Ring || !dev->DescriptorBuffer) {
        XhciReleaseDeviceDma(dev);
        return NTLSTATUS_NO_MEMORY;
    }

    dev->Ep0Enqueue = 0;
    dev->Ep0Cycle = 1;

    XHCI_TRB* link = &dev->Ep0Ring[XHCI_EP0_RING_TRBS - 1u];
    link->ParameterLo = (uint32_t)dev->Ep0RingDeviceAddress;
    link->ParameterHi = (uint32_t)(dev->Ep0RingDeviceAddress >> 32);
    link->Status = 0;
    link->Control = (XHCI_TRB_TYPE_LINK << XHCI_TRB_TYPE_SHIFT) |
                    XHCI_TRB_LINK_TC |
                    dev->Ep0Cycle;

    return NTLSTATUS_SUCCESS;
}

// Builds default Control Context.
void XhciBuildDefaultControlContext(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    memset(dev->InputContext, 0, PAGE_SIZE);
    memset(dev->DeviceContext, 0, PAGE_SIZE);

    uint32_t* icc = XhciInputControlContext(xhci, dev);
    uint32_t* slot = XhciInputSlotContext(xhci, dev);
    uint32_t* ep0 = XhciInputEndpointContext(xhci, dev, 1);

    icc[0] = 0;      // drop flags
    icc[1] = 0x3u;   // add slot context and endpoint 0 context

    uint32_t maxPacket = XhciEpMaxPacket(dev->Speed);
    dev->MaxPacketSize0 = (uint16_t)maxPacket;

    uint8_t rootPort = dev->Path.RootPort ? dev->Path.RootPort : dev->Port;
    uint32_t routeString = dev->Path.RouteString & 0x000fffffu;

    slot[0] = XHCI_CONTEXT_ENTRIES(1) | XHCI_SLOT_SPEED(dev->Speed) | routeString;
    // Slot context dword 1 owns Root Hub Port Number bits 23:16.
    // Putting the port number in dword 2 makes Address Device fail with a
    // context/TRB error on QEMU xHCI.
    slot[1] = XHCI_SLOT_ROOT_PORT(rootPort);
    slot[2] = 0;
    if (dev->Path.BehindHub && dev->Speed <= 2u && dev->Path.ParentHubSlot) {
        // TT Hub Slot ID bits 7:0 and TT Port Number bits 15:8.
        // This is required for low/full-speed devices behind high-speed hubs;
        // harmless for older full-speed hub paths on the controllers we target.
        slot[2] = (uint32_t)dev->Path.ParentHubSlot |
                  ((uint32_t)dev->Path.ParentHubPort << 8);
    }
    slot[3] = 0;

    ep0[0] = 0;
    ep0[1] = XHCI_EP_CERR_3 | XHCI_EP_TYPE_CONTROL | (maxPacket << 16);
    ep0[2] = (uint32_t)(dev->Ep0RingDeviceAddress | XHCI_EP_DCS);
    ep0[3] = (uint32_t)((dev->Ep0RingDeviceAddress | XHCI_EP_DCS) >> 32);
    ep0[4] = 8;

    xhci->Dcbaa[dev->SlotId] = dev->DeviceContextDeviceAddress;
    XhciBarrier();
}

// Returns the device address of one endpoint-zero TRB.
uint64_t XhciEp0TrbDeviceAddress(XHCI_DEVICE* dev, uint32_t index) {
    return dev->Ep0RingDeviceAddress + (uint64_t)index * sizeof(XHCI_TRB);
}

KSTATUS XhciEp0AppendTrb(XHCI_DEVICE* dev,
                                uint64_t parameter,
                                uint32_t status,
                                uint32_t controlWithoutCycle,
                                uint64_t* outDeviceAddress) {
    if (!dev || !dev->Ep0Ring)
        return NTLSTATUS_INVALID_PARAMETER;

    if (dev->Ep0Enqueue >= XHCI_EP0_RING_TRBS - 1u) {
        XHCI_TRB* link = &dev->Ep0Ring[XHCI_EP0_RING_TRBS - 1u];
        link->ParameterLo = (uint32_t)dev->Ep0RingDeviceAddress;
        link->ParameterHi = (uint32_t)(dev->Ep0RingDeviceAddress >> 32);
        link->Status = 0;
        XhciBarrier();
        link->Control = (XHCI_TRB_TYPE_LINK << XHCI_TRB_TYPE_SHIFT) |
                        XHCI_TRB_LINK_TC |
                        dev->Ep0Cycle;
        dev->Ep0Enqueue = 0;
        dev->Ep0Cycle ^= 1u;
    }

    uint32_t index = dev->Ep0Enqueue++;
    XHCI_TRB* trb = &dev->Ep0Ring[index];
    uint64_t deviceAddress = XhciEp0TrbDeviceAddress(dev, index);

    trb->ParameterLo = (uint32_t)parameter;
    trb->ParameterHi = (uint32_t)(parameter >> 32);
    trb->Status = status;
    XhciBarrier();
    trb->Control = controlWithoutCycle | dev->Ep0Cycle;

    if (outDeviceAddress) *outDeviceAddress = deviceAddress;
    return NTLSTATUS_SUCCESS;
}

KSTATUS XhciControlInTimeout(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                    const USB_SETUP_PACKET* setup,
                                    void* buffer,
                                    uint64_t bufferDeviceAddress,
                                    uint16_t length,
                                    uint32_t timeoutUs) {
    if (!buffer || length == 0 || length > PAGE_SIZE)
        return NTLSTATUS_INVALID_PARAMETER;

    memset(buffer, 0, length);

    uint64_t setupValue = UsbSetupPacketToU64(setup);
    KSTATUS st = XhciEp0AppendTrb(
        dev,
        setupValue,
        8,
        (XHCI_TRB_TYPE_SETUP << XHCI_TRB_TYPE_SHIFT) |
            XHCI_TRB_IDT |
            XHCI_TRB_SETUP_TRT_IN,
        NULL);
    if (!KSUCCESS(st)) return st;

    st = XhciEp0AppendTrb(
        dev,
        bufferDeviceAddress,
        length,
        (XHCI_TRB_TYPE_DATA << XHCI_TRB_TYPE_SHIFT) |
            XHCI_TRB_DATA_DIR_IN,
        NULL);
    if (!KSUCCESS(st)) return st;

    uint64_t statusTrbDeviceAddress = 0;
    st = XhciEp0AppendTrb(
        dev,
        0,
        0,
        (XHCI_TRB_TYPE_STATUS << XHCI_TRB_TYPE_SHIFT) |
            XHCI_TRB_IOC,
        &statusTrbDeviceAddress);
    if (!KSUCCESS(st)) return st;

    XHCI_WAIT_TARGET target;
    memset(&target, 0, sizeof(target));
    target.TransferDeviceAddress = statusTrbDeviceAddress;
    target.TransferSlotId = dev->SlotId;
    target.TransferEndpointId = 1;
    XhciRegisterWaitTarget(xhci, &target);

    XhciBarrier();
    XhciWrite32(xhci->Doorbells, (uint32_t)dev->SlotId * 4u, 1u); // endpoint ID 1 = EP0

    if (timeoutUs == 0) timeoutUs = USB_XHCI_CONTROL_TIMEOUT_US;
    st = XhciWaitForArmedTarget(xhci, &target, timeoutUs);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI control IN timeout slot=%u len=%u timeout=%uus\n",
                 dev->SlotId, length, timeoutUs);
        return st;
    }
    if (target.Transfer.CompletionCode == XHCI_CC_SUCCESS)
        return NTLSTATUS_SUCCESS;

    DbgPrint("usb: xHCI control IN failed slot=%u cc=%u residue=%u\n",
             dev->SlotId,
             target.Transfer.CompletionCode,
             target.Transfer.Residue);
    return NTLSTATUS_DEVICE_ERROR;
}

KSTATUS XhciControlIn(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                             const USB_SETUP_PACKET* setup,
                             void* buffer,
                             uint64_t bufferDeviceAddress,
                             uint16_t length) {
    return XhciControlInTimeout(xhci, dev, setup, buffer, bufferDeviceAddress, length,
                                USB_XHCI_CONTROL_TIMEOUT_US);
}

KSTATUS XhciControlNoData(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                      const USB_SETUP_PACKET* setup) {
    uint64_t setupValue = UsbSetupPacketToU64(setup);
    KSTATUS st = XhciEp0AppendTrb(
        dev,
        setupValue,
        8,
        (XHCI_TRB_TYPE_SETUP << XHCI_TRB_TYPE_SHIFT) |
            XHCI_TRB_IDT |
            XHCI_TRB_SETUP_TRT_NONE,
        NULL);
    if (!KSUCCESS(st)) return st;

    uint64_t statusTrbDeviceAddress = 0;
    st = XhciEp0AppendTrb(
        dev,
        0,
        0,
        (XHCI_TRB_TYPE_STATUS << XHCI_TRB_TYPE_SHIFT) |
            XHCI_TRB_STATUS_DIR_IN |
            XHCI_TRB_IOC,
        &statusTrbDeviceAddress);
    if (!KSUCCESS(st)) return st;

    XHCI_WAIT_TARGET target;
    memset(&target, 0, sizeof(target));
    target.TransferDeviceAddress = statusTrbDeviceAddress;
    target.TransferSlotId = dev->SlotId;
    target.TransferEndpointId = 1;
    XhciRegisterWaitTarget(xhci, &target);

    XhciBarrier();
    XhciWrite32(xhci->Doorbells, (uint32_t)dev->SlotId * 4u, 1u);

    st = XhciWaitForArmedTarget(xhci, &target, 1000000u);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI control no-data timeout slot=%u req=%u value=%04x\n",
                 dev->SlotId, setup->bRequest, setup->wValue);
        return st;
    }
    if (target.Transfer.CompletionCode == XHCI_CC_SUCCESS)
        return NTLSTATUS_SUCCESS;

    DbgPrint("usb: xHCI control no-data failed slot=%u cc=%u residue=%u\n",
             dev->SlotId,
             target.Transfer.CompletionCode,
             target.Transfer.Residue);
    return NTLSTATUS_DEVICE_ERROR;
}

KSTATUS XhciGetDeviceDescriptorTimeout(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                              USB_DEVICE_DESCRIPTOR* outDesc,
                                              uint32_t timeoutUs) {
    USB_SETUP_PACKET setup;
    setup.bmRequestType = USB_DIR_IN;
    setup.bRequest = USB_REQ_GET_DESCRIPTOR;
    setup.wValue = (uint16_t)(USB_DESC_DEVICE << 8);
    setup.wIndex = 0;
    setup.wLength = sizeof(USB_DEVICE_DESCRIPTOR);

    if (timeoutUs == 0) timeoutUs = USB_XHCI_CONTROL_TIMEOUT_US;
    KSTATUS st = XhciControlInTimeout(xhci, dev, &setup,
                                      dev->DescriptorBuffer,
                                      dev->DescriptorDeviceAddress,
                                      sizeof(USB_DEVICE_DESCRIPTOR),
                                      timeoutUs);
    if (!KSUCCESS(st)) return st;

    memcpy(outDesc, dev->DescriptorBuffer, sizeof(*outDesc));
    if (outDesc->bLength < sizeof(USB_DEVICE_DESCRIPTOR) ||
        outDesc->bDescriptorType != USB_DESC_DEVICE) {
        DbgPrint("usb: invalid device descriptor slot=%u len=%u type=%u\n",
                 dev->SlotId, outDesc->bLength, outDesc->bDescriptorType);
        return NTLSTATUS_DEVICE_ERROR;
    }

    return NTLSTATUS_SUCCESS;
}

KSTATUS XhciGetDeviceDescriptor(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                       USB_DEVICE_DESCRIPTOR* outDesc) {
    return XhciGetDeviceDescriptorTimeout(xhci, dev, outDesc,
                                          USB_XHCI_CONTROL_TIMEOUT_US);
}

KSTATUS XhciGetConfigurationDescriptor(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                            uint8_t configurationIndex,
                                            USB_CONFIGURATION_DESCRIPTOR* outConfig,
                                            uint16_t* outTotalLength) {
    USB_SETUP_PACKET setup;
    setup.bmRequestType = USB_DIR_IN;
    setup.bRequest = USB_REQ_GET_DESCRIPTOR;
    setup.wValue = (uint16_t)((USB_DESC_CONFIGURATION << 8) | configurationIndex);
    setup.wIndex = 0;
    setup.wLength = sizeof(USB_CONFIGURATION_DESCRIPTOR);

    KSTATUS st = XhciControlIn(xhci, dev, &setup,
                               dev->DescriptorBuffer,
                               dev->DescriptorDeviceAddress,
                               sizeof(USB_CONFIGURATION_DESCRIPTOR));
    if (!KSUCCESS(st)) return st;

    memcpy(outConfig, dev->DescriptorBuffer, sizeof(*outConfig));
    if (outConfig->bLength < sizeof(USB_CONFIGURATION_DESCRIPTOR) ||
        outConfig->bDescriptorType != USB_DESC_CONFIGURATION ||
        outConfig->wTotalLength < sizeof(USB_CONFIGURATION_DESCRIPTOR)) {
        DbgPrint("usb: invalid config descriptor slot=%u index=%u len=%u type=%u total=%u\n",
                 dev->SlotId,
                 configurationIndex,
                 outConfig->bLength,
                 outConfig->bDescriptorType,
                 outConfig->wTotalLength);
        return NTLSTATUS_DEVICE_ERROR;
    }

    uint16_t total = outConfig->wTotalLength;
    if (total > PAGE_SIZE) {
        DbgPrint("usb: config descriptor slot=%u index=%u total=%u too large, capped at %u\n",
                 dev->SlotId, configurationIndex, total, (uint32_t)PAGE_SIZE);
        total = (uint16_t)PAGE_SIZE;
    }

    setup.wLength = total;
    st = XhciControlIn(xhci, dev, &setup,
                       dev->DescriptorBuffer,
                       dev->DescriptorDeviceAddress,
                       total);
    if (!KSUCCESS(st)) return st;

    memcpy(outConfig, dev->DescriptorBuffer, sizeof(*outConfig));
    if (outConfig->bLength < sizeof(USB_CONFIGURATION_DESCRIPTOR) ||
        outConfig->bDescriptorType != USB_DESC_CONFIGURATION) {
        DbgPrint("usb: invalid full config descriptor slot=%u index=%u len=%u type=%u\n",
                 dev->SlotId,
                 configurationIndex,
                 outConfig->bLength,
                 outConfig->bDescriptorType);
        return NTLSTATUS_DEVICE_ERROR;
    }

    if (outTotalLength) *outTotalLength = total;
    return NTLSTATUS_SUCCESS;
}

// Implements the endpoint Id From Address operation.
uint8_t XhciEndpointIdFromAddress(uint8_t endpointAddress) {
    uint8_t number = endpointAddress & USB_ENDPOINT_NUMBER_MASK;
    if (number == 0) return 1;
    return (uint8_t)(number * 2u + ((endpointAddress & USB_ENDPOINT_DIR_IN) ? 1u : 0u));
}

// Implements the interrupt Interval operation.
uint8_t XhciInterruptInterval(uint8_t speed, uint8_t usbInterval) {
    if (usbInterval == 0) usbInterval = 1;

    if (speed == 3 || speed == 4 || speed == 5) {
        // High-speed and newer interrupt endpoints already encode bInterval as
        // an exponent in USB descriptors. xHCI stores the zero-based exponent.
        if (usbInterval > 16) usbInterval = 16;
        return (uint8_t)(usbInterval - 1u);
    }

    // Full/low-speed interrupt endpoints express bInterval in milliseconds.
    // xHCI wants the nearest exponent in 125us units.
    uint32_t units125us = (uint32_t)usbInterval * 8u;
    uint8_t interval = 0;
    uint32_t v = 1;
    while (v < units125us && interval < 15u) {
        v <<= 1;
        interval++;
    }
    return interval;
}

KSTATUS XhciRingAppend(XHCI_TRB* ring, uint64_t ringDeviceAddress,
                              uint32_t trbCount, uint32_t* enqueue, uint8_t* cycle,
                              uint64_t parameter, uint32_t status,
                              uint32_t controlWithoutCycle, uint64_t* outDeviceAddress) {
    if (!ring || !enqueue || !cycle || trbCount < 2)
        return NTLSTATUS_INVALID_PARAMETER;

    if (*enqueue >= trbCount - 1u) {
        XHCI_TRB* link = &ring[trbCount - 1u];
        link->ParameterLo = (uint32_t)ringDeviceAddress;
        link->ParameterHi = (uint32_t)(ringDeviceAddress >> 32);
        link->Status = 0;
        XhciBarrier();
        link->Control = (XHCI_TRB_TYPE_LINK << XHCI_TRB_TYPE_SHIFT) |
                        XHCI_TRB_LINK_TC |
                        *cycle;
        *enqueue = 0;
        *cycle ^= 1u;
    }

    uint32_t index = (*enqueue)++;
    XHCI_TRB* trb = &ring[index];
    uint64_t deviceAddress = ringDeviceAddress + (uint64_t)index * sizeof(XHCI_TRB);

    trb->ParameterLo = (uint32_t)parameter;
    trb->ParameterHi = (uint32_t)(parameter >> 32);
    trb->Status = status;
    XhciBarrier();
    trb->Control = controlWithoutCycle | *cycle;

    if (outDeviceAddress) *outDeviceAddress = deviceAddress;
    return NTLSTATUS_SUCCESS;
}

void XhciInitialiseTransferRing(XHCI_TRB* ring, uint64_t ringDeviceAddress,
                                       uint32_t trbCount, uint8_t cycle) {
    memset(ring, 0, trbCount * sizeof(XHCI_TRB));
    XHCI_TRB* link = &ring[trbCount - 1u];
    link->ParameterLo = (uint32_t)ringDeviceAddress;
    link->ParameterHi = (uint32_t)(ringDeviceAddress >> 32);
    link->Status = 0;
    link->Control = (XHCI_TRB_TYPE_LINK << XHCI_TRB_TYPE_SHIFT) |
                    XHCI_TRB_LINK_TC |
                    cycle;
}

KSTATUS XhciBulkTransfer(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                uint8_t endpointId,
                                XHCI_TRB* ring,
                                uint64_t ringDeviceAddress,
                                uint32_t* enqueue,
                                uint8_t* cycle,
                                uint64_t bufferDeviceAddress,
                                uint32_t length,
                                uint32_t* outActual) {
    if (!xhci || !dev || !ring || !enqueue || !cycle || !bufferDeviceAddress || length == 0)
        return NTLSTATUS_INVALID_PARAMETER;
    if (length > USB_STORAGE_MAX_BYTES)
        return NTLSTATUS_INVALID_PARAMETER;

    uint64_t trbDeviceAddress = 0;
    KSTATUS st = XhciRingAppend(ring,
                                ringDeviceAddress,
                                XHCI_BULK_RING_TRBS,
                                enqueue,
                                cycle,
                                bufferDeviceAddress,
                                length,
                                (XHCI_TRB_TYPE_NORMAL << XHCI_TRB_TYPE_SHIFT) |
                                    XHCI_TRB_ISP |
                                    XHCI_TRB_IOC,
                                &trbDeviceAddress);
    if (!KSUCCESS(st)) return st;

    XHCI_WAIT_TARGET target;
    memset(&target, 0, sizeof(target));
    target.TransferDeviceAddress = trbDeviceAddress;
    target.TransferSlotId = dev->SlotId;
    target.TransferEndpointId = endpointId;

    XhciRegisterWaitTarget(xhci, &target);
    XhciBarrier();
    XhciWrite32(xhci->Doorbells, (uint32_t)dev->SlotId * 4u, endpointId);

    st = XhciWaitForArmedTarget(xhci, &target, 1000000u);
    if (!KSUCCESS(st)) {
        DbgPrint("usb-storage: bulk transfer timeout slot=%u ep=%u len=%u\n",
                 dev->SlotId, endpointId, length);
        return st;
    }
    if (target.Transfer.CompletionCode != XHCI_CC_SUCCESS) {
        DbgPrint("usb-storage: bulk transfer failed slot=%u ep=%u cc=%u residue=%u\n",
                 dev->SlotId,
                 endpointId,
                 target.Transfer.CompletionCode,
                 target.Transfer.Residue);
        return NTLSTATUS_DEVICE_ERROR;
    }

    if (outActual)
        *outActual = length - target.Transfer.Residue;
    return NTLSTATUS_SUCCESS;
}

// Implements the configure Pipe Endpoint operation.
KSTATUS XhciConfigurePipeEndpoint(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev, USB_PIPE* pipe) {
    if (!xhci || !dev || !pipe || !pipe->Open)
        return NTLSTATUS_INVALID_PARAMETER;

    uint8_t dci = pipe->EndpointId;
    if (dci < 2 || dci > 31)
        return NTLSTATUS_INVALID_PARAMETER;

    if (pipe->HcdPipeContext) {
        XHCI_PIPE_CONTEXT* ctx = (XHCI_PIPE_CONTEXT*)pipe->HcdPipeContext;
        return ctx->Configured ? NTLSTATUS_SUCCESS : NTLSTATUS_DEVICE_ERROR;
    }

    XHCI_PIPE_CONTEXT* ctx = (XHCI_PIPE_CONTEXT*)MmAllocatePool(sizeof(XHCI_PIPE_CONTEXT));
    if (!ctx) return NTLSTATUS_NO_MEMORY;
    memset(ctx, 0, sizeof(*ctx));

    ctx->Ring = (XHCI_TRB*)UsbDmaAllocPages(xhci, 1, &ctx->RingDeviceAddress, &ctx->RingDma);
    if (!ctx->Ring) {
        MmFreePool(ctx);
        return NTLSTATUS_NO_MEMORY;
    }
    ctx->Enqueue = 0;
    ctx->Cycle = 1;
    ctx->EndpointId = dci;
    ctx->Type = pipe->Type;
    XhciInitialiseTransferRing(ctx->Ring, ctx->RingDeviceAddress, XHCI_PIPE_RING_TRBS, ctx->Cycle);

    memset(dev->InputContext, 0, PAGE_SIZE);
    uint32_t* icc = XhciInputControlContext(xhci, dev);
    uint32_t* slot = XhciInputSlotContext(xhci, dev);
    uint32_t* ep = XhciInputEndpointContext(xhci, dev, dci);

    icc[0] = 0;
    icc[1] = (1u << 0) | (1u << dci);

    XhciBuildConfigureSlotContext(xhci, dev, slot, dci);

    uint16_t maxPacket = pipe->MaxPacketSize & 0x07FFu;
    if (maxPacket == 0) maxPacket = 8u;

    uint32_t endpointType = 0;
    uint8_t interval = 0;
    if (pipe->Type == UsbPipeInterrupt && (pipe->EndpointAddress & USB_ENDPOINT_DIR_IN)) {
        endpointType = XHCI_EP_TYPE_INT_IN;
        interval = XhciInterruptInterval(dev->Speed, pipe->Interval);
    } else if (pipe->Type == UsbPipeBulk && (pipe->EndpointAddress & USB_ENDPOINT_DIR_IN)) {
        endpointType = XHCI_EP_TYPE_BULK_IN;
    } else if (pipe->Type == UsbPipeBulk) {
        endpointType = XHCI_EP_TYPE_BULK_OUT;
    } else {
        IoDmaFreeCommonBuffer(ctx->RingDma);
        MmFreePool(ctx);
        return NTLSTATUS_NOT_SUPPORTED;
    }

    ep[0] = (uint32_t)interval << 16;
    ep[1] = XHCI_EP_CERR_3 | endpointType | ((uint32_t)maxPacket << 16);
    ep[2] = (uint32_t)(ctx->RingDeviceAddress | XHCI_EP_DCS);
    ep[3] = (uint32_t)((ctx->RingDeviceAddress | XHCI_EP_DCS) >> 32);
    ep[4] = (uint32_t)maxPacket | ((uint32_t)maxPacket << 16);

    XHCI_COMMAND_RESULT result;
    memset(&result, 0, sizeof(result));
    KSTATUS st = XhciSubmitCommand(
        xhci,
        dev->InputContextDeviceAddress,
        0,
        (XHCI_TRB_TYPE_CONFIGURE_ENDPOINT << XHCI_TRB_TYPE_SHIFT) |
            ((uint32_t)dev->SlotId << 24),
        &result);
    if (!KSUCCESS(st) || result.CompletionCode != XHCI_CC_SUCCESS) {
        DbgPrint("usb: Configure pipe endpoint failed slot=%u ep=%u type=%u st=0x%x cc=%u\n",
                 dev->SlotId, dci, pipe->Type, st, result.CompletionCode);
        IoDmaFreeCommonBuffer(ctx->RingDma);
        MmFreePool(ctx);
        return KSUCCESS(st) ? NTLSTATUS_DEVICE_ERROR : st;
    }

    ctx->Configured = true;
    pipe->HcdPipeContext = ctx;
    dev->EndpointPipes[dci] = pipe;
    UsbReferencePipe(pipe);
    return NTLSTATUS_SUCCESS;
}

// Implements the hcd Submit Interrupt Urb operation.
static KSTATUS XhciHcdSubmitInterruptUrb(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev, USB_URB* urb) {
    if (!xhci || !dev || !urb || !urb->Pipe || !urb->BufferDeviceAddress || urb->Length == 0)
        return NTLSTATUS_INVALID_PARAMETER;
    if (urb->Pipe->Type != UsbPipeInterrupt || !(urb->Pipe->EndpointAddress & USB_ENDPOINT_DIR_IN))
        return NTLSTATUS_NOT_SUPPORTED;

    KSTATUS st = XhciConfigurePipeEndpoint(xhci, dev, urb->Pipe);
    if (!KSUCCESS(st)) return st;

    XHCI_PIPE_CONTEXT* ctx = (XHCI_PIPE_CONTEXT*)urb->Pipe->HcdPipeContext;
    if (!ctx || !ctx->Configured)
        return NTLSTATUS_DEVICE_ERROR;
    uintptr_t flags = KeAcquireSpinlock(&xhci->EventLock);
    if (ctx->ActiveTransfer) {
        KeReleaseSpinlock(&xhci->EventLock, flags);
        return NTLSTATUS_QUEUE_FULL;
    }
    if (urb->Pipe->Halted) {
        KeReleaseSpinlock(&xhci->EventLock, flags);
        return NTLSTATUS_OBJECT_DEAD;
    }

    uint64_t trbDeviceAddress = 0;
    st = XhciRingAppend(ctx->Ring,
                        ctx->RingDeviceAddress,
                        XHCI_PIPE_RING_TRBS,
                        &ctx->Enqueue,
                        &ctx->Cycle,
                        urb->BufferDeviceAddress,
                        urb->Length,
                        (XHCI_TRB_TYPE_NORMAL << XHCI_TRB_TYPE_SHIFT) |
                            XHCI_TRB_ISP |
                            XHCI_TRB_IOC,
                        &trbDeviceAddress);
    if (!KSUCCESS(st)) {
        KeReleaseSpinlock(&xhci->EventLock, flags);
        return st;
    }

    ctx->ActiveTransfer = true;
    ctx->ActiveTrbDeviceAddress = trbDeviceAddress;
    ctx->ActiveUrb = urb;
    urb->Pending = true;
    urb->Cancelled = false;
    urb->Status = NTLSTATUS_PENDING;
    urb->ActualLength = 0;
    urb->HcdCompletionCode = 0;

    XhciBarrier();
    XhciWrite32(xhci->Doorbells, (uint32_t)dev->SlotId * 4u, ctx->EndpointId);
    KeReleaseSpinlock(&xhci->EventLock, flags);
    return NTLSTATUS_PENDING;
}

// Completes ownership of an interrupt URB before its callback is dispatched.
static bool XhciClaimAsyncCompletion(XHCI_CONTROLLER* xhci, uint64_t trbDeviceAddress,
                                      uint32_t completionCode, uint8_t slotId,
                                      uint8_t endpointId, uint32_t residue,
                                      XHCI_ASYNC_COMPLETION* completed) {
    if (!xhci || !completed || slotId == 0 || slotId > XHCI_MAX_TRACKED_SLOTS || endpointId >= 32)
        return false;

    XHCI_DEVICE* dev = &xhci->Devices[slotId];
    USB_PIPE* pipe = dev->EndpointPipes[endpointId];
    if (!pipe || !pipe->HcdPipeContext)
        return false;

    XHCI_PIPE_CONTEXT* ctx = (XHCI_PIPE_CONTEXT*)pipe->HcdPipeContext;
    if (!ctx->ActiveTransfer || ctx->ActiveTrbDeviceAddress != trbDeviceAddress)
        return false;

    USB_URB* urb = ctx->ActiveUrb;
    ctx->ActiveTransfer = false;
    ctx->ActiveTrbDeviceAddress = 0;
    ctx->ActiveUrb = (USB_URB*)0;
    if (!urb)
        return false;

    KSTATUS st = NTLSTATUS_SUCCESS;
    if (urb->Cancelled)
        st = NTLSTATUS_CANCELLED;
    else if (completionCode != XHCI_CC_SUCCESS && completionCode != XHCI_CC_SHORT_PACKET)
        st = NTLSTATUS_DEVICE_ERROR;

    uint32_t actual = KSUCCESS(st) && residue <= urb->Length ? urb->Length - residue : 0;
    urb->Pending = false;
    urb->Status = st;
    urb->ActualLength = actual;
    urb->HcdCompletionCode = completionCode;
    completed->Urb = urb;
    completed->Irp = urb->Irp;
    completed->Status = st;
    completed->ActualLength = actual;
    urb->Irp = (IRP*)0;
    return true;
}

// Implements the hcd Submit Bulk Urb operation.
static KSTATUS XhciHcdSubmitBulkUrb(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev, USB_URB* urb) {
    if (!xhci || !dev || !urb || !urb->Pipe || !urb->BufferDeviceAddress || urb->Length == 0)
        return NTLSTATUS_INVALID_PARAMETER;

    uint8_t endpointId = urb->Pipe->EndpointId;
    XHCI_TRB* ring = (XHCI_TRB*)0;
    uint64_t ringDeviceAddress = 0;
    uint32_t* enqueue = (uint32_t*)0;
    uint8_t* cycle = (uint8_t*)0;

    if (endpointId == dev->MassBulkInEndpointId) {
        ring = dev->MassInRing;
        ringDeviceAddress = dev->MassInRingDeviceAddress;
        enqueue = &dev->MassInEnqueue;
        cycle = &dev->MassInCycle;
    } else if (endpointId == dev->MassBulkOutEndpointId) {
        ring = dev->MassOutRing;
        ringDeviceAddress = dev->MassOutRingDeviceAddress;
        enqueue = &dev->MassOutEnqueue;
        cycle = &dev->MassOutCycle;
    } else {
        return NTLSTATUS_NOT_SUPPORTED;
    }

    uint32_t actual = 0;
    KSTATUS st = XhciBulkTransfer(xhci,
                                  dev,
                                  endpointId,
                                  ring,
                                  ringDeviceAddress,
                                  enqueue,
                                  cycle,
                                  urb->BufferDeviceAddress,
                                  urb->Length,
                                  &actual);
    urb->ActualLength = KSUCCESS(st) ? actual : 0;
    return st;
}

// Implements the hcd Submit Urb operation.
static KSTATUS XhciHcdSubmitUrb(USB_HCD* hcd, USB_URB* urb) {
    if (!hcd || !urb || !urb->Device)
        return NTLSTATUS_INVALID_PARAMETER;

    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)hcd->ControllerContext;
    XHCI_DEVICE* dev = (XHCI_DEVICE*)urb->Device->HcdDeviceContext;
    if (!xhci || !dev || !dev->Allocated)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS st;
    switch (urb->Function) {
        case UrbFunctionControlTransfer:
            if (urb->Length == 0) {
                st = XhciControlNoData(xhci, dev, &urb->Setup);
                urb->ActualLength = 0;
            } else if ((urb->Setup.bmRequestType & USB_DIR_IN) != 0) {
                st = XhciControlIn(xhci, dev, &urb->Setup, urb->Buffer, urb->BufferDeviceAddress, (uint16_t)urb->Length);
                urb->ActualLength = KSUCCESS(st) ? urb->Length : 0;
            } else {
                st = NTLSTATUS_NOT_SUPPORTED;
                urb->ActualLength = 0;
            }
            break;

        case UrbFunctionBulkTransfer:
            st = XhciHcdSubmitBulkUrb(xhci, dev, urb);
            break;

        case UrbFunctionInterruptTransfer:
            st = XhciHcdSubmitInterruptUrb(xhci, dev, urb);
            break;

        default:
            st = NTLSTATUS_NOT_SUPPORTED;
            urb->ActualLength = 0;
            break;
    }

    urb->Status = st;
    return st;
}

static KSTATUS XhciHcdCancelUrb(USB_HCD* hcd, USB_URB* urb) {
    if (!hcd || !urb || !urb->Device || !urb->Pipe)
        return NTLSTATUS_INVALID_PARAMETER;

    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)hcd->ControllerContext;
    XHCI_DEVICE* dev = (XHCI_DEVICE*)urb->Device->HcdDeviceContext;
    if (!xhci || !dev || dev->SlotId == 0 || !urb->Pipe->HcdPipeContext)
        return NTLSTATUS_INVALID_PARAMETER;

    XHCI_PIPE_CONTEXT* ctx = (XHCI_PIPE_CONTEXT*)urb->Pipe->HcdPipeContext;
    uintptr_t flags = KeAcquireSpinlock(&xhci->EventLock);
    urb->Cancelled = true;
    bool active = ctx->ActiveTransfer && ctx->ActiveUrb == urb;
    KeReleaseSpinlock(&xhci->EventLock, flags);
    if (!active)
        return NTLSTATUS_SUCCESS;

    XHCI_COMMAND_RESULT result;
    memset(&result, 0, sizeof(result));
    KSTATUS st = XhciSubmitCommand(xhci, 0, 0,
                                   (XHCI_TRB_TYPE_STOP_ENDPOINT << XHCI_TRB_TYPE_SHIFT) |
                                       ((uint32_t)ctx->EndpointId << 16) |
                                       ((uint32_t)dev->SlotId << 24),
                                   &result);

    flags = KeAcquireSpinlock(&xhci->EventLock);
    if (!ctx->ActiveTransfer || ctx->ActiveUrb != urb) {
        KeReleaseSpinlock(&xhci->EventLock, flags);
        return NTLSTATUS_SUCCESS;
    }
    if (!KSUCCESS(st) || result.CompletionCode != XHCI_CC_SUCCESS) {
        KeReleaseSpinlock(&xhci->EventLock, flags);
        DbgPrint("usb: xHCI STOP_ENDPOINT failed slot=%u ep=%u st=0x%x cc=%u\n",
                 dev->SlotId, ctx->EndpointId, st, result.CompletionCode);
        return KSUCCESS(st) ? NTLSTATUS_DEVICE_ERROR : st;
    }

    ctx->ActiveTransfer = false;
    ctx->ActiveTrbDeviceAddress = 0;
    ctx->ActiveUrb = (USB_URB*)0;
    urb->Pipe->Halted = true;
    urb->Pending = false;
    urb->Status = NTLSTATUS_CANCELLED;
    urb->ActualLength = 0;
    urb->HcdCompletionCode = result.CompletionCode;

    IRP* completeIrp = urb->Irp;
    urb->Irp = (IRP*)0;
    KeReleaseSpinlock(&xhci->EventLock, flags);
    if (urb->Complete)
        urb->Complete(urb, urb->CompleteContext);
    if (completeIrp)
        IoCompleteRequest(completeIrp, NTLSTATUS_CANCELLED, 0);
    return NTLSTATUS_SUCCESS;
}

// Implements the hcd Allocate Dma operation.
static void* XhciHcdAllocateDma(USB_HCD* hcd,
                                size_t pages,
                                uint64_t* outDeviceAddress,
                                IO_DMA_BUFFER** outBuffer) {
    if (!hcd) return (void*)0;
    return UsbDmaAllocPages((XHCI_CONTROLLER*)hcd->ControllerContext,
                            pages,
                            outDeviceAddress,
                            outBuffer);
}

// Releases the host-controller context attached to a USB pipe.
static void XhciHcdDestroyPipe(USB_HCD* hcd, USB_PIPE* pipe) {
    (void)hcd;
    if (!pipe || !pipe->HcdPipeContext)
        return;
    XHCI_PIPE_CONTEXT* ctx = (XHCI_PIPE_CONTEXT*)pipe->HcdPipeContext;
    IoDmaFreeCommonBuffer(ctx->RingDma);
    ctx->RingDma = (IO_DMA_BUFFER*)0;
    ctx->Ring = (XHCI_TRB*)0;
    pipe->HcdPipeContext = (void*)0;
    MmFreePool(ctx);
}

// Implements the hcd Reset Root Port operation.
static KSTATUS XhciHcdResetRootPort(USB_HCD* hcd, uint8_t port, uint8_t* outSpeed) {
    if (!hcd) return NTLSTATUS_INVALID_PARAMETER;
    return XhciResetPort((XHCI_CONTROLLER*)hcd->ControllerContext, port, outSpeed);
}

// Implements the hcd Clear Root Port Changes operation.
static KSTATUS XhciHcdClearRootPortChanges(USB_HCD* hcd, uint8_t port) {
    if (!hcd) return NTLSTATUS_INVALID_PARAMETER;
    XhciClearPortChanges((XHCI_CONTROLLER*)hcd->ControllerContext, port);
    return NTLSTATUS_SUCCESS;
}

const USB_HCD_OPERATIONS gXhciHcdOperations = {
    XhciHcdSubmitUrb,
    XhciHcdCancelUrb,
    XhciHcdResetRootPort,
    XhciHcdClearRootPortChanges,
    XhciHcdDestroyPipe,
    XhciHcdAllocateDma,
};

// Starts controller.
KSTATUS XhciStartController(XHCI_CONTROLLER* xhci) {
    if (!xhci || xhci->Magic != XHCI_CONTROLLER_MAGIC || !xhci->Pdo || !xhci->Fdo)
        return NTLSTATUS_INVALID_PARAMETER;
    if (xhci->Running)
        return NTLSTATUS_SUCCESS;

    KSTATUS st = PciGetDeviceInfo(xhci->Pdo, &xhci->PciInfo);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI PciGetDeviceInfo failed st=0x%x\n", st);
        return st;
    }

    xhci->Bus = xhci->PciInfo.Address.Bus;
    xhci->Dev = xhci->PciInfo.Address.Device;
    xhci->Fn = xhci->PciInfo.Address.Function;

    st = PciGetBar(xhci->Pdo, 0, &xhci->MmioBar);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI PCI %02x:%02x.%x BAR0 missing st=0x%x\n",
                 xhci->Bus, xhci->Dev, xhci->Fn, st);
        return st;
    }
    if (xhci->MmioBar.Type != PciBarMemory32 && xhci->MmioBar.Type != PciBarMemory64) {
        DbgPrint("usb: xHCI PCI %02x:%02x.%x BAR0 is not MMIO type=%u\n",
                 xhci->Bus, xhci->Dev, xhci->Fn, xhci->MmioBar.Type);
        return NTLSTATUS_DEVICE_ERROR;
    }

    xhci->Irq = 0xFFu;
    xhci->UsingMsi = false;

    st = PciEnableMsiAuto(xhci->Pdo, &xhci->Irq);
    if (KSUCCESS(st)) {
        xhci->UsingMsi = true;
        DbgPrint("usb: xHCI using MSI IRQ slot %u\n", xhci->Irq);
    } else {
        KSTATUS intxStatus = PciGetInterrupt(xhci->Pdo, &xhci->Irq);
        if (KSUCCESS(intxStatus) && xhci->Irq != 0xFFu) {
            DbgPrint("usb: xHCI using legacy INTx IRQ %u\n", xhci->Irq);
        } else {
            xhci->Irq = 0xFFu;
            DbgPrint("usb: xHCI no MSI/INTx interrupt available; using polled event-ring mode (msi=0x%x intx=0x%x)\n",
                     st, intxStatus);
        }
    }

    uint32_t enableFlags = PCI_ENABLE_MEMORY | PCI_ENABLE_BUSMASTER;
    if (xhci->UsingMsi || xhci->Irq == 0xFFu)
        enableFlags |= PCI_DISABLE_INTX;
    else
        enableFlags |= PCI_ENABLE_INTX;

    st = PciEnableDevice(xhci->Pdo, enableFlags);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI PciEnableDevice failed st=0x%x\n", st);
        return st;
    }

    xhci->MmioPhys = xhci->MmioBar.Base;

    uint64_t mappedLength = 0;
    st = PciMapBar(xhci->Pdo,
                   0,
                   PAGE_PRESENT | PAGE_WRITE | PAGE_CACHE_UC,
                   (void**)&xhci->Mmio,
                   &mappedLength);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI PciMapBar(BAR0) failed st=0x%x\n", st);
        return st;
    }
    xhci->MmioBytes = (mappedLength > 0xFFFFFFFFull) ? 0xFFFFFFFFu : (uint32_t)mappedLength;

    xhci->Cap = xhci->Mmio;
    uint32_t cap0 = XhciRead32(xhci->Cap, XHCI_CAP_CAPLENGTH);
    xhci->CapLength = (uint8_t)(cap0 & 0xFFu);
    xhci->Version = (uint16_t)(cap0 >> 16);
    xhci->Hcs1 = XhciRead32(xhci->Cap, XHCI_CAP_HCSPARAMS1);
    xhci->Hcs2 = XhciRead32(xhci->Cap, XHCI_CAP_HCSPARAMS2);
    xhci->Hcc1 = XhciRead32(xhci->Cap, XHCI_CAP_HCCPARAMS1);
    xhci->DbOff = XhciRead32(xhci->Cap, XHCI_CAP_DBOFF) & ~0x3u;
    xhci->RtsOff = XhciRead32(xhci->Cap, XHCI_CAP_RTSOFF) & ~0x1Fu;
    xhci->MaxSlots = (uint8_t)(xhci->Hcs1 & 0xFFu);
    xhci->MaxIntrs = (uint16_t)((xhci->Hcs1 >> 8) & 0x7FFu);
    xhci->MaxPorts = (uint8_t)((xhci->Hcs1 >> 24) & 0xFFu);
    xhci->ScratchpadCount = XhciScratchpadCount(xhci->Hcs2);
    xhci->ContextBytes = (xhci->Hcc1 & XHCI_HCC1_CSZ) ? 64u : 32u;

    if (xhci->CapLength < 0x20 || xhci->DbOff == 0 || xhci->RtsOff == 0 ||
        xhci->MaxSlots == 0 || xhci->MaxPorts == 0) {
        DbgPrint("usb: xHCI invalid caps caplen=%u dboff=0x%x rtsoff=0x%x hcs1=%08x\n",
                 xhci->CapLength, xhci->DbOff, xhci->RtsOff, xhci->Hcs1);
        return NTLSTATUS_DEVICE_ERROR;
    }

    uint32_t portEnd = xhci->CapLength + XHCI_OP_PORT_BASE +
                       (uint32_t)xhci->MaxPorts * XHCI_OP_PORT_STRIDE;
    uint32_t needBytes = portEnd;
    if (xhci->DbOff + 0x100u > needBytes) needBytes = xhci->DbOff + 0x100u;
    if (xhci->RtsOff + 0x1000u > needBytes) needBytes = xhci->RtsOff + 0x1000u;
    needBytes = UsbAlignUp32(needBytes, (uint32_t)PAGE_SIZE);
    if (xhci->MmioBytes < needBytes) {
        DbgPrint("usb: xHCI BAR0 too small length=0x%x need=0x%x dboff=0x%x rtsoff=0x%x ports=%u\n",
                 xhci->MmioBytes, needBytes, xhci->DbOff, xhci->RtsOff, xhci->MaxPorts);
        return NTLSTATUS_DEVICE_ERROR;
    }

    xhci->Op = xhci->Mmio + xhci->CapLength;
    xhci->Runtime = xhci->Mmio + xhci->RtsOff;
    xhci->Doorbells = xhci->Mmio + xhci->DbOff;

    UsbInitializeHcd(&xhci->Hcd, &gXhciHcdOperations, xhci, 0, xhci->MaxPorts);
    xhci->Hcd.ControllerDevice = xhci->Fdo;

    DbgPrint("usb: xHCI init %02x:%02x.%x vendor=%04x device=%04x irq=%u mmio=0x%016llX caplen=%u version=%04x\n",
             xhci->Bus, xhci->Dev, xhci->Fn,
             xhci->PciInfo.VendorId, xhci->PciInfo.DeviceId,
             xhci->Irq,
             (unsigned long long)xhci->MmioPhys, xhci->CapLength, xhci->Version);
    DbgPrint("usb: xHCI caps slots=%u intrs=%u ports=%u ctx=%u scratchpads=%u hcs2=%08x hcc1=%08x dboff=0x%x rtsoff=0x%x pagesize=%08x\n",
             xhci->MaxSlots, xhci->MaxIntrs, xhci->MaxPorts, xhci->ContextBytes,
             xhci->ScratchpadCount, xhci->Hcs2, xhci->Hcc1, xhci->DbOff,
             xhci->RtsOff, XhciRead32(xhci->Op, XHCI_OP_PAGESIZE));

    XhciParseSupportedProtocols(xhci);

    st = XhciBiosHandoff(xhci);
    if (!KSUCCESS(st)) return st;

    st = XhciStopAndReset(xhci);
    if (!KSUCCESS(st)) return st;

    uint32_t pageSize = XhciRead32(xhci->Op, XHCI_OP_PAGESIZE);
    if ((pageSize & 1u) == 0) {
        DbgPrint("usb: xHCI does not advertise 4K pages pagesize=%08x\n", pageSize);
        return NTLSTATUS_NOT_SUPPORTED;
    }

    IO_DMA_REQUIREMENTS dmaRequirements;
    memset(&dmaRequirements, 0, sizeof(dmaRequirements));
    dmaRequirements.MaximumDeviceAddress =
        (xhci->Hcc1 & XHCI_HCC1_AC64) ? UINT64_MAX : UINT32_MAX;
    dmaRequirements.Alignment = 1;
    dmaRequirements.Boundary = 0x10000u;
    dmaRequirements.MaximumSegmentLength = 0x10000u;
    dmaRequirements.MaximumSegments = XHCI_PIPE_RING_TRBS - 1u;
    st = IoDmaGetAdapter(xhci->Fdo,
                                     &dmaRequirements,
                                     &xhci->DmaAdapter);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI DMA adapter creation failed st=0x%x\n", st);
        return st;
    }

    st = XhciAllocateRuntimeMemory(xhci);
    if (!KSUCCESS(st))
        goto failDma;

    st = XhciProgramRuntime(xhci);
    if (!KSUCCESS(st))
        goto failDma;

    st = XhciStart(xhci);
    if (!KSUCCESS(st))
        goto failDma;

    st = XhciSubmitNoopCommand(xhci);
    if (!KSUCCESS(st))
        goto failDma;

    XhciPowerRootPorts(xhci);
    XhciLogPortStatus(xhci);

    st = UsbCreateRootHubPdo(&xhci->Hcd);
    if (!KSUCCESS(st) && st != NTLSTATUS_ALREADY_EXISTS) {
        DbgPrint("usb: root hub PDO creation failed st=0x%x\n", st);
        goto failDma;
    }

    st = XhciRegisterController(xhci);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: xHCI controller registry full st=0x%x\n", st);
        goto failDma;
    }

    if (xhci->Irq != 0xFFu) {
        KeIrqRegister(xhci->Irq, XhciIsr);
        XhciEnableInterrupts(xhci);
        xhci->InterruptsEnabled = true;
        KeInsertQueueDpc(&xhci->EventDpc, (void*)0, (void*)0);
        DbgPrint("usb: xHCI ISR registered on IRQ %u (%s); event ring handled by DPC\n",
                 xhci->Irq, xhci->UsingMsi ? "MSI" : "INTx");
    } else {
        xhci->InterruptsEnabled = false;
        DbgPrint("usb: xHCI started in polling mode; synchronous waits drain event ring\n");
    }
    xhci->Fdo->PnpState = DevicePnpStarted;

    st = UsbHubScheduleRootScan(&xhci->Hcd);
    if (!KSUCCESS(st))
        DbgPrint("usb: root hub async scan scheduling failed st=0x%x\n", st);
    else
        DbgPrint("usb: root hub initial scan requested\n");

    DbgPrint("usb: xHCI host controller initialised; HID boot input active=%u mass-storage will arrive through hotplug\n",
             xhci->HidDeviceCount);
    return NTLSTATUS_SUCCESS;

failDma:
    XhciHaltForDmaRelease(xhci);
    XhciUnregisterController(xhci);
    if (xhci->Hcd.RootHubPdo) {
        (void)UsbRemovePdoSubtree(xhci->Hcd.RootHubPdo, true);
        xhci->Hcd.RootHubPdo = (DEVICE_OBJECT*)0;
    }
    XhciReleaseControllerDma(xhci);
    return st;
}

// Dispatches pnp.
KSTATUS XhciDispatchPnp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp || !device->DeviceExtension)
        return NTLSTATUS_INVALID_PARAMETER;

    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)device->DeviceExtension;
    if (xhci->Magic != XHCI_CONTROLLER_MAGIC)
        return NTLSTATUS_INVALID_PARAMETER;

    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios) return NTLSTATUS_INVALID_PARAMETER;

    switch ((IO_PNP_MINOR_FUNCTION)ios->MinorFunction) {
        case IRP_MN_START_DEVICE:
            return XhciStartController(xhci);

        case IRP_MN_QUERY_REMOVE_DEVICE:
            device->PnpState = DevicePnpRemovePending;
            return NTLSTATUS_SUCCESS;

        case IRP_MN_CANCEL_REMOVE_DEVICE:
            device->PnpState = xhci->Running ? DevicePnpStarted : DevicePnpNotStarted;
            return NTLSTATUS_SUCCESS;

        case IRP_MN_SURPRISE_REMOVAL:
            device->PnpState = DevicePnpSurpriseRemoved;
            xhci->SurpriseRemoved = true;
            xhci->Running = false;
            xhci->InterruptsEnabled = false;
            if (xhci->Hcd.RootHubPdo)
                (void)UsbRemovePdoSubtree(xhci->Hcd.RootHubPdo, true);
            return NTLSTATUS_SUCCESS;

        case IRP_MN_REMOVE_DEVICE:
            device->PnpState = DevicePnpDeleted;
            XhciHaltForDmaRelease(xhci);
            if (xhci->Hcd.RootHubPdo) {
                (void)UsbRemovePdoSubtree(xhci->Hcd.RootHubPdo, true);
                xhci->Hcd.RootHubPdo = (DEVICE_OBJECT*)0;
            }
            XhciUnregisterController(xhci);
            XhciReleaseControllerDma(xhci);
            IoDetachDevice(device);
            IoDeleteDevice(device);
            return NTLSTATUS_SUCCESS;

        default:
            return IoForwardIrpToLower(device, irp);
    }
}

// Implements the add Device operation.
KSTATUS XhciAddDevice(DRIVER_OBJECT* driver, DEVICE_OBJECT* physicalDeviceObject) {
    if (!driver || !physicalDeviceObject || physicalDeviceObject->Type != DeviceObjectPhysical)
        return NTLSTATUS_INVALID_PARAMETER;

    DEVICE_OBJECT* fdo = (DEVICE_OBJECT*)0;
    KSTATUS st = IoCreateDevice(driver, sizeof(XHCI_CONTROLLER), &fdo);
    if (!KSUCCESS(st))
        return st;
    (void)IoSetDeviceName(fdo, SYS_SV_LITERAL("xhci-controller"));

    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)fdo->DeviceExtension;
    memset(xhci, 0, sizeof(*xhci));
    xhci->Magic = XHCI_CONTROLLER_MAGIC;

    xhci->Fdo = fdo;
    xhci->Pdo = physicalDeviceObject;
    xhci->Irq = 0xFFu;
    KeInitializeDpc(&xhci->EventDpc, XhciEventDpcRoutine, xhci);

    xhci->LowerDevice = IoAttachDeviceToDeviceStack(fdo, physicalDeviceObject);
    if (!xhci->LowerDevice) {
        IoDeleteDevice(fdo);
        return NTLSTATUS_FAILURE;
    }

    DbgPrint("usb: xHCI FDO attached to %v\n", physicalDeviceObject->Name);
    return NTLSTATUS_SUCCESS;
}