// USB HID class driver. This file owns HID interface probing and input report translation.

#include "usb.h"

// Releases the HID class DMA buffer after all pending transfers are cancelled.
static void UsbHidReleaseDma(USB_HID_EXTENSION* ext) {
    if (!ext)
        return;
    UsbFreeDmaBuffer(ext->ReportBufferDma);
    ext->ReportBufferDma = (IO_DMA_BUFFER*)0;
    ext->ReportBuffer = (uint8_t*)0;
    ext->ReportBufferDeviceAddress = 0;
}

// Implements the hid Key In Report operation.
bool XhciHidKeyInReport(const uint8_t* keys, uint8_t key) {
    if (key == 0 || key == 1) return false;
    for (uint8_t i = 0; i < 6; i++) {
        if (keys[i] == key) return true;
    }
    return false;
}

INPUT_KEY XhciHidModifierKey(uint8_t bit) {
    switch (bit) {
        case 0: return INPUT_KEY_LEFT_CTRL;
        case 1: return INPUT_KEY_LEFT_SHIFT;
        case 2: return INPUT_KEY_LEFT_ALT;
        case 3: return INPUT_KEY_LEFT_GUI;
        case 4: return INPUT_KEY_RIGHT_CTRL;
        case 5: return INPUT_KEY_RIGHT_SHIFT;
        case 6: return INPUT_KEY_RIGHT_ALT;
        case 7: return INPUT_KEY_RIGHT_GUI;
        default: return INPUT_KEY_NONE;
    }
}

// Submits hid Key.
void XhciSubmitHidKey(INPUT_KEY key, bool pressed) {
    if (key == INPUT_KEY_NONE) return;
    INPUT_KEY_EVENT ev;
    ev.Key = key;
    ev.Pressed = pressed;
    KeInputSubmitKey(&ev);
}

// Processes boot Keyboard Report.
void XhciProcessBootKeyboardReport(XHCI_DEVICE* dev, const uint8_t* report, uint16_t len) {
    if (len < 8) return;

    uint8_t modifiers = report[0];
    for (uint8_t bit = 0; bit < 8; bit++) {
        uint8_t mask = (uint8_t)(1u << bit);
        bool wasDown = (dev->HidLastModifiers & mask) != 0;
        bool isDown = (modifiers & mask) != 0;
        if (wasDown != isDown)
            XhciSubmitHidKey(XhciHidModifierKey(bit), isDown);
    }

    const uint8_t* keys = &report[2];
    for (uint8_t i = 0; i < 6; i++) {
        uint8_t oldKey = dev->HidLastKeys[i];
        if (oldKey > 1 && !XhciHidKeyInReport(keys, oldKey))
            XhciSubmitHidKey((INPUT_KEY)oldKey, false);
    }
    for (uint8_t i = 0; i < 6; i++) {
        uint8_t newKey = keys[i];
        if (newKey > 1 && !XhciHidKeyInReport(dev->HidLastKeys, newKey))
            XhciSubmitHidKey((INPUT_KEY)newKey, true);
    }

    dev->HidLastModifiers = modifiers;
    memcpy(dev->HidLastKeys, keys, sizeof(dev->HidLastKeys));
}

// Processes boot Mouse Report.
void XhciProcessBootMouseReport(XHCI_DEVICE* dev, const uint8_t* report, uint16_t len) {
    if (len < 3) return;

    int32_t dx = (int32_t)(int8_t)report[1];
    int32_t dy = (int32_t)(int8_t)report[2];
    int32_t wheel = (len >= 4) ? (int32_t)(int8_t)report[3] : 0;
    uint8_t buttons = report[0] & 0x07u;

    if (dx || dy || wheel || buttons != dev->HidLastButtons) {
        INPUT_MOUSE_EVENT ev;
        ev.dx = dx;
        ev.dy = dy;
        ev.wheel = wheel;
        ev.buttons = buttons;
        KeInputSubmitMouse(&ev);
    }

    dev->HidLastButtons = buttons;
}

// Implements the arm Hid Interrupt In operation.
KSTATUS XhciArmHidInterruptIn(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    if (!dev->HidActive || dev->HidArmed || !dev->HidRing || !dev->HidBuffer)
        return NTLSTATUS_SUCCESS;

    memset(dev->HidBuffer, 0, dev->HidReportLength);

    uint64_t trbDeviceAddress = 0;
    KSTATUS st = XhciRingAppend(dev->HidRing,
                                dev->HidRingDeviceAddress,
                                XHCI_HID_RING_TRBS,
                                &dev->HidEnqueue,
                                &dev->HidCycle,
                                dev->HidBufferDeviceAddress,
                                dev->HidReportLength,
                                (XHCI_TRB_TYPE_NORMAL << XHCI_TRB_TYPE_SHIFT) |
                                    XHCI_TRB_ISP |
                                    XHCI_TRB_IOC,
                                &trbDeviceAddress);
    if (!KSUCCESS(st)) return st;

    dev->HidLastTrbDeviceAddress = trbDeviceAddress;
    dev->HidArmed = true;
    XhciBarrier();
    XhciWrite32(xhci->Doorbells, (uint32_t)dev->SlotId * 4u, dev->HidEndpointId);
    return NTLSTATUS_SUCCESS;
}

