// usb_hub.c - USB hub function driver.
//
// Root hubs and external hubs are represented by the same FDO abstraction.  The
// xHCI root hub feeds port-change events directly into this driver, while an
// external hub feeds the same port state machine from its interrupt-IN change
// bitmap.
#include "usb.h"

DRIVER_OBJECT gUsbHubDriver;

static const IO_DEVICE_ID gUsbHubIds[] = {
    { IoBusUsb, { "USB\\ROOT_HUB", sizeof("USB\\ROOT_HUB") - 1 } },
    { IoBusUsb, { "USB\\CLASS_09", sizeof("USB\\CLASS_09") - 1 } },
};

static const IO_DRIVER_ID_TABLE gUsbHubIdTable = {
    gUsbHubIds,
    sizeof(gUsbHubIds) / sizeof(gUsbHubIds[0]),
};

#define USB_HUB_PORT_ENUMERATING 0xffu
#define USB_HUB_PORT_FAILED      0xfeu

static bool UsbHubSlotBusy(uint8_t slot) {
    return slot != 0 && slot != USB_HUB_PORT_ENUMERATING && slot != USB_HUB_PORT_FAILED;
}

static bool UsbHubRootPortConnected(XHCI_CONTROLLER* xhci, uint8_t port) {
    if (!xhci || !xhci->Op || port == 0) return false;
    uint32_t portsc = XhciRead32(xhci->Op, XhciPortReg(port));
    return (portsc & XHCI_PORTSC_CCS) != 0;
}

static void UsbHubClearRootPortFailure(USB_HUB_EXTENSION* hub, uint8_t port) {
    if (!hub || port == 0 || port > USB_ROOT_HUB_MAX_PORTS) return;
    hub->RootPortFailCount[port] = 0;
    if (hub->RootPortSlot[port] == USB_HUB_PORT_FAILED)
        hub->RootPortSlot[port] = 0;
}

static void UsbHubMarkRootPortFailed(USB_HUB_EXTENSION* hub, uint8_t port) {
    if (!hub || port == 0 || port > USB_ROOT_HUB_MAX_PORTS) return;
    hub->RootPortSlot[port] = USB_HUB_PORT_FAILED;
    hub->RootPortState[port] = UsbHubPortFailedConnected;
}

static KSTATUS UsbHubEnumerateRootPortWithRetries(USB_HUB_EXTENSION* hub,
                                                   XHCI_CONTROLLER* xhci,
                                                   uint8_t port,
                                                   uint8_t* outSlot) {
    if (outSlot) *outSlot = 0;
    if (!hub || !xhci || port == 0 || port > USB_ROOT_HUB_MAX_PORTS)
        return STATUS_INVALID_PARAMETER;

    KSTATUS last = STATUS_DEVICE_ERROR;
    uint8_t maxAttempts = USB_ROOT_PORT_ENUM_RETRIES;
    if (maxAttempts == 0) maxAttempts = 1;

    for (uint8_t attempt = 1; attempt <= maxAttempts; attempt++) {
        if (!UsbHubRootPortConnected(xhci, port))
            return STATUS_NOT_FOUND;

        uint8_t slot = 0;
        last = XhciEnumeratePortEx(xhci, port, &slot);
        if (KSUCCESS(last)) {
            if (outSlot) *outSlot = slot;
            hub->RootPortFailCount[port] = 0;
            hub->RootPortState[port] = UsbHubPortChildPresent;
            return STATUS_SUCCESS;
        }

        hub->RootPortFailCount[port] = attempt;

        /*
         * On real Intel laptop controllers some internal full-speed devices
         * respond with xHCI USB Transaction Error during Address Device while
         * the port is still settling.  Retry that boundedly while the port is
         * held in ENUMERATING state so queued PSC events cannot start parallel
         * enumerations.  Descriptor/control timeouts are not retried here; they
         * are usually dead/firmware-owned ports and retrying them stalls boot.
         */
        if (last != STATUS_DEVICE_ERROR || attempt == maxAttempts)
            break;

        DbgPrint("usb-hub: root port %u enumeration attempt %u failed st=0x%x; retrying\n",
                 port, attempt, last);
        UsbDelayUs(xhci->Hal, 50000);
    }

    if (UsbHubRootPortConnected(xhci, port)) {
        UsbHubMarkRootPortFailed(hub, port);
    } else {
        UsbHubClearRootPortFailure(hub, port);
        hub->RootPortState[port] = UsbHubPortDisconnected;
    }

    return last;
}

static bool UsbHubPopPendingRootPort(USB_HUB_EXTENSION* hub, uint8_t* outPort) {
    if (!hub || !outPort) return false;

    uintptr_t flags = KeAcquireSpinlock(&hub->PendingLock);
    for (uint32_t word = 0; word < 8; word++) {
        uint32_t bits = hub->PendingRootPorts[word];
        if (!bits)
            continue;
        for (uint8_t bit = 0; bit < 32; bit++) {
            if (!(bits & (1u << bit)))
                continue;
            hub->PendingRootPorts[word] &= ~(1u << bit);
            KeReleaseSpinlock(&hub->PendingLock, flags);
            *outPort = (uint8_t)(word * 32u + bit);
            return true;
        }
    }
    KeReleaseSpinlock(&hub->PendingLock, flags);
    return false;
}

static bool UsbHubPopPendingExternalPort(USB_HUB_EXTENSION* hub, uint8_t* outPort) {
    if (!hub || !outPort) return false;

    uintptr_t flags = KeAcquireSpinlock(&hub->PendingLock);
    uint32_t bits = hub->PendingHubPorts[0];
    if (!bits) {
        KeReleaseSpinlock(&hub->PendingLock, flags);
        return false;
    }

    for (uint8_t port = 1; port <= hub->PortCount && port <= USB_EXTERNAL_HUB_MAX_PORTS; port++) {
        if (!(bits & (1u << port)))
            continue;
        hub->PendingHubPorts[0] &= ~(1u << port);
        KeReleaseSpinlock(&hub->PendingLock, flags);
        *outPort = port;
        return true;
    }

    hub->PendingHubPorts[0] = 0;
    KeReleaseSpinlock(&hub->PendingLock, flags);
    return false;
}

