// usb_core.c - USB descriptor/configuration parsing, interface publishing, and root-port enumeration.
#include "usb.h"

KSTATUS XhciSetConfiguration(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                    uint8_t configurationValue) {
    USB_SETUP_PACKET setup;
    setup.bmRequestType = USB_DIR_OUT;
    setup.bRequest = USB_REQ_SET_CONFIGURATION;
    setup.wValue = configurationValue;
    setup.wIndex = 0;
    setup.wLength = 0;

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

    dev->ConfigurationValue = configurationValue;
    dev->Configured = true;
    UsbTrace("usb: SET_CONFIGURATION slot=%u value=%u complete\n",
             dev->SlotId, configurationValue);
    return STATUS_SUCCESS;
}


static USB_TRACKED_INTERFACE* UsbFindTrackedInterfaceByNumber(XHCI_DEVICE* dev, uint8_t interfaceNumber) {
    if (!dev) return (USB_TRACKED_INTERFACE*)0;
    for (uint8_t i = 0; i < dev->InterfaceCount && i < USB_MAX_TRACKED_INTERFACES; i++) {
        if (dev->Interfaces[i].Present &&
            dev->Interfaces[i].Descriptor.bInterfaceNumber == interfaceNumber)
            return &dev->Interfaces[i];
    }
    return (USB_TRACKED_INTERFACE*)0;
}