// Implements the arm All Hid operation.
void XhciArmAllHid(XHCI_CONTROLLER* xhci) {
    for (uint8_t slot = 1; slot <= XHCI_MAX_TRACKED_SLOTS; slot++) {
        XHCI_DEVICE* dev = &xhci->Devices[slot];
        if (!dev->HidActive) continue;
        KSTATUS st = XhciArmHidInterruptIn(xhci, dev);
        if (!KSUCCESS(st)) {
            DbgPrint("usb: HID arm failed slot=%u ep=%u st=0x%x\n",
                     dev->SlotId, dev->HidEndpointId, st);
        }
    }
}

static void XhciHandleLegacyHidTransferEvent(XHCI_CONTROLLER* xhci, uint64_t trbDeviceAddress,
                                         uint32_t completionCode, uint8_t slotId,
                                         uint8_t endpointId, uint32_t residue) {
    if (slotId == 0 || slotId > XHCI_MAX_TRACKED_SLOTS) return;

    XHCI_DEVICE* dev = &xhci->Devices[slotId];
    if (!dev->HidActive || endpointId != dev->HidEndpointId || trbDeviceAddress != dev->HidLastTrbDeviceAddress)
        return;

    dev->HidArmed = false;
    if (completionCode != XHCI_CC_SUCCESS && completionCode != XHCI_CC_SHORT_PACKET) {
        DbgPrint("usb: HID transfer failed slot=%u ep=%u cc=%u residue=%u\n",
                 slotId, endpointId, completionCode, residue);
        return;
    }

    uint16_t length = dev->HidReportLength;
    if (residue < length) length = (uint16_t)(length - residue);

    if (dev->HidKeyboard)
        XhciProcessBootKeyboardReport(dev, dev->HidBuffer, length);
    else if (dev->HidMouse)
        XhciProcessBootMouseReport(dev, dev->HidBuffer, length);

    KSTATUS st = XhciArmHidInterruptIn(xhci, dev);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: HID re-arm failed slot=%u ep=%u st=0x%x\n",
                 slotId, endpointId, st);
    }
}

// Starts hid Interrupts.
KSTATUS XhciStartHidInterrupts(XHCI_CONTROLLER* xhci) {
    if (!xhci || xhci->HidDeviceCount == 0)
        return NTLSTATUS_SUCCESS;

    DbgPrint("usb: HID interrupt URBs active devices=%u\n", xhci->HidDeviceCount);
    return NTLSTATUS_SUCCESS;
}

// Updates hid Boot Protocol.
KSTATUS XhciSetHidBootProtocol(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    USB_SETUP_PACKET setup;
    setup.bmRequestType = USB_DIR_OUT | USB_TYPE_CLASS | USB_RECIP_INTERFACE;
    setup.bRequest = USB_REQ_SET_PROTOCOL;
    setup.wValue = 0; // boot protocol
    setup.wIndex = dev->HidInterface;
    setup.wLength = 0;

    KSTATUS st = XhciControlNoData(xhci, dev, &setup);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: HID SET_PROTOCOL boot failed slot=%u if=%u st=0x%x\n",
                 dev->SlotId, dev->HidInterface, st);
        return st;
    }

    setup.bRequest = USB_REQ_SET_IDLE;
    setup.wValue = 0;
    st = XhciControlNoData(xhci, dev, &setup);
    if (!KSUCCESS(st)) {
        // SET_IDLE is useful for keyboards but not worth rejecting an otherwise
        // usable boot device in the early stack.
        DbgPrint("usb: HID SET_IDLE failed slot=%u if=%u st=0x%x; continuing\n",
                 dev->SlotId, dev->HidInterface, st);
    }

    DbgPrint("usb: HID boot protocol selected slot=%u if=%u proto=%u\n",
             dev->SlotId, dev->HidInterface, dev->HidProtocol);
    return NTLSTATUS_SUCCESS;
}

// Creates hid Memory.
KSTATUS XhciAllocateHidMemory(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    if (dev->HidRing && dev->HidBuffer) return NTLSTATUS_SUCCESS;

    dev->HidRing = (XHCI_TRB*)UsbDmaAllocPages(xhci, 1, &dev->HidRingDeviceAddress, &dev->HidRingDma);
    dev->HidBuffer = (uint8_t*)UsbDmaAllocPages(xhci, 1, &dev->HidBufferDeviceAddress, &dev->HidBufferDma);
    if (!dev->HidRing || !dev->HidBuffer) {
        {
            IoDmaFreeCommonBuffer(dev->HidBufferDma);
            IoDmaFreeCommonBuffer(dev->HidRingDma);
        }
        dev->HidBufferDma = (IO_DMA_BUFFER*)0;
        dev->HidRingDma = (IO_DMA_BUFFER*)0;
        dev->HidBuffer = (uint8_t*)0;
        dev->HidRing = (XHCI_TRB*)0;
        return NTLSTATUS_NO_MEMORY;
    }

    dev->HidEnqueue = 0;
    dev->HidCycle = 1;

    XHCI_TRB* link = &dev->HidRing[XHCI_HID_RING_TRBS - 1u];
    link->ParameterLo = (uint32_t)dev->HidRingDeviceAddress;
    link->ParameterHi = (uint32_t)(dev->HidRingDeviceAddress >> 32);
    link->Status = 0;
    link->Control = (XHCI_TRB_TYPE_LINK << XHCI_TRB_TYPE_SHIFT) |
                    XHCI_TRB_LINK_TC |
                    dev->HidCycle;

    memset(dev->HidLastKeys, 0, sizeof(dev->HidLastKeys));
    dev->HidLastModifiers = 0;
    dev->HidLastButtons = 0;
    return NTLSTATUS_SUCCESS;
}