static void UsbHubRecordRootPortChange(USB_HUB_EXTENSION* hub, uint8_t port) {
    if (!hub || port == 0 || port > hub->PortCount)
        return;
    uint32_t word = (uint32_t)port / 32u;
    uint32_t bit = (uint32_t)port & 31u;
    uintptr_t flags = KeAcquireSpinlock(&hub->PendingLock);
    hub->PendingRootPorts[word] |= (1u << bit);
    KeReleaseSpinlock(&hub->PendingLock, flags);
    if (hub->WorkerSemaphore)
        KeSignalSemaphore(hub->WorkerSemaphore);
}

void UsbHubQueueExternalPortChange(USB_HUB_EXTENSION* hub, uint8_t port) {
    if (!hub || port == 0 || port > hub->PortCount || port > USB_EXTERNAL_HUB_MAX_PORTS)
        return;
    uintptr_t flags = KeAcquireSpinlock(&hub->PendingLock);
    hub->PendingHubPorts[0] |= (1u << port);
    KeReleaseSpinlock(&hub->PendingLock, flags);
    if (hub->WorkerSemaphore)
        KeSignalSemaphore(hub->WorkerSemaphore);
}

static void UsbHubRemoveSlotChild(USB_HUB_EXTENSION* hub, uint8_t port, uint8_t slot) {
    if (!hub || !hub->Hcd || slot == 0)
        return;

    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)hub->Hcd->ControllerContext;
    if (!xhci || slot > XHCI_MAX_TRACKED_SLOTS)
        return;

    XHCI_DEVICE* xdev = &xhci->Devices[slot];
    DbgPrint("usb-hub: %s port %u removing slot=%u\n",
             hub->Kind == UsbHubKindRoot ? "root" : "hub", port, slot);

    if (xdev->UsbPdo)
        (void)UsbRemovePdoSubtree(xdev->UsbPdo, true);

    if (xdev->Allocated)
        (void)XhciDisableSlot(xhci, slot);

    if (xdev->UsbDevice) {
        xdev->UsbDevice->State = UsbDeviceRemoved;
        UsbDereferenceDevice(xdev->UsbDevice);
    }

    memset(xdev, 0, sizeof(*xdev));
    if (hub->Kind == UsbHubKindRoot && port <= USB_ROOT_HUB_MAX_PORTS) {
        hub->RootPortSlot[port] = 0;
        hub->RootPortState[port] = UsbHubPortDisconnected;
    } else if (port <= USB_EXTERNAL_HUB_MAX_PORTS) {
        hub->ChildSlot[port] = 0;
        hub->ChildPortState[port] = UsbHubPortDisconnected;
    }
}

static void UsbHubRemoveRootChild(USB_HUB_EXTENSION* hub, uint8_t port, uint8_t slot) {
    UsbHubRemoveSlotChild(hub, port, slot);
}

static KSTATUS UsbHubControlIn(USB_HUB_EXTENSION* hub,
                               uint8_t requestType,
                               uint8_t request,
                               uint16_t value,
                               uint16_t index,
                               void* buffer,
                               PHYSICAL_ADDRESS bufferPhys,
                               uint16_t length) {
    if (!hub || !hub->LowerDevice || !buffer || !bufferPhys || !length)
        return STATUS_INVALID_PARAMETER;

    USB_URB urb;
    memset(&urb, 0, sizeof(urb));
    urb.Function = UrbFunctionControlTransfer;
    urb.Direction = UsbTransferIn;
    urb.Device = hub->Device;
    urb.Buffer = buffer;
    urb.BufferPhys = bufferPhys;
    urb.Length = length;
    urb.Setup.bmRequestType = requestType;
    urb.Setup.bRequest = request;
    urb.Setup.wValue = value;
    urb.Setup.wIndex = index;
    urb.Setup.wLength = length;
    return UsbSubmitUrbSync(hub->LowerDevice, &urb);
}

static KSTATUS UsbHubControlNoData(USB_HUB_EXTENSION* hub,
                                   uint8_t requestType,
                                   uint8_t request,
                                   uint16_t value,
                                   uint16_t index) {
    if (!hub || !hub->LowerDevice)
        return STATUS_INVALID_PARAMETER;

    USB_URB urb;
    memset(&urb, 0, sizeof(urb));
    urb.Function = UrbFunctionControlTransfer;
    urb.Direction = UsbTransferNoData;
    urb.Device = hub->Device;
    urb.Setup.bmRequestType = requestType;
    urb.Setup.bRequest = request;
    urb.Setup.wValue = value;
    urb.Setup.wIndex = index;
    urb.Setup.wLength = 0;
    return UsbSubmitUrbSync(hub->LowerDevice, &urb);
}

static KSTATUS UsbHubReadDescriptor(USB_HUB_EXTENSION* hub, USB_HUB_DESCRIPTOR* outDesc) {
    if (!hub || !outDesc || !hub->ChangeBitmap || !hub->ChangeBitmapPhys)
        return STATUS_INVALID_PARAMETER;

    memset(hub->ChangeBitmap, 0, USB_HUB_CHANGE_BITMAP_BYTES);
    KSTATUS st = UsbHubControlIn(hub,
                                 USB_DIR_IN | USB_TYPE_CLASS | USB_RECIP_DEVICE,
                                 USB_REQ_GET_DESCRIPTOR,
                                 (uint16_t)(USB_DESC_HUB << 8),
                                 0,
                                 hub->ChangeBitmap,
                                 hub->ChangeBitmapPhys,
                                 USB_HUB_CHANGE_BITMAP_BYTES);
    if (!KSUCCESS(st))
        return st;

    memcpy(outDesc, hub->ChangeBitmap, sizeof(USB_HUB_DESCRIPTOR));
    if (outDesc->bDescriptorType != USB_DESC_HUB || outDesc->bNbrPorts == 0)
        return STATUS_DEVICE_ERROR;
    return STATUS_SUCCESS;
}