void XhciLogConfigurationTree(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                     const USB_CONFIGURATION_DESCRIPTOR* config,
                                     uint16_t totalLength) {
    (void)xhci;

    memset(dev->Interfaces, 0, sizeof(dev->Interfaces));
    dev->InterfaceCount = 0;

    UsbTrace("usb: config slot=%u value=%u total=%u interfaces=%u attrs=%02x maxpower=%umA\n",
             dev->SlotId,
             config->bConfigurationValue,
             totalLength,
             config->bNumInterfaces,
             config->bmAttributes,
             (uint32_t)config->bMaxPower * 2u);

    uint16_t off = config->bLength;
    uint8_t currentInterface = 0xFFu;
    uint8_t currentClass = 0;
    uint8_t currentSubClass = 0;
    uint8_t currentProtocol = 0;
    bool hidCandidate = false;
    bool massCandidate = false;
    uint8_t hidInterruptIn = 0;
    uint8_t massBulkIn = 0;
    uint8_t massBulkOut = 0;
    USB_TRACKED_INTERFACE* tracked = (USB_TRACKED_INTERFACE*)0;

    while (off + 2u <= totalLength) {
        const uint8_t* desc = dev->DescriptorBuffer + off;
        uint8_t len = desc[0];
        uint8_t type = desc[1];

        if (len < 2 || off + len > totalLength) {
            DbgPrint("usb: descriptor parse stopped slot=%u off=%u len=%u type=%u total=%u\n",
                     dev->SlotId, off, len, type, totalLength);
            break;
        }

        if (type == USB_DESC_INTERFACE && len >= sizeof(USB_INTERFACE_DESCRIPTOR)) {
            const USB_INTERFACE_DESCRIPTOR* iface = (const USB_INTERFACE_DESCRIPTOR*)desc;
            currentInterface = iface->bInterfaceNumber;
            currentClass = iface->bInterfaceClass;
            currentSubClass = iface->bInterfaceSubClass;
            currentProtocol = iface->bInterfaceProtocol;
            hidCandidate = (currentClass == USB_CLASS_HID);
            massCandidate = (currentClass == USB_CLASS_MASS_STORAGE &&
                             currentSubClass == USB_MASS_SUBCLASS_SCSI &&
                             currentProtocol == USB_MASS_PROTOCOL_BOT);
            hidInterruptIn = 0;
            massBulkIn = 0;
            massBulkOut = 0;
            tracked = (USB_TRACKED_INTERFACE*)0;

            if (iface->bAlternateSetting == 0) {
                USB_TRACKED_INTERFACE* existing = UsbFindTrackedInterfaceByNumber(dev, iface->bInterfaceNumber);
                if (existing) {
                    tracked = existing;
                    memset(tracked->Endpoints, 0, sizeof(tracked->Endpoints));
                    tracked->EndpointCount = 0;
                    memcpy(&tracked->Descriptor, iface, sizeof(*iface));
                    UsbTrace("usb: slot=%u duplicate alt0 interface %u replaced existing PDO candidate\n",
                             dev->SlotId, iface->bInterfaceNumber);
                } else if (dev->InterfaceCount < USB_MAX_TRACKED_INTERFACES) {
                    tracked = &dev->Interfaces[dev->InterfaceCount++];
                    memset(tracked, 0, sizeof(*tracked));
                    tracked->Present = true;
                    memcpy(&tracked->Descriptor, iface, sizeof(*iface));
                } else {
                    DbgPrint("usb: slot=%u too many interfaces, interface %u alt=0 not published\n",
                             dev->SlotId, iface->bInterfaceNumber);
                }
            } else {
                /*
                 * Alternate settings are not independent USB interfaces.  Real
                 * UVC cameras often expose one VideoStreaming interface with
                 * many alternate settings, and publishing each alternate setting
                 * as a separate PDO causes duplicate interface devices and can
                 * exhaust USB_MAX_TRACKED_INTERFACES.  Track only alternate
                 * setting zero until a function driver explicitly selects a
                 * different setting.
                 */
                tracked = (USB_TRACKED_INTERFACE*)0;
            }

            UsbTrace("usb: interface slot=%u if=%u alt=%u eps=%u class=%02x subclass=%02x proto=%02x (%s)%s\n",
                     dev->SlotId,
                     iface->bInterfaceNumber,
                     iface->bAlternateSetting,
                     iface->bNumEndpoints,
                     currentClass,
                     currentSubClass,
                     currentProtocol,
                     UsbInterfaceClassName(currentClass, currentSubClass, currentProtocol),
                     iface->bAlternateSetting ? " alt-not-published" : "");
        } else if (type == USB_DESC_ENDPOINT && len >= sizeof(USB_ENDPOINT_DESCRIPTOR)) {
            const USB_ENDPOINT_DESCRIPTOR* ep = (const USB_ENDPOINT_DESCRIPTOR*)desc;
            uint8_t epType = ep->bmAttributes & USB_ENDPOINT_ATTR_TYPE_MASK;
            uint8_t epNumber = ep->bEndpointAddress & USB_ENDPOINT_NUMBER_MASK;
            bool in = (ep->bEndpointAddress & USB_ENDPOINT_DIR_IN) != 0;

            if (tracked && tracked->EndpointCount < USB_MAX_TRACKED_ENDPOINTS) {
                memcpy(&tracked->Endpoints[tracked->EndpointCount++], ep, sizeof(*ep));
            } else if (tracked) {
                DbgPrint("usb: slot=%u if=%u too many endpoints, endpoint %02x not published\n",
                         dev->SlotId, currentInterface, ep->bEndpointAddress);
            }

            UsbTrace("usb: endpoint slot=%u if=%u addr=%02x ep=%u dir=%s type=%s maxpkt=%u interval=%u\n",
                     dev->SlotId,
                     currentInterface,
                     ep->bEndpointAddress,
                     epNumber,
                     UsbDirectionName(ep->bEndpointAddress),
                     UsbEndpointTypeName(epType),
                     ep->wMaxPacketSize,
                     ep->bInterval);

            if (hidCandidate && in && epType == USB_ENDPOINT_TYPE_INTERRUPT && hidInterruptIn == 0) {
                hidInterruptIn = ep->bEndpointAddress;
            }
            if (massCandidate && epType == USB_ENDPOINT_TYPE_BULK) {
                if (in && massBulkIn == 0) massBulkIn = ep->bEndpointAddress;
                if (!in && massBulkOut == 0) massBulkOut = ep->bEndpointAddress;
            }
        } else if (type == USB_DESC_HID) {
            UsbTrace("usb: hid descriptor slot=%u if=%u len=%u\n",
                     dev->SlotId, currentInterface, len);
        } else if (type == USB_DESC_CS_INTERFACE) {
            UsbTrace("usb: class-specific interface descriptor slot=%u if=%u len=%u subtype=%u\n",
                     dev->SlotId, currentInterface, len, len >= 3 ? desc[2] : 0);
        }

        uint16_t next = (uint16_t)(off + len);
        if (hidCandidate && hidInterruptIn &&
            (next >= totalLength || (next + 1u < totalLength && dev->DescriptorBuffer[next + 1u] == USB_DESC_INTERFACE))) {
            UsbTrace("usb: class candidate slot=%u if=%u hid interrupt-in=%02x subclass=%02x proto=%02x\n",
                     dev->SlotId,
                     currentInterface,
                     hidInterruptIn,
                     currentSubClass,
                     currentProtocol);
            hidCandidate = false;
        }
        if (massCandidate && massBulkIn && massBulkOut &&
            (next >= totalLength || (next + 1u < totalLength && dev->DescriptorBuffer[next + 1u] == USB_DESC_INTERFACE))) {
            UsbTrace("usb: class candidate slot=%u if=%u mass-storage bulk-in=%02x bulk-out=%02x\n",
                     dev->SlotId,
                     currentInterface,
                     massBulkIn,
                     massBulkOut);
            massCandidate = false;
        }

        off = next;
    }

    if (hidCandidate && hidInterruptIn) {
        UsbTrace("usb: class candidate slot=%u if=%u hid interrupt-in=%02x subclass=%02x proto=%02x\n",
                 dev->SlotId,
                 currentInterface,
                 hidInterruptIn,
                 currentSubClass,
                 currentProtocol);
    }
    if (massCandidate && massBulkIn && massBulkOut) {
        UsbTrace("usb: class candidate slot=%u if=%u mass-storage bulk-in=%02x bulk-out=%02x\n",
                 dev->SlotId,
                 currentInterface,
                 massBulkIn,
                 massBulkOut);
    }
}



