// USB composite driver. This file owns multi-interface device attachment and function handoff.

#include "usb.h"

DRIVER_OBJECT gUsbCompositeDriver;

static const IO_DEVICE_ID gUsbCompositeIds[] = {
    { IoBusUsb, { "USB\\COMPOSITE", sizeof("USB\\COMPOSITE") - 1 } },
};

static const IO_DRIVER_ID_TABLE gUsbCompositeIdTable = {
    gUsbCompositeIds,
    sizeof(gUsbCompositeIds) / sizeof(gUsbCompositeIds[0]),
};

// Implements the composite Create Interfaces operation.
static KSTATUS UsbCompositeCreateInterfaces(USB_COMPOSITE_EXTENSION* ext) {
    if (!ext || !ext->Device || ext->InterfacesCreated)
        return NTLSTATUS_SUCCESS;

    XHCI_DEVICE* xdev = (XHCI_DEVICE*)ext->Device->HcdDeviceContext;
    if (!xdev)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS st = UsbCreateInterfacePdos(ext->Device->Hcd, xdev, ext->Self);
    if (KSUCCESS(st))
        ext->InterfacesCreated = true;
    return st;
}

// Implements the composite Forward Pnp operation.
static KSTATUS UsbCompositeForwardPnp(DEVICE_OBJECT* device, IRP* irp, KSTATUS localStatus) {
    if (!KSUCCESS(localStatus)) {
        IoCompleteRequest(irp, localStatus, 0);
        return NTLSTATUS_PENDING;
    }
    return IoForwardIrpToLower(device, irp);
}

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

    // The composite driver owns two different device-object roles:
    //   * an FDO attached above a USB device PDO, and
    //   * interface PDOs that it creates for class drivers.
    // Forwarded PnP IRPs from HID/storage land on the interface PDO, so the
    // PDO must use the generic USB PDO PnP handler. Treating its
    // USB_PDO_EXTENSION as USB_COMPOSITE_EXTENSION corrupts the dispatch path.
    if (device->Type == DeviceObjectPhysical)
        return UsbDispatchPdoPnp(device, irp);

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

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

    switch ((IO_PNP_MINOR_FUNCTION)ios->MinorFunction) {
        case IRP_MN_START_DEVICE:
            device->PnpState = DevicePnpStarted;
            if (ext) {
                if (ext->DevicePdo)
                    ext->DevicePdo->PnpState = DevicePnpStarted;
                ext->Started = true;
                st = UsbCompositeCreateInterfaces(ext);
            }
            return UsbCompositeForwardPnp(device, irp, st);

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

        case IRP_MN_CANCEL_REMOVE_DEVICE:
            device->PnpState = DevicePnpStarted;
            return IoForwardIrpToLower(device, irp);

        case IRP_MN_SURPRISE_REMOVAL:
            device->PnpState = DevicePnpSurpriseRemoved;
            if (ext)
                ext->Started = false;
            (void)UsbRemoveDeviceObjectChildren(device, true);
            return IoForwardIrpToLower(device, irp);

        case IRP_MN_REMOVE_DEVICE:
            device->PnpState = DevicePnpDeleted;
            if (ext) {
                ext->Started = false;
                (void)UsbRemoveDeviceObjectChildren(device, false);
                if (ext->LowerDevice) {
                    KSTATUS lower = IoForwardIrpToLower(device, irp);
                    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;
                }
            }
            IoDeleteDevice(device);
            IoCompleteRequest(irp, NTLSTATUS_SUCCESS, 0);
            return NTLSTATUS_PENDING;

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

// Implements the composite Add Device operation.
KSTATUS UsbCompositeAddDevice(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 != UsbPdoDevice || !pdoExt->Device)
        return NTLSTATUS_INVALID_PARAMETER;

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

    (void)IoSetDeviceName(fdo, SYS_SV_LITERAL("usb-composite"));

    USB_COMPOSITE_EXTENSION* ext = (USB_COMPOSITE_EXTENSION*)fdo->DeviceExtension;
    ext->Self = fdo;
    ext->DevicePdo = physicalDeviceObject;
    ext->Device = pdoExt->Device;
    ext->LowerDevice = IoAttachDeviceToDeviceStack(fdo, physicalDeviceObject);
    if (!ext->LowerDevice) {
        IoDeleteDevice(fdo);
        return NTLSTATUS_FAILURE;
    }

    DbgPrint("usb: composite attached slot=%u\n", ext->Device->SlotId);
    return NTLSTATUS_SUCCESS;
}

// Initializes composite Driver.
KSTATUS UsbInitializeCompositeDriver(void) {
    memset(&gUsbCompositeDriver, 0, sizeof(gUsbCompositeDriver));
    IoInitializeDriverObject(&gUsbCompositeDriver, SYS_SV_LITERAL("usb-composite"));
    gUsbCompositeDriver.IdTable = &gUsbCompositeIdTable;
    gUsbCompositeDriver.DriverAddDevice = UsbCompositeAddDevice;
    gUsbCompositeDriver.MajorFunction[IRP_MJ_PNP] = UsbCompositeDispatchPnp;
    gUsbCompositeDriver.MajorFunction[IRP_MJ_INTERNAL_DEVICE_CONTROL] = UsbDispatchInternalDeviceControl;
    return NTLSTATUS_SUCCESS;
}