static KSTATUS UsbHubGetPortStatus(USB_HUB_EXTENSION* hub, uint8_t port, USB_HUB_PORT_STATUS* outStatus) {
    if (!hub || !outStatus || port == 0)
        return STATUS_INVALID_PARAMETER;

    memset(hub->ChangeBitmap, 0, USB_HUB_CHANGE_BITMAP_BYTES);
    KSTATUS st = UsbHubControlIn(hub,
                                 USB_DIR_IN | USB_TYPE_CLASS | USB_RECIP_OTHER,
                                 USB_REQ_GET_STATUS,
                                 0,
                                 port,
                                 hub->ChangeBitmap,
                                 hub->ChangeBitmapPhys,
                                 sizeof(USB_HUB_PORT_STATUS));
    if (!KSUCCESS(st))
        return st;

    memcpy(outStatus, hub->ChangeBitmap, sizeof(USB_HUB_PORT_STATUS));
    return STATUS_SUCCESS;
}

static KSTATUS UsbHubSetPortFeature(USB_HUB_EXTENSION* hub, uint8_t port, uint16_t feature) {
    return UsbHubControlNoData(hub,
                               USB_DIR_OUT | USB_TYPE_CLASS | USB_RECIP_OTHER,
                               USB_REQ_SET_FEATURE,
                               feature,
                               port);
}

static KSTATUS UsbHubClearPortFeature(USB_HUB_EXTENSION* hub, uint8_t port, uint16_t feature) {
    return UsbHubControlNoData(hub,
                               USB_DIR_OUT | USB_TYPE_CLASS | USB_RECIP_OTHER,
                               USB_REQ_CLEAR_FEATURE,
                               feature,
                               port);
}

static uint8_t UsbHubSpeedFromPortStatus(uint16_t status) {
    if (status & USB_HUB_PORT_STATUS_LOW_SPEED)
        return 2u;
    if (status & USB_HUB_PORT_STATUS_HIGH_SPEED)
        return 3u;
    return 1u;
}

static KSTATUS UsbHubResetExternalPort(USB_HUB_EXTENSION* hub, uint8_t port, uint8_t* outSpeed) {
    if (outSpeed) *outSpeed = 0;
    if (!hub || port == 0)
        return STATUS_INVALID_PARAMETER;

    (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_CONNECTION);
    (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_ENABLE);
    (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_RESET);

    KSTATUS st = UsbHubSetPortFeature(hub, port, USB_HUB_PORT_FEAT_RESET);
    if (!KSUCCESS(st))
        return st;

    USB_HUB_PORT_STATUS ps;
    memset(&ps, 0, sizeof(ps));
    for (uint32_t i = 0; i < 500u; i++) {
        UsbDelayUs(((XHCI_CONTROLLER*)hub->Hcd->ControllerContext)->Hal, 1000u);
        st = UsbHubGetPortStatus(hub, port, &ps);
        if (!KSUCCESS(st))
            return st;
        if (ps.Change & USB_HUB_PORT_CHANGE_RESET)
            break;
    }

    (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_RESET);
    (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_ENABLE);

    st = UsbHubGetPortStatus(hub, port, &ps);
    if (!KSUCCESS(st))
        return st;
    if (!(ps.Status & USB_HUB_PORT_STATUS_CONNECTION))
        return STATUS_NOT_FOUND;
    if (!(ps.Status & USB_HUB_PORT_STATUS_ENABLE))
        return STATUS_DEVICE_ERROR;

    if (outSpeed)
        *outSpeed = UsbHubSpeedFromPortStatus(ps.Status);
    return STATUS_SUCCESS;
}