KSTATUS XhciDiscoverAndSetConfiguration(XHCI_CONTROLLER* xhci, XHCI_DEVICE* dev,
                                               const USB_DEVICE_DESCRIPTOR* devDesc) {
    if (devDesc->bNumConfigurations == 0) {
        DbgPrint("usb: slot=%u has no configurations\n", dev->SlotId);
        return STATUS_NOT_FOUND;
    }

    USB_CONFIGURATION_DESCRIPTOR config;
    memset(&config, 0, sizeof(config));
    uint16_t totalLength = 0;
    KSTATUS st = XhciGetConfigurationDescriptor(xhci, dev, 0, &config, &totalLength);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: slot=%u configuration descriptor read failed st=0x%x\n",
                 dev->SlotId, st);
        return st;
    }

    XhciLogConfigurationTree(xhci, dev, &config, totalLength);

    if (config.bConfigurationValue == 0) {
        DbgPrint("usb: slot=%u configuration value 0 is invalid for SET_CONFIGURATION\n",
                 dev->SlotId);
        return STATUS_DEVICE_ERROR;
    }

    st = XhciSetConfiguration(xhci, dev, config.bConfigurationValue);
    if (!KSUCCESS(st)) return st;

    if (!dev->UsbPdo) {
        st = UsbCreateDevicePdo(&xhci->Hcd, dev,
                                 dev->ParentDeviceObject ? dev->ParentDeviceObject :
                                     (xhci->Hcd.RootHubPdo ? xhci->Hcd.RootHubPdo : xhci->Hcd.ControllerPdo));
        if (!KSUCCESS(st)) {
            DbgPrint("usb: create device PDO slot=%u failed st=0x%x\n", dev->SlotId, st);
            return st;
        }
    }

    return STATUS_SUCCESS;
}


