// usb_hid.c - USB HID boot keyboard/mouse class binding and report handling.
#include "usb.h"

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;
    }
}


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);
}


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));
}


void XhciProcessBootMouseReport(XHCI_DEVICE* dev, const uint8_t* report, uint16_t len) {
    if (len < 3) return;

    int dx = (int)(int8_t)report[1];
    int dy = (int)(int8_t)report[2];
    int wheel = (len >= 4) ? (int)(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;
}


KSTATUS XhciArmHidInterruptIn(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    if (!dev->HidActive || dev->HidArmed || !dev->HidRing || !dev->HidBuffer)
        return STATUS_SUCCESS;

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

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

    dev->HidLastTrbPhys = trbPhys;
    dev->HidArmed = true;
    XhciBarrier();
    XhciWrite32(xhci->Doorbells, (uint32_t)dev->SlotId * 4u, dev->HidEndpointId);
    return STATUS_SUCCESS;
}


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 trbPhys,
                                         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 || trbPhys != dev->HidLastTrbPhys)
        return;

    dev->HidArmed = false;
    if (completionCode != XHCI_CC_SUCCESS) {
        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);
    }
}


KSTATUS XhciStartHidInterrupts(XHCI_CONTROLLER* xhci) {
    if (!xhci || xhci->HidDeviceCount == 0)
        return STATUS_SUCCESS;

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


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 STATUS_SUCCESS;
}


KSTATUS XhciAllocateHidMemory(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    if (dev->HidRing && dev->HidBuffer) return STATUS_SUCCESS;

    dev->HidRing = (XHCI_TRB*)UsbDmaAllocPages(xhci, 1, &dev->HidRingPhys);
    dev->HidBuffer = (uint8_t*)UsbDmaAllocPages(xhci, 1, &dev->HidBufferPhys);
    if (!dev->HidRing || !dev->HidBuffer)
        return STATUS_NO_MEMORY;

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

    XHCI_TRB* link = &dev->HidRing[XHCI_HID_RING_TRBS - 1u];
    link->ParameterLo = (uint32_t)dev->HidRingPhys;
    link->ParameterHi = (uint32_t)(dev->HidRingPhys >> 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 STATUS_SUCCESS;
}


KSTATUS XhciConfigureHidEndpoint(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev) {
    if (!dev->HidCandidate)
        return STATUS_SUCCESS;

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

    uint8_t dci = dev->HidEndpointId;
    if (dci < 2 || dci > 31)
        return STATUS_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);

    slot[0] = XHCI_CONTEXT_ENTRIES(dci) | XHCI_SLOT_SPEED(dev->Speed);
    slot[1] = XHCI_SLOT_ROOT_PORT(dev->Port);
    slot[2] = 0;
    slot[3] = 0;

    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->HidRingPhys | XHCI_EP_DCS);
    ep[3] = (uint32_t)((dev->HidRingPhys | 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->InputContextPhys,
        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 STATUS_SUCCESS;
}


const IO_PUBLISHED_USB_ENDPOINT* UsbFindEndpoint(const IO_PUBLISHED_DEVICE* pub,
                                                        uint8_t endpointType,
                                                        bool wantIn) {
    if (!pub || pub->BusType != IoBusUsb) return (const IO_PUBLISHED_USB_ENDPOINT*)0;

    for (uint8_t i = 0; i < pub->u.Usb.EndpointCount; i++) {
        const IO_PUBLISHED_USB_ENDPOINT* ep = &pub->u.Usb.Endpoints[i];
        uint8_t type = ep->Attributes & USB_ENDPOINT_ATTR_TYPE_MASK;
        bool in = (ep->Address & USB_ENDPOINT_DIR_IN) != 0;
        if (type == endpointType && in == wantIn)
            return ep;
    }

    return (const IO_PUBLISHED_USB_ENDPOINT*)0;
}


bool UsbHidMatch(const IO_PUBLISHED_DEVICE* pub, void* context) {
    (void)context;
    if (!pub || pub->BusType != IoBusUsb)
        return false;
    if (pub->u.Usb.InterfaceClass != USB_CLASS_HID)
        return false;
    if (pub->u.Usb.InterfaceSubclass != USB_HID_SUBCLASS_BOOT)
        return false;
    if (pub->u.Usb.InterfaceProtocol != USB_HID_PROTOCOL_KEYBOARD &&
        pub->u.Usb.InterfaceProtocol != USB_HID_PROTOCOL_MOUSE)
        return false;

    return UsbFindEndpoint(pub, USB_ENDPOINT_TYPE_INTERRUPT, true) != (const IO_PUBLISHED_USB_ENDPOINT*)0;
}


KSTATUS UsbHidBind(IO_PUBLISHED_DEVICE* pub, void* context) {
    (void)context;
    if (!pub || !pub->BusContext || !pub->DeviceContext)
        return STATUS_INVALID_PARAMETER;

    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)pub->BusContext;
    XHCI_DEVICE* dev = (XHCI_DEVICE*)pub->DeviceContext;

    if (dev->HidActive || dev->HidCandidate)
        return STATUS_ALREADY_EXISTS;

    const IO_PUBLISHED_USB_ENDPOINT* ep = UsbFindEndpoint(pub, USB_ENDPOINT_TYPE_INTERRUPT, true);
    if (!ep)
        return STATUS_NOT_FOUND;

    dev->HidCandidate = true;
    dev->HidKeyboard = (pub->u.Usb.InterfaceProtocol == USB_HID_PROTOCOL_KEYBOARD);
    dev->HidMouse = (pub->u.Usb.InterfaceProtocol == USB_HID_PROTOCOL_MOUSE);
    dev->HidInterface = pub->u.Usb.InterfaceNumber;
    dev->HidProtocol = pub->u.Usb.InterfaceProtocol;
    dev->HidEndpointAddress = ep->Address;
    dev->HidEndpointId = XhciEndpointIdFromAddress(ep->Address);
    dev->HidInterval = ep->Interval;
    dev->HidMaxPacket = ep->MaxPacketSize;

    KSTATUS st = XhciConfigureHidEndpoint(xhci, dev);
    if (!KSUCCESS(st)) {
        dev->HidCandidate = false;
        dev->HidKeyboard = false;
        dev->HidMouse = false;
    }
    return st;
}