static void UsbHubHandleExternalPortChange(USB_HUB_EXTENSION* hub, uint8_t port) {
    if (!hub || hub->Kind != UsbHubKindExternal || port == 0 || port > hub->PortCount || port > USB_EXTERNAL_HUB_MAX_PORTS)
        return;

    USB_HUB_PORT_STATUS ps;
    memset(&ps, 0, sizeof(ps));
    KSTATUS st = UsbHubGetPortStatus(hub, port, &ps);
    if (!KSUCCESS(st)) {
        DbgPrint("usb-hub: slot=%u port=%u GET_STATUS failed st=0x%x\n",
                 hub->Device ? hub->Device->SlotId : 0, port, st);
        return;
    }

    if (ps.Change & USB_HUB_PORT_CHANGE_CONNECTION)
        (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_CONNECTION);
    if (ps.Change & USB_HUB_PORT_CHANGE_ENABLE)
        (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_ENABLE);
    if (ps.Change & USB_HUB_PORT_CHANGE_RESET)
        (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_RESET);

    uint8_t oldSlot = hub->ChildSlot[port];
    if (!(ps.Status & USB_HUB_PORT_STATUS_CONNECTION)) {
        if (UsbHubSlotBusy(oldSlot)) {
            hub->ChildPortState[port] = UsbHubPortRemoving;
            UsbHubRemoveSlotChild(hub, port, oldSlot);
        } else {
            hub->ChildPortState[port] = UsbHubPortDisconnected;
            UsbTrace("usb-hub: slot=%u port=%u disconnected\n", hub->Device ? hub->Device->SlotId : 0, port);
        }
        return;
    }

    if (oldSlot == USB_HUB_PORT_ENUMERATING) {
        UsbTrace("usb-hub: slot=%u port=%u change ignored; enumeration already in progress\n",
                 hub->Device ? hub->Device->SlotId : 0, port);
        return;
    }
    if (UsbHubSlotBusy(oldSlot)) {
        UsbTrace("usb-hub: slot=%u port=%u change ignored; child slot=%u still present\n",
                 hub->Device ? hub->Device->SlotId : 0, port, oldSlot);
        return;
    }

    hub->ChildSlot[port] = USB_HUB_PORT_ENUMERATING;
    hub->ChildPortState[port] = UsbHubPortConnectedDebounce;
    UsbDelayUs(((XHCI_CONTROLLER*)hub->Hcd->ControllerContext)->Hal, 30000u);

    uint8_t speed = 0;
    hub->ChildPortState[port] = UsbHubPortResetting;
    st = UsbHubResetExternalPort(hub, port, &speed);
    if (!KSUCCESS(st)) {
        hub->ChildSlot[port] = 0;
        hub->ChildPortState[port] = UsbHubPortFailedConnected;
        DbgPrint("usb-hub: slot=%u port=%u reset failed st=0x%x\n",
                 hub->Device ? hub->Device->SlotId : 0, port, st);
        return;
    }

    hub->ChildPortState[port] = UsbHubPortEnumerating;
    uint8_t slot = 0;
    st = XhciEnumerateHubPortEx((XHCI_CONTROLLER*)hub->Hcd->ControllerContext, hub, port, speed, &slot);
    if (KSUCCESS(st)) {
        hub->ChildSlot[port] = slot;
        hub->ChildPortState[port] = UsbHubPortChildPresent;
        DbgPrint("usb-hub: slot=%u port=%u attached child-slot=%u speed=%s(%u)\n",
                 hub->Device ? hub->Device->SlotId : 0,
                 port,
                 slot,
                 XhciPortSpeedName(speed),
                 speed);
    } else {
        hub->ChildSlot[port] = 0;
        hub->ChildPortState[port] = UsbHubPortFailedConnected;
        DbgPrint("usb-hub: slot=%u port=%u enumeration failed st=0x%x\n",
                 hub->Device ? hub->Device->SlotId : 0, port, st);
    }
}

static void UsbHubHandleRootPortChange(USB_HUB_EXTENSION* hub, uint8_t port) {
    if (!hub || !hub->Hcd || port == 0 || port > hub->PortCount || port > USB_ROOT_HUB_MAX_PORTS)
        return;

    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)hub->Hcd->ControllerContext;
    if (!xhci || !xhci->Op)
        return;

    uint32_t portsc = XhciRead32(xhci->Op, XhciPortReg(port));
    uint32_t changes = portsc & XHCI_PORTSC_CHANGE_MASK;
    bool connected = (portsc & XHCI_PORTSC_CCS) != 0;
    uint8_t oldSlot = hub->RootPortSlot[port];

    if (changes)
        XhciClearPortChanges(xhci, port);

    if (!connected) {
        if (UsbHubSlotBusy(oldSlot)) {
            hub->RootPortState[port] = UsbHubPortRemoving;
            UsbHubRemoveRootChild(hub, port, oldSlot);
        } else if (oldSlot == USB_HUB_PORT_ENUMERATING)
            UsbTrace("usb-hub: root port %u disconnect while enumeration is already in progress\n", port);
        else {
            hub->RootPortState[port] = UsbHubPortDisconnected;
            DbgPrint("usb-hub: root port %u disconnected\n", port);
        }
        UsbHubClearRootPortFailure(hub, port);
        return;
    }

    if (oldSlot == USB_HUB_PORT_ENUMERATING) {
        UsbTrace("usb-hub: root port %u change ignored; enumeration already in progress\n", port);
        return;
    }

    if (oldSlot == USB_HUB_PORT_FAILED) {
        if (changes & XHCI_PORTSC_CSC) {
            UsbTrace("usb-hub: root port %u retrying after connection-change on failed port\n", port);
            UsbHubClearRootPortFailure(hub, port);
            oldSlot = 0;
        } else {
            UsbTrace("usb-hub: root port %u change ignored; previous enumeration failed and device is still connected\n", port);
            return;
        }
    }

    if (UsbHubSlotBusy(oldSlot)) {
        UsbTrace("usb-hub: root port %u change ignored; slot=%u still present\n", port, oldSlot);
        return;
    }

    hub->RootPortSlot[port] = USB_HUB_PORT_ENUMERATING;
    hub->RootPortState[port] = UsbHubPortConnectedDebounce;

    UsbDelayUs(xhci->Hal, 30000);
    portsc = XhciRead32(xhci->Op, XhciPortReg(port));
    if (!(portsc & XHCI_PORTSC_CCS)) {
        UsbHubClearRootPortFailure(hub, port);
        return;
    }

    hub->RootPortState[port] = UsbHubPortEnumerating;
    uint8_t slot = 0;
    KSTATUS st = UsbHubEnumerateRootPortWithRetries(hub, xhci, port, &slot);
    if (KSUCCESS(st)) {
        hub->RootPortSlot[port] = slot;
        hub->RootPortState[port] = UsbHubPortChildPresent;
        DbgPrint("usb-hub: root port %u attached slot=%u\n", port, slot);
    } else {
        DbgPrint("usb-hub: root port %u hotplug enumeration failed st=0x%x after %u attempt(s)\n",
                 port, st, hub->RootPortFailCount[port]);
    }
}
static void UsbHubArmExternalInterrupt(USB_HUB_EXTENSION* hub);
static KSTATUS UsbHubDoRootPortScan(USB_HCD* hcd);