static void XhciRollbackEnumeratedSlot(XHCI_CONTROLLER* xhci, uint8_t slotId) {
    if (!xhci || slotId == 0 || slotId > XHCI_MAX_TRACKED_SLOTS)
        return;

    XHCI_DEVICE* dev = &xhci->Devices[slotId];

    if (dev->UsbPdo) {
        (void)UsbRemovePdoSubtree(dev->UsbPdo, true);
        dev->UsbPdo = (DEVICE_OBJECT*)0;
    }

    if (dev->UsbDevice) {
        dev->UsbDevice->State = UsbDeviceRemoved;
        UsbDereferenceDevice(dev->UsbDevice);
        dev->UsbDevice = (USB_DEVICE*)0;
    }

    if (dev->Allocated)
        (void)XhciDisableSlot(xhci, slotId);

    memset(dev, 0, sizeof(*dev));
}


static KSTATUS XhciEnumeratePathEx(XHCI_CONTROLLER* xhci,
                                    uint8_t logicalPort,
                                    uint8_t speed,
                                    const USB_PORT_PATH* path,
                                    DEVICE_OBJECT* parentDevice,
                                    uint8_t* outSlotId) {
    if (outSlotId) *outSlotId = 0;
    if (!xhci || !path || path->RootPort == 0)
        return STATUS_INVALID_PARAMETER;

    uint8_t slotId = 0;
    KSTATUS st = XhciEnableSlot(xhci, &slotId);
    if (!KSUCCESS(st)) return st;

    XHCI_DEVICE* dev = &xhci->Devices[slotId];
    memset(dev, 0, sizeof(*dev));
    dev->Allocated = true;
    dev->SlotId = slotId;
    dev->Port = path->RootPort;
    dev->Speed = speed;
    dev->Path = *path;
    dev->ParentDeviceObject = parentDevice;

    st = XhciAllocateDeviceMemory(xhci, dev);
    if (!KSUCCESS(st)) {
        XhciRollbackEnumeratedSlot(xhci, slotId);
        return st;
    }

    XhciBuildDefaultControlContext(xhci, dev);

    st = XhciAddressDevice(xhci, dev);
    if (!KSUCCESS(st)) {
        XhciRollbackEnumeratedSlot(xhci, slotId);
        return st;
    }
    dev->Addressed = true;

    USB_DEVICE_DESCRIPTOR desc;
    memset(&desc, 0, sizeof(desc));
    st = XhciGetDeviceDescriptorTimeout(xhci, dev, &desc, USB_XHCI_ENUM_CONTROL_TIMEOUT_US);
    if (!KSUCCESS(st)) {
        DbgPrint("usb: port %u slot %u descriptor read failed st=0x%x\n",
                 logicalPort, slotId, st);
        XhciRollbackEnumeratedSlot(xhci, slotId);
        return st;
    }

    dev->MaxPacketSize0 = desc.bMaxPacketSize0;
    dev->VendorId = desc.idVendor;
    dev->ProductId = desc.idProduct;
    dev->BcdUsb = desc.bcdUSB;
    dev->BcdDevice = desc.bcdDevice;
    dev->DeviceClass = desc.bDeviceClass;
    dev->DeviceSubclass = desc.bDeviceSubClass;
    dev->DeviceProtocol = desc.bDeviceProtocol;

    if (path->BehindHub) {
        DbgPrint("usb: device root-port=%u hub-slot=%u hub-port=%u slot=%u speed=%s(%u) vid=%04x pid=%04x class=%02x subclass=%02x proto=%02x mps0=%u configs=%u usb=%04x route=%05x\n",
                 path->RootPort,
                 path->ParentHubSlot,
                 path->ParentHubPort,
                 slotId,
                 XhciPortSpeedName(speed),
                 speed,
                 desc.idVendor,
                 desc.idProduct,
                 desc.bDeviceClass,
                 desc.bDeviceSubClass,
                 desc.bDeviceProtocol,
                 desc.bMaxPacketSize0,
                 desc.bNumConfigurations,
                 desc.bcdUSB,
                 path->RouteString);
    } else {
        DbgPrint("usb: device port=%u slot=%u speed=%s(%u) vid=%04x pid=%04x class=%02x subclass=%02x proto=%02x mps0=%u configs=%u usb=%04x\n",
                 logicalPort,
                 slotId,
                 XhciPortSpeedName(speed),
                 speed,
                 desc.idVendor,
                 desc.idProduct,
                 desc.bDeviceClass,
                 desc.bDeviceSubClass,
                 desc.bDeviceProtocol,
                 desc.bMaxPacketSize0,
                 desc.bNumConfigurations,
                 desc.bcdUSB);
    }

    st = XhciDiscoverAndSetConfiguration(xhci, dev, &desc);
    if (!KSUCCESS(st)) {
        XhciRollbackEnumeratedSlot(xhci, slotId);
        return st;
    }

    if (outSlotId) *outSlotId = slotId;
    return STATUS_SUCCESS;
}