// -------------------------------------------------------------------------
// 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 } },
};

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

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;
}


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));
}

static void UsbHidProcessBootMouseReport(USB_HID_EXTENSION* ext, const uint8_t* report, uint16_t len) {
    if (!ext || len < 3) return;

    int dx = (int)(int8_t)report[1];
    int dy = (int)(int8_t)report[2];
    int wheel = (len >= 4) ? (int)(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;
}

static KSTATUS UsbHidArmInterrupt(USB_HID_EXTENSION* ext);

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)
        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->Mouse)
            UsbHidProcessBootMouseReport(ext, ext->ReportBuffer, length);
    } else if (urb->Status != STATUS_CANCELLED) {
        DbgPrint("usb: HID interrupt URB failed slot=%u if=%u ep=%02x st=0x%x\n",
                 ext->Device ? ext->Device->SlotId : 0,
                 ext->InterfaceNumber,
                 ext->EndpointAddress,
                 urb->Status);
    }

    if (!ext->SurpriseRemoved && !ext->Removed)
        (void)UsbHidArmInterrupt(ext);
}

static KSTATUS UsbHidSendBootProtocol(USB_HID_EXTENSION* ext) {
    if (!ext || !ext->LowerDevice || !ext->Device)
        return STATUS_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 STATUS_SUCCESS;
}

static KSTATUS UsbHidArmInterrupt(USB_HID_EXTENSION* ext) {
    if (!ext || !ext->Started || ext->SurpriseRemoved || ext->Removed)
        return STATUS_OBJECT_DEAD;
    if (!ext->LowerDevice || !ext->InterruptPipe || !ext->ReportBuffer || !ext->ReportBufferPhys)
        return STATUS_INVALID_PARAMETER;
    if (ext->InterruptPending)
        return STATUS_SUCCESS;

    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.BufferPhys = ext->ReportBufferPhys;
    ext->InterruptUrb.Length = ext->ReportLength;
    ext->InterruptUrb.Complete = UsbHidInterruptComplete;
    ext->InterruptUrb.CompleteContext = ext;

    KSTATUS st = UsbSubmitUrbAsync(ext->LowerDevice, &ext->InterruptUrb);
    if (st != STATUS_PENDING && !KSUCCESS(st))
        return st;
    if (ext->InterruptUrb.Status != STATUS_PENDING && ext->InterruptUrb.Status != STATUS_SUCCESS)
        return ext->InterruptUrb.Status;

    ext->InterruptPending = true;
    return STATUS_SUCCESS;
}