static bool UsbHubTakeInitialRootScan(USB_HUB_EXTENSION* hub) {
    if (!hub || hub->Kind != UsbHubKindRoot)
        return false;
    uintptr_t flags = KeAcquireSpinlock(&hub->PendingLock);
    bool take = hub->InitialScanPending && !hub->InitialScanInProgress && !hub->InitialScanComplete;
    if (take) {
        hub->InitialScanPending = false;
        hub->InitialScanInProgress = true;
    }
    KeReleaseSpinlock(&hub->PendingLock, flags);
    return take;
}

static void UsbHubWorker(void* arg) {
    USB_HUB_EXTENSION* hub = (USB_HUB_EXTENSION*)arg;
    while (hub && hub->WorkerRunning) {
        KSTATUS st = KeWaitSemaphore(hub->WorkerSemaphore);
        if (!KSUCCESS(st))
            continue;

        if (hub->Kind == UsbHubKindRoot) {
            if (UsbHubTakeInitialRootScan(hub))
                (void)UsbHubDoRootPortScan(hub->Hcd);

            uint8_t port = 0;
            while (UsbHubPopPendingRootPort(hub, &port))
                UsbHubHandleRootPortChange(hub, port);
        } else {
            uint8_t port = 0;
            while (UsbHubPopPendingExternalPort(hub, &port))
                UsbHubHandleExternalPortChange(hub, port);
            if (hub->Started && !hub->Removed && !hub->SurpriseRemoved && !hub->ChangeUrbPending)
                UsbHubArmExternalInterrupt(hub);
        }
    }
}

static KSTATUS UsbHubDoRootPortScan(USB_HCD* hcd) {
    if (!hcd || !hcd->RootHubContext)
        return STATUS_INVALID_PARAMETER;

    USB_HUB_EXTENSION* hub = (USB_HUB_EXTENSION*)hcd->RootHubContext;
    if (!hub->Started)
        return STATUS_OBJECT_DEAD;

    uint32_t attempted = 0;
    uint32_t configured = 0;
    XHCI_CONTROLLER* xhci = (XHCI_CONTROLLER*)hcd->ControllerContext;
    if (!xhci || !xhci->Op)
        return STATUS_INVALID_PARAMETER;

    for (uint8_t port = 1; port <= hub->PortCount; port++) {
        uint32_t portsc = XhciRead32(xhci->Op, XhciPortReg(port));
        if (!(portsc & XHCI_PORTSC_CCS))
            continue;

        if (hub->RootPortSlot[port]) {
            UsbTrace("usb-hub: root port %u initial scan skipped; state=%u\n",
                     port, hub->RootPortSlot[port]);
            continue;
        }

        attempted++;
        hub->RootPortSlot[port] = USB_HUB_PORT_ENUMERATING;
        hub->RootPortState[port] = UsbHubPortEnumerating;
        uint8_t slot = 0;
        KSTATUS st = UsbHubEnumerateRootPortWithRetries(hub, xhci, port, &slot);
        if (KSUCCESS(st)) {
            hub->RootPortSlot[port] = slot;
            hub->RootPortState[port] = UsbHubPortChildPresent;
            configured++;
        } else {
            DbgPrint("usb-hub: root port %u enumeration failed st=0x%x after %u attempt(s)\n",
                     port, st, hub->RootPortFailCount[port]);
        }
    }

    DbgPrint("usb-hub: root-port scan complete; attempted=%u configured=%u\n", attempted, configured);
    hub->InitialScanComplete = true;
    hub->InitialScanInProgress = false;
    return STATUS_SUCCESS;
}

KSTATUS UsbHubScheduleRootScan(USB_HCD* hcd) {
    if (!hcd || !hcd->RootHubContext)
        return STATUS_INVALID_PARAMETER;
    USB_HUB_EXTENSION* hub = (USB_HUB_EXTENSION*)hcd->RootHubContext;
    if (hub->Kind != UsbHubKindRoot || !hub->Started)
        return STATUS_OBJECT_DEAD;

    uintptr_t flags = KeAcquireSpinlock(&hub->PendingLock);
    if (!hub->InitialScanComplete && !hub->InitialScanInProgress)
        hub->InitialScanPending = true;
    KeReleaseSpinlock(&hub->PendingLock, flags);

    if (hub->WorkerSemaphore)
        KeSignalSemaphore(hub->WorkerSemaphore);
    return STATUS_SUCCESS;
}

void UsbHubQueueRootPortChange(USB_HCD* hcd, uint8_t port) {
    if (!hcd || !hcd->RootHubContext)
        return;
    USB_HUB_EXTENSION* hub = (USB_HUB_EXTENSION*)hcd->RootHubContext;
    UsbHubRecordRootPortChange(hub, port);
}

static KSTATUS UsbHubScanExternalPorts(USB_HUB_EXTENSION* hub) {
    if (!hub || hub->Kind != UsbHubKindExternal)
        return STATUS_INVALID_PARAMETER;

    uint32_t attempted = 0;
    uint32_t configured = 0;
    for (uint8_t port = 1; port <= hub->PortCount && port <= USB_EXTERNAL_HUB_MAX_PORTS; port++) {
        USB_HUB_PORT_STATUS ps;
        memset(&ps, 0, sizeof(ps));
        KSTATUS st = UsbHubGetPortStatus(hub, port, &ps);
        if (!KSUCCESS(st))
            continue;
        (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_CONNECTION);
        (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_ENABLE);
        (void)UsbHubClearPortFeature(hub, port, USB_HUB_PORT_FEAT_C_RESET);
        if (!(ps.Status & USB_HUB_PORT_STATUS_CONNECTION))
            continue;
        attempted++;
        UsbHubHandleExternalPortChange(hub, port);
        if (UsbHubSlotBusy(hub->ChildSlot[port]))
            configured++;
    }
    DbgPrint("usb-hub: external hub slot=%u scan complete; attempted=%u configured=%u\n",
             hub->Device ? hub->Device->SlotId : 0, attempted, configured);
    return STATUS_SUCCESS;
}