KSTATUS XhciEnumeratePortEx(XHCI_CONTROLLER* xhci, uint8_t port, uint8_t* outSlotId) {
    if (outSlotId) *outSlotId = 0;

    uint8_t speed = 0;
    KSTATUS st = XhciResetPort(xhci, port, &speed);
    if (!KSUCCESS(st)) return st;

    USB_PORT_PATH path;
    memset(&path, 0, sizeof(path));
    path.RootPort = port;
    path.Depth = 0;
    path.RouteString = 0;
    path.BehindHub = false;

    DEVICE_OBJECT* parent = xhci && xhci->Hcd.RootHubPdo ? xhci->Hcd.RootHubPdo : (DEVICE_OBJECT*)0;
    return XhciEnumeratePathEx(xhci, port, speed, &path, parent, outSlotId);
}

KSTATUS XhciEnumerateHubPortEx(XHCI_CONTROLLER* xhci,
                               USB_HUB_EXTENSION* hub,
                               uint8_t hubPort,
                               uint8_t speed,
                               uint8_t* outSlotId) {
    if (outSlotId) *outSlotId = 0;
    if (!xhci || !hub || !hub->Device || !hub->Self || hubPort == 0)
        return STATUS_INVALID_PARAMETER;

    USB_DEVICE* parent = hub->Device;
    if (parent->Path.Depth >= 5u) {
        DbgPrint("usb-hub: hub slot=%u port=%u exceeds xHCI route depth\n",
                 parent->SlotId, hubPort);
        return STATUS_NOT_SUPPORTED;
    }
    if (hubPort > 15u) {
        DbgPrint("usb-hub: hub slot=%u port=%u cannot be represented in xHCI route string yet\n",
                 parent->SlotId, hubPort);
        return STATUS_NOT_SUPPORTED;
    }

    USB_PORT_PATH path = parent->Path;
    path.BehindHub = true;
    path.ParentHubSlot = parent->SlotId;
    path.ParentHubPort = hubPort;
    path.RouteString |= ((uint32_t)hubPort << (path.Depth * 4u));
    path.Depth++;

    return XhciEnumeratePathEx(xhci, hubPort, speed, &path, hub->Self, outSlotId);
}

KSTATUS XhciEnumeratePort(XHCI_CONTROLLER* xhci, uint8_t port) {
    return XhciEnumeratePortEx(xhci, port, (uint8_t*)0);
}


void XhciEnumerateRootPorts(XHCI_CONTROLLER* xhci) {
    uint32_t attempted = 0;
    uint32_t configured = 0;

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

        attempted++;
        uint8_t slot = 0;
        KSTATUS st = XhciEnumeratePortEx(xhci, port, &slot);
        if (KSUCCESS(st)) configured++;
        else DbgPrint("usb: root-port %u enumeration failed st=0x%x\n", port, st);
    }

    DbgPrint("usb: xHCI root-port enumeration complete; attempted=%u configured=%u\n",
             attempted, configured);
}