// Implements the configure Hid Endpoint operation.
KSTATUS XhciConfigureHidEndpoint(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    if (!dev->HidCandidate)
        return NTLSTATUS_SUCCESS;

    KSTATUS st = XhciAllocateHidMemory(xhci, dev);
    if (!KSUCCESS(st)) return st;

    uint8_t dci = dev->HidEndpointId;
    if (dci < 2 || dci > 31)
        return NTLSTATUS_INVALID_PARAMETER;

    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 = dev->HidMaxPacket & 0x07FFu;
    if (maxPacket == 0) maxPacket = dev->HidKeyboard ? 8u : 4u;
    uint8_t interval = XhciInterruptInterval(dev->Speed, dev->HidInterval);

    ep[0] = (uint32_t)interval << 16;
    ep[1] = XHCI_EP_CERR_3 | XHCI_EP_TYPE_INT_IN | ((uint32_t)maxPacket << 16);
    ep[2] = (uint32_t)(dev->HidRingDeviceAddress | XHCI_EP_DCS);
    ep[3] = (uint32_t)((dev->HidRingDeviceAddress | XHCI_EP_DCS) >> 32);
    ep[4] = (uint32_t)maxPacket | ((uint32_t)maxPacket << 16);

    XHCI_COMMAND_RESULT result;
    memset(&result, 0, sizeof(result));
    st = XhciSubmitCommand(
        xhci,
        dev->InputContextDeviceAddress,
        0,
        (XHCI_TRB_TYPE_CONFIGURE_ENDPOINT << XHCI_TRB_TYPE_SHIFT) |
            ((uint32_t)dev->SlotId << 24),
        &result);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: Configure Endpoint failed slot=%u ep=%u st=0x%x cc=%u\n",
                 dev->SlotId, dci, st, result.CompletionCode);
        return st;
    }

    dev->HidReportLength = maxPacket;
    if (dev->HidKeyboard && dev->HidReportLength < 8) dev->HidReportLength = 8;
    if (dev->HidMouse && dev->HidReportLength < 3) dev->HidReportLength = 3;
    if (dev->HidReportLength > XHCI_MAX_HID_REPORT) dev->HidReportLength = XHCI_MAX_HID_REPORT;

    st = XhciSetHidBootProtocol(xhci, dev);
    if (!KSUCCESS(st)) return st;

    dev->HidActive = true;
    xhci->HidDeviceCount++;
    uint32_t inputId = 0;
    KeInputRegisterDevice(dev->HidKeyboard ? "usb-hid-keyboard" : "usb-hid-mouse", 0, &inputId);

    DbgPrint("usb: HID boot %s bound slot=%u if=%u ep=%02x dci=%u maxpkt=%u interval=%u input=%u\n",
             dev->HidKeyboard ? "keyboard" : "mouse",
             dev->SlotId,
             dev->HidInterface,
             dev->HidEndpointAddress,
             dci,
             maxPacket,
             dev->HidInterval,
             inputId);
    return NTLSTATUS_SUCCESS;
}

// PDO/FDO boot HID function driver path.

DRIVER_OBJECT gUsbHidDriver;

static const IO_DEVICE_ID gUsbHidIds[] = {
    { IoBusUsb, { "USB\\CLASS_03&SUBCLASS_01&PROT_01", sizeof("USB\\CLASS_03&SUBCLASS_01&PROT_01") - 1 } },
    { IoBusUsb, { "USB\\CLASS_03&SUBCLASS_01&PROT_02", sizeof("USB\\CLASS_03&SUBCLASS_01&PROT_02") - 1 } },
    { IoBusUsb, { "USB\\CLASS_03&SUBCLASS_00&PROT_00", sizeof("USB\\CLASS_03&SUBCLASS_00&PROT_00") - 1 } },
};

static const IO_DRIVER_ID_TABLE gUsbHidIdTable = {
    gUsbHidIds,
    sizeof(gUsbHidIds) / sizeof(gUsbHidIds[0]),
};