static const USB_ENDPOINT_DESCRIPTOR* UsbHubFindInterruptIn(USB_PDO_EXTENSION* pdoExt) {
    if (!pdoExt)
        return (const USB_ENDPOINT_DESCRIPTOR*)0;

    if (pdoExt->Kind == UsbPdoInterface) {
        for (uint8_t i = 0; i < pdoExt->EndpointCount; i++) {
            const USB_ENDPOINT_DESCRIPTOR* ep = &pdoExt->Endpoints[i];
            if ((ep->bEndpointAddress & USB_ENDPOINT_DIR_IN) &&
                ((ep->bmAttributes & USB_ENDPOINT_ATTR_TYPE_MASK) == USB_ENDPOINT_TYPE_INTERRUPT))
                return ep;
        }
        return (const USB_ENDPOINT_DESCRIPTOR*)0;
    }

    if (pdoExt->Kind == UsbPdoDevice && pdoExt->Device) {
        XHCI_DEVICE* xdev = (XHCI_DEVICE*)pdoExt->Device->HcdDeviceContext;
        if (!xdev)
            return (const USB_ENDPOINT_DESCRIPTOR*)0;
        for (uint8_t i = 0; i < xdev->InterfaceCount; i++) {
            USB_TRACKED_INTERFACE* iface = &xdev->Interfaces[i];
            if (!iface->Present || iface->Descriptor.bInterfaceClass != USB_CLASS_HUB)
                continue;
            for (uint8_t epn = 0; epn < iface->EndpointCount; epn++) {
                const USB_ENDPOINT_DESCRIPTOR* ep = &iface->Endpoints[epn];
                if ((ep->bEndpointAddress & USB_ENDPOINT_DIR_IN) &&
                    ((ep->bmAttributes & USB_ENDPOINT_ATTR_TYPE_MASK) == USB_ENDPOINT_TYPE_INTERRUPT))
                    return ep;
            }
        }
    }
    return (const USB_ENDPOINT_DESCRIPTOR*)0;
}

static void UsbHubExternalChangeComplete(USB_URB* urb, void* context) {
    USB_HUB_EXTENSION* hub = (USB_HUB_EXTENSION*)context;
    if (!hub || urb != &hub->ChangeUrb)
        return;

    hub->ChangeUrbPending = false;
    if (!KSUCCESS(urb->Status) || hub->Removed || hub->SurpriseRemoved)
        return;

    uint32_t actual = urb->ActualLength;
    if (actual > USB_HUB_CHANGE_BITMAP_BYTES)
        actual = USB_HUB_CHANGE_BITMAP_BYTES;

    for (uint8_t port = 1; port <= hub->PortCount && port <= USB_EXTERNAL_HUB_MAX_PORTS; port++) {
        uint8_t byteIndex = port / 8u;
        uint8_t bit = port & 7u;
        if (byteIndex < actual && (hub->ChangeBitmap[byteIndex] & (1u << bit)))
            UsbHubQueueExternalPortChange(hub, port);
    }
}

static void UsbHubArmExternalInterrupt(USB_HUB_EXTENSION* hub) {
    if (!hub || !hub->InterruptPipe || !hub->ChangeBitmap || hub->ChangeUrbPending)
        return;

    memset(hub->ChangeBitmap, 0, USB_HUB_CHANGE_BITMAP_BYTES);
    memset(&hub->ChangeUrb, 0, sizeof(hub->ChangeUrb));
    hub->ChangeUrb.Function = UrbFunctionInterruptTransfer;
    hub->ChangeUrb.Direction = UsbTransferIn;
    hub->ChangeUrb.Device = hub->Device;
    hub->ChangeUrb.Pipe = hub->InterruptPipe;
    hub->ChangeUrb.Buffer = hub->ChangeBitmap;
    hub->ChangeUrb.BufferPhys = hub->ChangeBitmapPhys;
    hub->ChangeUrb.Length = hub->InterruptMaxPacket;
    if (hub->ChangeUrb.Length > USB_HUB_CHANGE_BITMAP_BYTES)
        hub->ChangeUrb.Length = USB_HUB_CHANGE_BITMAP_BYTES;
    hub->ChangeUrb.Complete = UsbHubExternalChangeComplete;
    hub->ChangeUrb.CompleteContext = hub;

    KSTATUS st = UsbSubmitUrbAsync(hub->LowerDevice, &hub->ChangeUrb);
    if (st == STATUS_PENDING) {
        hub->ChangeUrbPending = true;
    } else if (!KSUCCESS(st)) {
        DbgPrint("usb-hub: slot=%u interrupt arm failed st=0x%x\n",
                 hub->Device ? hub->Device->SlotId : 0, st);
    }
}

static KSTATUS UsbHubStartRootDevice(DEVICE_OBJECT* device, USB_HUB_EXTENSION* hub) {
    hub->PortCount = hub->Hcd ? hub->Hcd->MaxRootPorts : 0;
    if (hub->PortCount > USB_ROOT_HUB_MAX_PORTS)
        hub->PortCount = USB_ROOT_HUB_MAX_PORTS;

    KSTATUS st = KeSemaphoreCreate(0, 0x7fffffff, &hub->WorkerSemaphore);
    if (!KSUCCESS(st))
        return st;

    hub->WorkerRunning = true;
    st = KeCreateKernelThread((uintptr_t)UsbHubWorker,
                              (uintptr_t)hub,
                              PriorityNormal,
                              &hub->WorkerThread);
    if (!KSUCCESS(st)) {
        hub->WorkerRunning = false;
        return st;
    }

    hub->Started = true;
    device->PnpState = DevicePnpStarted;
    if (hub->HubPdo)
        hub->HubPdo->PnpState = DevicePnpStarted;

    DbgPrint("usb-hub: root hub started ports=%u\n", hub->PortCount);
    return STATUS_SUCCESS;
}