static KSTATUS UsbHidStartDevice(DEVICE_OBJECT* device, USB_HID_EXTENSION* ext) {
    if (!device || !ext || !ext->Device || !ext->InterfacePdo)
        return STATUS_INVALID_PARAMETER;
    if (ext->Started)
        return STATUS_SUCCESS;
    if (ext->SurpriseRemoved || ext->Removed)
        return STATUS_OBJECT_DEAD;

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

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

    ext->Keyboard = (pdoExt->InterfaceDescriptor.bInterfaceProtocol == USB_HID_PROTOCOL_KEYBOARD);
    ext->Mouse = (pdoExt->InterfaceDescriptor.bInterfaceProtocol == USB_HID_PROTOCOL_MOUSE);
    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 : 4u;
    if (ext->Keyboard && reportLength < 8) reportLength = 8;
    if (ext->Mouse && 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->ReportBufferPhys);
    if (!ext->ReportBuffer)
        return STATUS_NO_MEMORY;

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

    st = UsbHidSendBootProtocol(ext);
    if (!KSUCCESS(st))
        return st;

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

    KeInputRegisterDevice(ext->Keyboard ? "usb-hid-keyboard" : "usb-hid-mouse", 0, &ext->InputDeviceId);

    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))
        return st;

    if (xhci)
        xhci->HidDeviceCount++;

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


KSTATUS UsbHidDispatchPnp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp)
        return STATUS_INVALID_PARAMETER;

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

    USB_HID_EXTENSION* ext = (USB_HID_EXTENSION*)device->DeviceExtension;
    KSTATUS st = STATUS_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 STATUS_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)
                    (void)UsbCancelUrb(&ext->InterruptUrb);
            }
            return IoForwardIrpToLower(device, irp);

        case IRP_MN_REMOVE_DEVICE: {
            device->PnpState = DevicePnpDeleted;
            if (ext) {
                ext->Removed = true;
                ext->Started = false;
                if (ext->InterruptPending)
                    (void)UsbCancelUrb(&ext->InterruptUrb);
                if (ext->InterruptPipe) {
                    ext->InterruptPipe->Open = false;
                    UsbDereferencePipe(ext->InterruptPipe);
                    ext->InterruptPipe = (USB_PIPE*)0;
                }
                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;
                }
            }

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

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

KSTATUS UsbHidAddDevice(DRIVER_OBJECT* driver, DEVICE_OBJECT* physicalDeviceObject) {
    if (!driver || !physicalDeviceObject || physicalDeviceObject->Type != DeviceObjectPhysical)
        return STATUS_INVALID_PARAMETER;

    USB_PDO_EXTENSION* pdoExt = (USB_PDO_EXTENSION*)physicalDeviceObject->DeviceExtension;
    if (!pdoExt || pdoExt->Kind != UsbPdoInterface || !pdoExt->Device)
        return STATUS_INVALID_PARAMETER;
    if (pdoExt->InterfaceDescriptor.bInterfaceClass != USB_CLASS_HID ||
        pdoExt->InterfaceDescriptor.bInterfaceSubClass != USB_HID_SUBCLASS_BOOT)
        return STATUS_INVALID_PARAMETER;
    if (pdoExt->InterfaceDescriptor.bInterfaceProtocol != USB_HID_PROTOCOL_KEYBOARD &&
        pdoExt->InterfaceDescriptor.bInterfaceProtocol != USB_HID_PROTOCOL_MOUSE)
        return STATUS_INVALID_PARAMETER;
    if (!UsbHidFindInterruptInPdo(pdoExt))
        return STATUS_NOT_FOUND;

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

    (void)IoSetDeviceName(fdo,
                          pdoExt->InterfaceDescriptor.bInterfaceProtocol == USB_HID_PROTOCOL_KEYBOARD
                              ? SYS_SV_LITERAL("usb-hid-keyboard")
                              : 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 STATUS_FAILURE;
    }

    DbgPrint("usb: HID FDO attached slot=%u if=%u proto=%u\n",
             ext->Device->SlotId,
             pdoExt->InterfaceDescriptor.bInterfaceNumber,
             pdoExt->InterfaceDescriptor.bInterfaceProtocol);
    return STATUS_SUCCESS;
}

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 IoRegisterDriverObject(&gUsbHidDriver);
}