// Implements the hid Is Qemu Absolute Tablet Pdo operation.
static bool UsbHidIsQemuAbsoluteTabletPdo(const USB_PDO_EXTENSION* pdoExt, const USB_ENDPOINT_DESCRIPTOR* ep) {
    if (!pdoExt || !pdoExt->Device || !ep)
        return false;
    if (pdoExt->InterfaceDescriptor.bInterfaceClass != USB_CLASS_HID)
        return false;
    if (pdoExt->InterfaceDescriptor.bInterfaceSubClass != 0 ||
        pdoExt->InterfaceDescriptor.bInterfaceProtocol != 0)
        return false;
    if (pdoExt->Device->VendorId != 0x0627u || pdoExt->Device->ProductId != 0x0001u)
        return false;

    // QEMU's usb-kbd and usb-tablet commonly share VID:PID 0627:0001.
    // Do not classify by VID:PID alone: the boot keyboard is subclass/protocol
    // 01/01, while usb-tablet is the non-boot HID interface 00/00.

    return (ep->wMaxPacketSize & 0x07FFu) >= 6u;
}

// Implements the hid Find Interrupt In Pdo operation.
static const USB_ENDPOINT_DESCRIPTOR* UsbHidFindInterruptInPdo(const USB_PDO_EXTENSION* ext) {
    if (!ext) return (const USB_ENDPOINT_DESCRIPTOR*)0;

    for (uint8_t i = 0; i < ext->EndpointCount; i++) {
        const USB_ENDPOINT_DESCRIPTOR* ep = &ext->Endpoints[i];
        uint8_t type = ep->bmAttributes & USB_ENDPOINT_ATTR_TYPE_MASK;
        bool in = (ep->bEndpointAddress & USB_ENDPOINT_DIR_IN) != 0;
        if (type == USB_ENDPOINT_TYPE_INTERRUPT && in)
            return ep;
    }
    return (const USB_ENDPOINT_DESCRIPTOR*)0;
}

// Implements the hid Process Boot Keyboard Report operation.
static void UsbHidProcessBootKeyboardReport(USB_HID_EXTENSION* ext, const uint8_t* report, uint16_t len) {
    if (!ext || len < 8) return;

    uint8_t modifiers = report[0];
    for (uint8_t bit = 0; bit < 8; bit++) {
        uint8_t mask = (uint8_t)(1u << bit);
        bool wasDown = (ext->LastModifiers & mask) != 0;
        bool isDown = (modifiers & mask) != 0;
        if (wasDown != isDown)
            XhciSubmitHidKey(XhciHidModifierKey(bit), isDown);
    }

    const uint8_t* keys = &report[2];
    for (uint8_t i = 0; i < 6; i++) {
        uint8_t oldKey = ext->LastKeys[i];
        if (oldKey > 1 && !XhciHidKeyInReport(keys, oldKey))
            XhciSubmitHidKey((INPUT_KEY)oldKey, false);
    }
    for (uint8_t i = 0; i < 6; i++) {
        uint8_t newKey = keys[i];
        if (newKey > 1 && !XhciHidKeyInReport(ext->LastKeys, newKey))
            XhciSubmitHidKey((INPUT_KEY)newKey, true);
    }

    ext->LastModifiers = modifiers;
    memcpy(ext->LastKeys, keys, sizeof(ext->LastKeys));
}

// Implements the hid Process Boot Mouse Report operation.
static void UsbHidProcessBootMouseReport(USB_HID_EXTENSION* ext, const uint8_t* report, uint16_t len) {
    if (!ext || len < 3) return;

    int32_t dx = (int32_t)(int8_t)report[1];
    int32_t dy = (int32_t)(int8_t)report[2];
    int32_t wheel = (len >= 4) ? (int32_t)(int8_t)report[3] : 0;
    uint8_t buttons = report[0] & 0x07u;

    if (dx || dy || wheel || buttons != ext->LastButtons) {
        INPUT_MOUSE_EVENT ev;
        ev.dx = dx;
        ev.dy = dy;
        ev.wheel = wheel;
        ev.buttons = buttons;
        KeInputSubmitMouse(&ev);
    }

    ext->LastButtons = buttons;
}

// Implements the hid Process Qemu Tablet Report operation.
static void UsbHidProcessQemuTabletReport(USB_HID_EXTENSION* ext, const uint8_t* report, uint16_t len) {
    if (!ext || len < 5) return;

    uint8_t buttons = report[0] & 0x07u;
    uint16_t x = (uint16_t)report[1] | ((uint16_t)report[2] << 8);
    uint16_t y = (uint16_t)report[3] | ((uint16_t)report[4] << 8);
    int32_t wheel = (len >= 6) ? (int32_t)(int8_t)report[5] : 0;

    if (x != ext->LastAbsoluteX || y != ext->LastAbsoluteY ||
        wheel || buttons != ext->LastButtons) {
        KeInputSubmitMouseAbsoluteScaled(x, y, ext->AbsoluteMaxX, ext->AbsoluteMaxY, wheel, buttons);
    }

    ext->LastAbsoluteX = x;
    ext->LastAbsoluteY = y;
    ext->LastButtons = buttons;
}

static KSTATUS UsbHidArmInterrupt(USB_HID_EXTENSION* ext);