static KSTATUS UsbHubStartExternalDevice(DEVICE_OBJECT* device, USB_HUB_EXTENSION* hub) {
    if (!hub || !hub->Device)
        return STATUS_INVALID_PARAMETER;

    hub->ChangeBitmap = (uint8_t*)UsbAllocateDmaBuffer(hub->Device, 1, &hub->ChangeBitmapPhys);
    if (!hub->ChangeBitmap)
        return STATUS_NO_MEMORY;

    USB_HUB_DESCRIPTOR desc;
    memset(&desc, 0, sizeof(desc));
    KSTATUS st = UsbHubReadDescriptor(hub, &desc);
    if (!KSUCCESS(st)) {
        DbgPrint("usb-hub: slot=%u descriptor read failed st=0x%x\n", hub->Device->SlotId, st);
        return st;
    }

    hub->PortCount = desc.bNbrPorts;
    if (hub->PortCount > USB_EXTERNAL_HUB_MAX_PORTS) {
        DbgPrint("usb-hub: slot=%u reports %u ports; limiting to %u\n",
                 hub->Device->SlotId, hub->PortCount, USB_EXTERNAL_HUB_MAX_PORTS);
        hub->PortCount = USB_EXTERNAL_HUB_MAX_PORTS;
    }
    hub->PowerOnToPowerGoodMs = desc.bPwrOn2PwrGood ? (uint8_t)(desc.bPwrOn2PwrGood * 2u) : 100u;

    for (uint8_t port = 1; port <= hub->PortCount; port++)
        (void)UsbHubSetPortFeature(hub, port, USB_HUB_PORT_FEAT_POWER);
    UsbDelayUs(((XHCI_CONTROLLER*)hub->Hcd->ControllerContext)->Hal,
               (uint32_t)hub->PowerOnToPowerGoodMs * 1000u);

    st = KeSemaphoreCreate(0, 0x7fffffff, &hub->WorkerSemaphore);
    if (!KSUCCESS(st))
        return st;

    hub->WorkerRunning = true;
    st = KeCreateKernelThread((uintptr_t)UsbHubWorker,
                              (uintptr_t)hub,
                              PriorityNormal,
                              &hub->WorkerThread);
    if (!KSUCCESS(st)) {
        hub->WorkerRunning = false;
        return st;
    }

    st = UsbCreatePipe(hub->Device,
                       UsbPipeInterrupt,
                       hub->InterruptEndpointAddress,
                       hub->InterruptMaxPacket,
                       hub->InterruptInterval,
                       &hub->InterruptPipe);
    if (!KSUCCESS(st))
        return st;

    hub->Started = true;
    device->PnpState = DevicePnpStarted;
    if (hub->HubPdo)
        hub->HubPdo->PnpState = DevicePnpStarted;

    DbgPrint("usb-hub: external hub started slot=%u ports=%u ep=%02x maxpkt=%u interval=%u\n",
             hub->Device->SlotId,
             hub->PortCount,
             hub->InterruptEndpointAddress,
             hub->InterruptMaxPacket,
             hub->InterruptInterval);

    (void)UsbHubScanExternalPorts(hub);
    UsbHubArmExternalInterrupt(hub);
    return STATUS_SUCCESS;
}

static KSTATUS UsbHubStartDevice(DEVICE_OBJECT* device, USB_HUB_EXTENSION* hub) {
    if (!device || !hub || hub->Started)
        return STATUS_SUCCESS;

    if (hub->Kind == UsbHubKindRoot)
        return UsbHubStartRootDevice(device, hub);
    return UsbHubStartExternalDevice(device, hub);
}

KSTATUS UsbHubDispatchPnp(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_HUB_EXTENSION* hub = (USB_HUB_EXTENSION*)device->DeviceExtension;
    KSTATUS st = STATUS_SUCCESS;

    switch ((IO_PNP_MINOR_FUNCTION)ios->MinorFunction) {
        case IRP_MN_START_DEVICE:
            st = UsbHubStartDevice(device, hub);
            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 (!hub || !hub->SurpriseRemoved)
                device->PnpState = DevicePnpStarted;
            return IoForwardIrpToLower(device, irp);

        case IRP_MN_SURPRISE_REMOVAL:
            device->PnpState = DevicePnpSurpriseRemoved;
            if (hub) {
                hub->SurpriseRemoved = true;
                hub->WorkerRunning = false;
                if (hub->ChangeUrbPending)
                    (void)UsbCancelUrb(&hub->ChangeUrb);
                if (hub->WorkerSemaphore)
                    KeSignalSemaphore(hub->WorkerSemaphore);

                // Remove downstream PDOs owned by this hub FDO.  The explicit
                // slot arrays are still used as the hardware lookup table; the
                // child PDO tree is the authoritative software lifetime tree.
                (void)UsbRemoveDeviceObjectChildren(device, true);
                if (hub->Kind == UsbHubKindExternal) {
                    for (uint8_t port = 1; port <= hub->PortCount && port <= USB_EXTERNAL_HUB_MAX_PORTS; port++) {
                        if (UsbHubSlotBusy(hub->ChildSlot[port]))
                            UsbHubRemoveSlotChild(hub, port, hub->ChildSlot[port]);
                    }
                } else if (hub->Kind == UsbHubKindRoot) {
                    for (uint8_t port = 1; port <= hub->PortCount && port <= USB_ROOT_HUB_MAX_PORTS; port++) {
                        if (UsbHubSlotBusy(hub->RootPortSlot[port]))
                            UsbHubRemoveRootChild(hub, port, hub->RootPortSlot[port]);
                    }
                }
            }
            return IoForwardIrpToLower(device, irp);

        case IRP_MN_REMOVE_DEVICE: {
            device->PnpState = DevicePnpDeleted;
            if (hub) {
                hub->Removed = true;
                hub->WorkerRunning = false;
                if (hub->ChangeUrbPending)
                    (void)UsbCancelUrb(&hub->ChangeUrb);
                (void)UsbRemoveDeviceObjectChildren(device, false);
                if (hub->InterruptPipe) {
                    hub->InterruptPipe->Open = false;
                    UsbDereferencePipe(hub->InterruptPipe);
                    hub->InterruptPipe = (USB_PIPE*)0;
                }
                if (hub->WorkerSemaphore)
                    KeSignalSemaphore(hub->WorkerSemaphore);
                if (hub->Kind == UsbHubKindRoot && hub->Hcd && hub->Hcd->RootHubContext == hub)
                    hub->Hcd->RootHubContext = (void*)0;
                if (hub->RemoveLockReady) {
                    IoReleaseRemoveLockAndWait(&hub->RemoveLock, irp);
                    IoDestroyRemoveLock(&hub->RemoveLock);
                    hub->RemoveLockReady = false;
                }
            }

            KSTATUS lower = hub && hub->LowerDevice ? IoForwardIrpToLower(device, irp) : STATUS_SUCCESS;
            if (hub && hub->LowerDevice) {
                IoDetachDevice(device);
                hub->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 UsbHubAddDevice(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)
        return STATUS_INVALID_PARAMETER;

    USB_HUB_KIND kind;
    USB_DEVICE* usbDevice = (USB_DEVICE*)0;
    const USB_ENDPOINT_DESCRIPTOR* intr = (const USB_ENDPOINT_DESCRIPTOR*)0;

    if (pdoExt->Kind == UsbPdoRootHub && pdoExt->Hcd) {
        kind = UsbHubKindRoot;
    } else if (pdoExt->Kind == UsbPdoDevice && pdoExt->Device &&
               pdoExt->Device->DeviceClass == USB_CLASS_HUB) {
        kind = UsbHubKindExternal;
        usbDevice = pdoExt->Device;
        intr = UsbHubFindInterruptIn(pdoExt);
        if (!intr)
            return STATUS_NOT_FOUND;
    } else if (pdoExt->Kind == UsbPdoInterface && pdoExt->Device &&
               pdoExt->InterfaceDescriptor.bInterfaceClass == USB_CLASS_HUB) {
        kind = UsbHubKindExternal;
        usbDevice = pdoExt->Device;
        intr = UsbHubFindInterruptIn(pdoExt);
        if (!intr)
            return STATUS_NOT_FOUND;
    } else {
        return STATUS_INVALID_PARAMETER;
    }

    DEVICE_OBJECT* fdo = (DEVICE_OBJECT*)0;
    KSTATUS st = IoCreateDevice(driver, sizeof(USB_HUB_EXTENSION), &fdo);
    if (!KSUCCESS(st))
        return st;
    (void)IoSetDeviceName(fdo, kind == UsbHubKindRoot ? SYS_SV_LITERAL("usb-root-hub-driver")
                                                      : SYS_SV_LITERAL("usb-hub"));

    USB_HUB_EXTENSION* hub = (USB_HUB_EXTENSION*)fdo->DeviceExtension;
    hub->Self = fdo;
    hub->HubPdo = physicalDeviceObject;
    hub->Kind = kind;
    hub->Hcd = pdoExt->Hcd;
    hub->Device = usbDevice;
    if (kind == UsbHubKindExternal && hub->Device)
        hub->Hcd = hub->Device->Hcd;
    hub->PendingLock = (SPINLOCK)SPINLOCK_INIT;

    if (intr) {
        hub->InterruptEndpointAddress = intr->bEndpointAddress;
        hub->InterruptMaxPacket = intr->wMaxPacketSize;
        if (hub->InterruptMaxPacket == 0 || hub->InterruptMaxPacket > USB_HUB_CHANGE_BITMAP_BYTES)
            hub->InterruptMaxPacket = USB_HUB_CHANGE_BITMAP_BYTES;
        hub->InterruptInterval = intr->bInterval ? intr->bInterval : 12u;
        if (pdoExt->Kind == UsbPdoInterface)
            hub->InterfaceNumber = pdoExt->InterfaceDescriptor.bInterfaceNumber;
    }

    st = IoInitializeRemoveLock(&hub->RemoveLock);
    if (!KSUCCESS(st)) {
        IoDeleteDevice(fdo);
        return st;
    }
    hub->RemoveLockReady = true;

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

    if (kind == UsbHubKindRoot) {
        hub->Hcd->RootHubContext = hub;
        DbgPrint("usb-hub: root hub FDO attached\n");
    } else {
        DbgPrint("usb-hub: external hub FDO attached slot=%u ep=%02x\n",
                 hub->Device ? hub->Device->SlotId : 0,
                 hub->InterruptEndpointAddress);
    }
    return STATUS_SUCCESS;
}

KSTATUS UsbInitializeHubDriver(void) {
    memset(&gUsbHubDriver, 0, sizeof(gUsbHubDriver));
    IoInitializeDriverObject(&gUsbHubDriver, SYS_SV_LITERAL("usb-hub"));
    gUsbHubDriver.IdTable = &gUsbHubIdTable;
    gUsbHubDriver.DriverAddDevice = UsbHubAddDevice;
    gUsbHubDriver.MajorFunction[IRP_MJ_PNP] = UsbHubDispatchPnp;
    return IoRegisterDriverObject(&gUsbHubDriver);
}