// Implements the hid Interrupt Complete operation.
static void UsbHidInterruptComplete(USB_URB* urb, void* context) {
    USB_HID_EXTENSION* ext = (USB_HID_EXTENSION*)context;
    if (!ext || urb != &ext->InterruptUrb)
        return;

    ext->InterruptPending = false;
    if (ext->SurpriseRemoved || ext->Removed || !ext->Started) {
        IoReleaseRemoveLock(&ext->RemoveLock, (IRP*)0);
        return;
    }

    if (KSUCCESS(urb->Status)) {
        uint16_t length = (uint16_t)urb->ActualLength;
        if (length > ext->ReportLength) length = ext->ReportLength;

        if (ext->Keyboard)
            UsbHidProcessBootKeyboardReport(ext, ext->ReportBuffer, length);
        else if (ext->AbsoluteMouse)
            UsbHidProcessQemuTabletReport(ext, ext->ReportBuffer, length);
        else if (ext->Mouse)
            UsbHidProcessBootMouseReport(ext, ext->ReportBuffer, length);
    } else if (urb->Status != NTLSTATUS_CANCELLED) {
        DbgPrint("usb: HID interrupt URB failed slot=%u if=%u ep=%02x st=0x%x cc=%u actual=%u/%u\n",
                 ext->Device ? ext->Device->SlotId : 0,
                 ext->InterfaceNumber,
                 ext->EndpointAddress,
                 urb->Status,
                 urb->HcdCompletionCode,
                 urb->ActualLength,
                 urb->Length);
    }

    if (!ext->SurpriseRemoved && !ext->Removed)
        (void)UsbHidArmInterrupt(ext);
    IoReleaseRemoveLock(&ext->RemoveLock, (IRP*)0);
}

// Implements the hid Send Boot Protocol operation.
static KSTATUS UsbHidSendBootProtocol(USB_HID_EXTENSION* ext) {
    if (!ext || !ext->LowerDevice || !ext->Device)
        return NTLSTATUS_INVALID_PARAMETER;

    USB_URB urb;
    memset(&urb, 0, sizeof(urb));
    urb.Function = UrbFunctionControlTransfer;
    urb.Direction = UsbTransferNoData;
    urb.Device = ext->Device;
    urb.Setup.bmRequestType = USB_DIR_OUT | USB_TYPE_CLASS | USB_RECIP_INTERFACE;
    urb.Setup.bRequest = USB_REQ_SET_PROTOCOL;
    urb.Setup.wValue = 0;
    urb.Setup.wIndex = ext->InterfaceNumber;
    urb.Setup.wLength = 0;

    KSTATUS st = UsbSubmitUrbSync(ext->LowerDevice, &urb);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: HID SET_PROTOCOL boot failed slot=%u if=%u st=0x%x\n",
                 ext->Device->SlotId, ext->InterfaceNumber, st);
        return st;
    }

    memset(&urb, 0, sizeof(urb));
    urb.Function = UrbFunctionControlTransfer;
    urb.Direction = UsbTransferNoData;
    urb.Device = ext->Device;
    urb.Setup.bmRequestType = USB_DIR_OUT | USB_TYPE_CLASS | USB_RECIP_INTERFACE;
    urb.Setup.bRequest = USB_REQ_SET_IDLE;
    urb.Setup.wValue = 0;
    urb.Setup.wIndex = ext->InterfaceNumber;
    urb.Setup.wLength = 0;

    st = UsbSubmitUrbSync(ext->LowerDevice, &urb);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: HID SET_IDLE failed slot=%u if=%u st=0x%x; continuing\n",
                 ext->Device->SlotId, ext->InterfaceNumber, st);
    }

    DbgPrint("usb: HID boot protocol selected slot=%u if=%u proto=%u\n",
             ext->Device->SlotId, ext->InterfaceNumber, ext->Protocol);
    return NTLSTATUS_SUCCESS;
}

// Implements the hid Arm Interrupt operation.
static KSTATUS UsbHidArmInterrupt(USB_HID_EXTENSION* ext) {
    if (!ext || !ext->Started || ext->SurpriseRemoved || ext->Removed)
        return NTLSTATUS_OBJECT_DEAD;
    if (!ext->LowerDevice || !ext->InterruptPipe || !ext->ReportBuffer || !ext->ReportBufferDeviceAddress)
        return NTLSTATUS_INVALID_PARAMETER;
    if (ext->InterruptPending)
        return NTLSTATUS_SUCCESS;

    // One in-flight callback owns a remove reference through dispatch or cancel.
    KSTATUS refStatus = IoAcquireRemoveLock(&ext->RemoveLock, (IRP*)0);
    if (!KSUCCESS(refStatus))
        return refStatus;
    ext->InterruptPending = true;
    memset(ext->ReportBuffer, 0, ext->ReportLength);
    memset(&ext->InterruptUrb, 0, sizeof(ext->InterruptUrb));
    ext->InterruptUrb.Function = UrbFunctionInterruptTransfer;
    ext->InterruptUrb.Direction = UsbTransferIn;
    ext->InterruptUrb.Device = ext->Device;
    ext->InterruptUrb.Pipe = ext->InterruptPipe;
    ext->InterruptUrb.Buffer = ext->ReportBuffer;
    ext->InterruptUrb.BufferDeviceAddress = ext->ReportBufferDeviceAddress;
    ext->InterruptUrb.Length = ext->ReportLength;
    ext->InterruptUrb.Complete = UsbHidInterruptComplete;
    ext->InterruptUrb.CompleteContext = ext;

    KSTATUS st = UsbSubmitUrbAsync(ext->LowerDevice, &ext->InterruptUrb);
    // The USB PDO returns PENDING even when the HCD rejected the URB and
    // completed its IRP synchronously. Inspect the URB's actual status.
    if (st != NTLSTATUS_PENDING || ext->InterruptUrb.Status != NTLSTATUS_PENDING) {
        ext->InterruptPending = false;
        IoReleaseRemoveLock(&ext->RemoveLock, (IRP*)0);
        return st == NTLSTATUS_PENDING ? ext->InterruptUrb.Status : st;
    }
    return NTLSTATUS_SUCCESS;
}

// Implements the hid Start Device operation.
static KSTATUS UsbHidStartDevice(DEVICE_OBJECT* device, USB_HID_EXTENSION* ext) {
    if (!device || !ext || !ext->Device || !ext->InterfacePdo)
        return NTLSTATUS_INVALID_PARAMETER;
    if (ext->Started)
        return NTLSTATUS_SUCCESS;
    if (ext->SurpriseRemoved || ext->Removed)
        return NTLSTATUS_OBJECT_DEAD;

    USB_PDO_EXTENSION* pdoExt = (USB_PDO_EXTENSION*)ext->InterfacePdo->DeviceExtension;
    if (!pdoExt || pdoExt->Kind != UsbPdoInterface)
        return NTLSTATUS_INVALID_PARAMETER;

    const USB_ENDPOINT_DESCRIPTOR* ep = UsbHidFindInterruptInPdo(pdoExt);
    if (!ep)
        return NTLSTATUS_NOT_FOUND;

    ext->Keyboard = (pdoExt->InterfaceDescriptor.bInterfaceSubClass == USB_HID_SUBCLASS_BOOT &&
                     pdoExt->InterfaceDescriptor.bInterfaceProtocol == USB_HID_PROTOCOL_KEYBOARD);
    ext->AbsoluteMouse = UsbHidIsQemuAbsoluteTabletPdo(pdoExt, ep);
    ext->Mouse = ext->AbsoluteMouse ||
                 (pdoExt->InterfaceDescriptor.bInterfaceSubClass == USB_HID_SUBCLASS_BOOT &&
                  pdoExt->InterfaceDescriptor.bInterfaceProtocol == USB_HID_PROTOCOL_MOUSE);
    ext->AbsoluteMaxX = 0x7FFFu;
    ext->AbsoluteMaxY = 0x7FFFu;
    ext->InterfaceNumber = pdoExt->InterfaceDescriptor.bInterfaceNumber;
    ext->Protocol = pdoExt->InterfaceDescriptor.bInterfaceProtocol;
    ext->EndpointAddress = ep->bEndpointAddress;
    ext->MaxPacketSize = ep->wMaxPacketSize;
    ext->Interval = ep->bInterval;

    uint16_t reportLength = ext->MaxPacketSize & 0x07FFu;
    if (reportLength == 0) reportLength = ext->Keyboard ? 8u : (ext->AbsoluteMouse ? 6u : 4u);
    if (ext->Keyboard && reportLength < 8) reportLength = 8;
    if (ext->AbsoluteMouse && reportLength < 6) reportLength = 6;
    if (ext->Mouse && !ext->AbsoluteMouse && reportLength < 3) reportLength = 3;
    if (reportLength > XHCI_MAX_HID_REPORT) reportLength = XHCI_MAX_HID_REPORT;
    ext->ReportLength = reportLength;

    ext->ReportBuffer = (uint8_t*)UsbAllocateDmaBuffer(ext->Device, 1, &ext->ReportBufferDeviceAddress, &ext->ReportBufferDma);
    if (!ext->ReportBuffer)
        return NTLSTATUS_NO_MEMORY;

    KSTATUS st = UsbCreatePipe(ext->Device,
                               UsbPipeInterrupt,
                               ext->EndpointAddress,
                               ext->MaxPacketSize,
                               ext->Interval,
                               &ext->InterruptPipe);
    if (!KSUCCESS(st)) {
        UsbHidReleaseDma(ext);
        return st;
    }

    if (!ext->AbsoluteMouse) {
        st = UsbHidSendBootProtocol(ext);
        if (!KSUCCESS(st)) {
            if (ext->InterruptPipe) {
                ext->InterruptPipe->Open = false;
                UsbDereferencePipe(ext->InterruptPipe);
                ext->InterruptPipe = (USB_PIPE*)0;
            }
            UsbHidReleaseDma(ext);
            return st;
        }
    }

    memset(ext->LastKeys, 0, sizeof(ext->LastKeys));
    ext->LastModifiers = 0;
    ext->LastButtons = 0;

    ext->Started = true;
    device->PnpState = DevicePnpStarted;
    if (ext->InterfacePdo)
        ext->InterfacePdo->PnpState = DevicePnpStarted;

    XHCI_CONTROLLER* xhci = ext->Device && ext->Device->Hcd ?
        (XHCI_CONTROLLER*)ext->Device->Hcd->ControllerContext : (XHCI_CONTROLLER*)0;

    st = UsbHidArmInterrupt(ext);
    if (!KSUCCESS(st)) {
        ext->Started = false;
        device->PnpState = DevicePnpStopped;
        if (ext->InterfacePdo)
            ext->InterfacePdo->PnpState = DevicePnpStopped;
        if (ext->InterruptPipe) {
            ext->InterruptPipe->Open = false;
            UsbDereferencePipe(ext->InterruptPipe);
            ext->InterruptPipe = (USB_PIPE*)0;
        }
        UsbHidReleaseDma(ext);
        return st;
    }

    KeInputRegisterDevice(ext->Keyboard ? "usb-hid-keyboard" : (ext->AbsoluteMouse ? "usb-hid-tablet" : "usb-hid-mouse"), 0, &ext->InputDeviceId);
    if (xhci)
        xhci->HidDeviceCount++;

    DbgPrint("usb: HID %s bound slot=%u if=%u ep=%02x maxpkt=%u interval=%u input=%u urb\n",
             ext->Keyboard ? "boot-keyboard" : (ext->AbsoluteMouse ? "qemu-tablet" : "boot-mouse"),
             ext->Device->SlotId,
             ext->InterfaceNumber,
             ext->EndpointAddress,
             ext->MaxPacketSize,
             ext->Interval,
             ext->InputDeviceId);
    return NTLSTATUS_SUCCESS;
}

// Implements the hid Dispatch Pnp operation.
KSTATUS UsbHidDispatchPnp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp)
        return NTLSTATUS_INVALID_PARAMETER;

    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios || ios->MajorFunction != IRP_MJ_PNP)
        return NTLSTATUS_INVALID_PARAMETER;

    USB_HID_EXTENSION* ext = (USB_HID_EXTENSION*)device->DeviceExtension;
    KSTATUS st = NTLSTATUS_SUCCESS;

    switch ((IO_PNP_MINOR_FUNCTION)ios->MinorFunction) {
        case IRP_MN_START_DEVICE:
            st = UsbHidStartDevice(device, ext);
            if (!KSUCCESS(st)) {
                IoCompleteRequest(irp, st, 0);
                return NTLSTATUS_PENDING;
            }
            return IoForwardIrpToLower(device, irp);

        case IRP_MN_QUERY_REMOVE_DEVICE:
            device->PnpState = DevicePnpRemovePending;
            return IoForwardIrpToLower(device, irp);

        case IRP_MN_CANCEL_REMOVE_DEVICE:
            if (!ext || !ext->SurpriseRemoved)
                device->PnpState = DevicePnpStarted;
            return IoForwardIrpToLower(device, irp);

        case IRP_MN_SURPRISE_REMOVAL:
            device->PnpState = DevicePnpSurpriseRemoved;
            if (ext) {
                ext->SurpriseRemoved = true;
                ext->Started = false;
                if (ext->InterruptPending) {
                    st = UsbCancelUrb(&ext->InterruptUrb);
                    if (!KSUCCESS(st)) {
                        IoCompleteRequest(irp, st, 0);
                        return NTLSTATUS_PENDING;
                    }
                }
            }
            return IoForwardIrpToLower(device, irp);

        case IRP_MN_REMOVE_DEVICE: {
            device->PnpState = DevicePnpDeleted;
            if (ext) {
                ext->Removed = true;
                ext->Started = false;
                if (ext->InterruptPending) {
                    st = UsbCancelUrb(&ext->InterruptUrb);
                    if (!KSUCCESS(st)) {
                        IoCompleteRequest(irp, st, 0);
                        return NTLSTATUS_PENDING;
                    }
                }
                if (ext->InterruptPipe)
                    ext->InterruptPipe->Open = false;
                if (ext->Device && ext->Device->Hcd) {
                    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)ext->Device->Hcd->ControllerContext;
                    if (xhci && xhci->HidDeviceCount)
                        xhci->HidDeviceCount--;
                }
                if (ext->RemoveLockReady) {
                    IoReleaseRemoveLockAndWait(&ext->RemoveLock, irp);
                    IoDestroyRemoveLock(&ext->RemoveLock);
                    ext->RemoveLockReady = false;
                }
                if (ext->InterruptPipe) {
                    UsbDereferencePipe(ext->InterruptPipe);
                    ext->InterruptPipe = (USB_PIPE*)0;
                }
                UsbHidReleaseDma(ext);
            }

            KSTATUS lower = ext && ext->LowerDevice ? IoForwardIrpToLower(device, irp) : NTLSTATUS_SUCCESS;
            if (ext && ext->LowerDevice) {
                IoDetachDevice(device);
                ext->LowerDevice = (DEVICE_OBJECT*)0;
            }
            IoDeleteDevice(device);
            if (lower == NTLSTATUS_NOT_SUPPORTED) {
                IoCompleteRequest(irp, NTLSTATUS_SUCCESS, 0);
                return NTLSTATUS_PENDING;
            }
            if (lower != NTLSTATUS_PENDING && !KSUCCESS(lower)) {
                IoCompleteRequest(irp, lower, 0);
                return NTLSTATUS_PENDING;
            }
            return lower;
        }

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

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

    USB_PDO_EXTENSION* pdoExt = (USB_PDO_EXTENSION*)physicalDeviceObject->DeviceExtension;
    if (!pdoExt || pdoExt->Kind != UsbPdoInterface || !pdoExt->Device)
        return NTLSTATUS_INVALID_PARAMETER;
    if (pdoExt->InterfaceDescriptor.bInterfaceClass != USB_CLASS_HID)
        return NTLSTATUS_INVALID_PARAMETER;

    const USB_ENDPOINT_DESCRIPTOR* ep = UsbHidFindInterruptInPdo(pdoExt);
    if (!ep)
        return NTLSTATUS_NOT_FOUND;

    bool bootKeyboard = pdoExt->InterfaceDescriptor.bInterfaceSubClass == USB_HID_SUBCLASS_BOOT &&
                        pdoExt->InterfaceDescriptor.bInterfaceProtocol == USB_HID_PROTOCOL_KEYBOARD;
    bool bootMouse = pdoExt->InterfaceDescriptor.bInterfaceSubClass == USB_HID_SUBCLASS_BOOT &&
                     pdoExt->InterfaceDescriptor.bInterfaceProtocol == USB_HID_PROTOCOL_MOUSE;
    bool qemuTablet = UsbHidIsQemuAbsoluteTabletPdo(pdoExt, ep);
    if (!bootKeyboard && !bootMouse && !qemuTablet) {
        // A nonboot HID interface requires its report descriptor to be parsed.
        // Avoid treating opaque reports as boot mouse packets.
        DbgPrint("usb: HID slot=%u if=%u subclass=%u protocol=%u requires report parser\n",
                 pdoExt->Device->SlotId,
                 pdoExt->InterfaceDescriptor.bInterfaceNumber,
                 pdoExt->InterfaceDescriptor.bInterfaceSubClass,
                 pdoExt->InterfaceDescriptor.bInterfaceProtocol);
        return NTLSTATUS_NOT_SUPPORTED;
    }

    DEVICE_OBJECT* fdo = (DEVICE_OBJECT*)0;
    KSTATUS st = IoCreateDevice(driver, sizeof(USB_HID_EXTENSION), &fdo);
    if (!KSUCCESS(st))
        return st;

    (void)IoSetDeviceName(fdo,
                          bootKeyboard
                              ? SYS_SV_LITERAL("usb-hid-keyboard")
                              : (qemuTablet ? SYS_SV_LITERAL("usb-hid-tablet")
                                            : SYS_SV_LITERAL("usb-hid-mouse")));

    USB_HID_EXTENSION* ext = (USB_HID_EXTENSION*)fdo->DeviceExtension;
    ext->Self = fdo;
    ext->InterfacePdo = physicalDeviceObject;
    ext->Device = pdoExt->Device;
    st = IoInitializeRemoveLock(&ext->RemoveLock);
    if (!KSUCCESS(st)) {
        IoDeleteDevice(fdo);
        return st;
    }
    ext->RemoveLockReady = true;

    ext->LowerDevice = IoAttachDeviceToDeviceStack(fdo, physicalDeviceObject);
    if (!ext->LowerDevice) {
        IoDestroyRemoveLock(&ext->RemoveLock);
        ext->RemoveLockReady = false;
        IoDeleteDevice(fdo);
        return NTLSTATUS_FAILURE;
    }

    DbgPrint("usb: HID FDO attached slot=%u if=%u proto=%u kind=%s\n",
             ext->Device->SlotId,
             pdoExt->InterfaceDescriptor.bInterfaceNumber,
             pdoExt->InterfaceDescriptor.bInterfaceProtocol,
             qemuTablet ? "qemu-tablet" : (bootKeyboard ? "keyboard" : "mouse"));
    return NTLSTATUS_SUCCESS;
}

// Initializes hid Driver.
KSTATUS UsbInitializeHidDriver(void) {
    memset(&gUsbHidDriver, 0, sizeof(gUsbHidDriver));
    IoInitializeDriverObject(&gUsbHidDriver, SYS_SV_LITERAL("usb-hid"));
    gUsbHidDriver.IdTable = &gUsbHidIdTable;
    gUsbHidDriver.DriverAddDevice = UsbHidAddDevice;
    gUsbHidDriver.MajorFunction[IRP_MJ_PNP] = UsbHidDispatchPnp;
    return NTLSTATUS_SUCCESS;
}