#include "usbkbd.h"
#include "fb.h"
#include "arch.h"
#include "paging.h"
#include "printf.h"
#include "util.h"
#include "dbg.h"
#include "mouse.h"

#define PCI_CONF_ADDR     0xCF8
#define PCI_CONF_DATA     0xCFC

#define XHCI_CMD_RUN      (1U << 0)
#define XHCI_CMD_HCRST    (1U << 1)

#define XHCI_STS_HCH      (1U << 0)
#define XHCI_STS_CNR      (1U << 11)

#define PORTSC_CCS        (1U << 0)
#define PORTSC_PED        (1U << 1)
#define PORTSC_PR         (1U << 4)
#define PORTSC_PP         (1U << 6)
#define PORTSC_CSC        (1U << 9)
#define PORTSC_WPR        (1U << 11)
#define PORTSC_LWS        (1U << 16)
#define PORTSC_PRC        (1U << 21)
#define PORTSC_WRC        (1U << 23)
#define PORTSC_SPEED_SH   10
#define PORTSC_SPEED_M   (0xF << 10)

#define XHCI_PORT_SPEED_FS  1
#define XHCI_PORT_SPEED_LS  2
#define XHCI_PORT_SPEED_HS  3
#define XHCI_PORT_SPEED_SS  4

#define TRB_TYPE_SHIFT     10

#define TRB_ENABLE_SLOT   (9  << TRB_TYPE_SHIFT)
#define TRB_DISABLE_SLOT  (10 << TRB_TYPE_SHIFT)
#define TRB_ADDRESS_DEV   (11 << TRB_TYPE_SHIFT)
#define TRB_BSR           (1U << 9)
#define TRB_CONFIG_EP     (12 << TRB_TYPE_SHIFT)
#define TRB_RESET_EP      (14 << TRB_TYPE_SHIFT)
#define TRB_LINK          (6  << TRB_TYPE_SHIFT)
#define TRB_SETUP_STAGE   (2  << TRB_TYPE_SHIFT)
#define TRB_STATUS_STAGE  (4  << TRB_TYPE_SHIFT)
#define TRB_NORMAL        (1  << TRB_TYPE_SHIFT)
#define TRB_DATA          (3  << TRB_TYPE_SHIFT)

#define TRB_DIR_IN        (1U << 16)
#define TRB_DIR_OUT       (0U << 16)
#define TRB_IOC           (1U << 5)
#define TRB_IDT           (1U << 6)
#define TRB_CHAIN         (1U << 4)
#define TRB_TC            (1U << 1)
#define TRB_TRT_NO_DATA   0
#define TRB_TRT_OUT       1
#define TRB_TRT_IN        2
#define TRB_TRT_SHIFT     16

#define EVT_TRANSFER      32
#define EVT_CMD_COMPLETE  33
#define EVT_PORT_STATUS   34

#define CMP_SUCCESS       1
#define COMP_CODE_SHIFT   24
#define COMP_CODE_MASK    0xFF

#define ERDP_EHB          (1U << 3)
#define ERDP_DCS          (1U << 0)

#define IMOD_INTERVAL     0xFA0

#define RING_SIZE 64
#define RING_MASK         (RING_SIZE - 1)

#define MAX_PORTS         16
#define CTRL_BUF_SIZE     256
#define MAX_REPORT_SIZE   64

#define USB_CLASS_HID     3
#define USB_SUBCLASS_BOOT 1
#define USB_PROTOCOL_KBD  1

#define USB_REQ_GET_STATUS     0x00
#define USB_REQ_CLEAR_FEATURE  0x01
#define USB_REQ_SET_FEATURE    0x03
#define USB_REQ_SET_ADDRESS    0x05
#define USB_REQ_GET_DESCRIPTOR 0x06
#define USB_REQ_SET_CONFIG     0x09
#define USB_REQ_SET_PROTOCOL   0x0B
#define USB_REQ_SET_IDLE       0x0A

#define USB_DESC_DEVICE       0x01
#define USB_DESC_CONFIG       0x02
#define USB_DESC_INTERFACE    0x04
#define USB_DESC_ENDPOINT     0x05
#define USB_DESC_HID          0x21
#define USB_DESC_HID_REPORT   0x22

struct usb_setup_pkt {
    uint8_t  bmReqType;
    uint8_t  bRequest;
    uint16_t wValue;
    uint16_t wIndex;
    uint16_t wLength;
} __attribute__((packed));

struct xhci_caps {
    volatile uint8_t  caplength;
    uint8_t  _rsvd0;
    volatile uint16_t hciversion;
    volatile uint32_t hcsparams1;
    volatile uint32_t hcsparams2;
    volatile uint32_t hcsparams3;
    volatile uint32_t hccparams1;
    volatile uint32_t dboff;
    volatile uint32_t rtsoff;
    volatile uint32_t hccparams2;
};

struct xhci_port_regs {
    volatile uint32_t portsc;
    volatile uint32_t portpmsc;
    volatile uint32_t portli;
    uint32_t _rsvd;
};

struct xhci_ops {
    volatile uint32_t usbcmd;
    volatile uint32_t usbsts;
    volatile uint32_t pagesize;
    uint8_t  _rsvd0[8];
    volatile uint32_t dnctrl;
    volatile uint64_t crcr;
    volatile uint32_t _rsvd_cfg[4];
    volatile uint64_t dcbaap;
    volatile uint32_t config;
    uint8_t  _gap2[0x3C4];
    volatile struct xhci_port_regs portregs[];
};

struct xhci_intr {
    volatile uint32_t iman;
    volatile uint32_t imod;
    volatile uint32_t erstsz;
    uint8_t  _rsvd0[4];
    volatile uint64_t erstba;
    volatile uint64_t erdp;
};

struct xhci_trb {
    volatile uint32_t dwords[4];
};

struct xhci_erst {
    volatile uint64_t seg_addr;
    volatile uint32_t seg_size;
    volatile uint32_t rsvd;
};

struct xhci_slot_ctx {
    volatile uint32_t word0;
    volatile uint32_t word1;
    volatile uint32_t word2;
    volatile uint32_t word3;
    volatile uint32_t word4;
    volatile uint32_t word5;
    volatile uint32_t word6;
    volatile uint32_t word7;
};

struct xhci_ep_ctx {
    volatile uint32_t word0;
    volatile uint32_t word1;
    volatile uint64_t dequeue;
    volatile uint32_t word4;
    volatile uint32_t word5;
    volatile uint32_t word6;
    volatile uint32_t word7;
};

/* Context stride: 32 or 64 bytes per entry, set from HCCPARAMS1 bit 2 */
static int ctx_stride = 32;

#define DEV_SLOT(buf)    ((struct xhci_slot_ctx*)(uint8_t*)(buf))
#define DEV_EP(buf, n)   ((struct xhci_ep_ctx*)((uint8_t*)(buf) + ctx_stride * (1 + (n))))
#define IN_CTRL(buf)     ((uint32_t*)(uint8_t*)(buf))
#define IN_SLOT(buf)     ((struct xhci_slot_ctx*)((uint8_t*)(buf) + 32))
#define IN_EP(buf, n)    ((struct xhci_ep_ctx*)((uint8_t*)(buf) + 32 + ctx_stride * (1 + (n))))

struct usb_desc_device {
    uint8_t  bLength;
    uint8_t  bDescriptorType;
    uint16_t bcdUSB;
    uint8_t  bDeviceClass;
    uint8_t  bDeviceSubClass;
    uint8_t  bDeviceProtocol;
    uint8_t  bMaxPacketSize0;
    uint16_t idVendor;
    uint16_t idProduct;
    uint16_t bcdDevice;
    uint8_t  iManufacturer;
    uint8_t  iProduct;
    uint8_t  iSerialNumber;
    uint8_t  bNumConfigurations;
};

struct usb_desc_config {
    uint8_t  bLength;
    uint8_t  bDescriptorType;
    uint16_t wTotalLength;
    uint8_t  bNumInterfaces;
    uint8_t  bConfigurationValue;
    uint8_t  iConfiguration;
    uint8_t  bmAttributes;
    uint8_t  bMaxPower;
};

struct usb_desc_interface {
    uint8_t  bLength;
    uint8_t  bDescriptorType;
    uint8_t  bInterfaceNumber;
    uint8_t  bAlternateSetting;
    uint8_t  bNumEndpoints;
    uint8_t  bInterfaceClass;
    uint8_t  bInterfaceSubClass;
    uint8_t  bInterfaceProtocol;
    uint8_t  iInterface;
};

struct usb_desc_endpoint {
    uint8_t  bLength;
    uint8_t  bDescriptorType;
    uint8_t  bEndpointAddress;
    uint8_t  bmAttributes;
    uint16_t wMaxPacketSize;
    uint8_t  bInterval;
};

struct usb_hid_descriptor {
    uint8_t  bLength;
    uint8_t  bDescriptorType;
    uint16_t bcdHID;
    uint8_t  bCountryCode;
    uint8_t  bNumDescriptors;
    uint8_t  bDescriptorType2;
    uint16_t wDescriptorLength;
} __attribute__((packed));

static uint64_t kernel_phys_base;
static int xhci_initialized;
static int kbd_found;

static volatile struct xhci_caps *caps;
static volatile struct xhci_ops  *ops;
static volatile uint32_t         *dbreg;
static volatile struct xhci_intr *ir;

static struct xhci_trb cmd_ring[RING_SIZE]     __attribute__((aligned(64)));
static struct xhci_trb evt_ring[RING_SIZE]     __attribute__((aligned(64)));
static struct xhci_erst erst                    __attribute__((aligned(64)));
static struct xhci_trb ep0_ring[RING_SIZE]     __attribute__((aligned(64)));

static uint64_t dcbaap[256]           __attribute__((aligned(64)));
static uint8_t dev_ctx_buf[64 * 32]   __attribute__((aligned(64)));
static uint8_t in_ctx_buf[64 * 33]    __attribute__((aligned(64)));

static uint64_t scratchpad_ptr_array[256] __attribute__((aligned(64)));
static uint8_t  scratchpad_pages[256][4096] __attribute__((aligned(4096)));

static int xhci_num_ports;

static int device_slot_id;
static int device_speed;
static int device_port;
static int ep0_mps;

/* Per-device slot IDs — prevent cross-contamination */
static int kbd_device_slot_id;
static int kbd_port;

static int kbd_ep_addr;
static int kbd_ep_interval;
static int kbd_eps_idx;
static int kbd_report_size;
static int kbd_use_boot_protocol;
static int kbd_hid_itf_num;

static uint8_t ctrl_buf[CTRL_BUF_SIZE] __attribute__((aligned(64)));
static uint8_t hid_report_desc_buf[256] __attribute__((aligned(64)));

static struct xhci_trb kbd_ring[RING_SIZE]   __attribute__((aligned(64)));
static uint8_t kbd_report_buf[MAX_REPORT_SIZE] __attribute__((aligned(64)));
static uint8_t prev_kbd_report[MAX_REPORT_SIZE] __attribute__((aligned(64)));

static int cmd_pcs;
static int evt_ccs;
static int evt_deq_idx;
static int ep0_pcs;
static int ep0_write_pos;
static int kbd_ring_pcs;
static int kbd_ring_write_pos;
static int kbd_ring_active;

static uint8_t last_processed_report[MAX_REPORT_SIZE];
static int report_ready;

/* ══════════════════════════════════════════════
   USB Mouse support
   ══════════════════════════════════════════════ */
#define USB_PROTOCOL_MOUSE 2
#define MOUSE_REPORT_SIZE  4
#define MOUSE_EP_RING_SIZE 64

static int mouse_found;
static int mouse_slot_id;
static int mouse_port;
static int mouse_eps_idx;
static int mouse_ep_addr;
static int mouse_ep_interval;
static int mouse_report_size;

static uint8_t mouse_dev_ctx_buf[64 * 32] __attribute__((aligned(64)));
static struct xhci_trb      mouse_ep_ring[MOUSE_EP_RING_SIZE] __attribute__((aligned(64)));
static uint8_t mouse_report_buf[MOUSE_REPORT_SIZE] __attribute__((aligned(64)));
static uint8_t prev_mouse_report[MOUSE_REPORT_SIZE] __attribute__((aligned(64)));
static uint8_t last_mouse_processed[MOUSE_REPORT_SIZE] __attribute__((aligned(64)));

static int mouse_ring_pcs;
static int mouse_ring_write_pos;
static int mouse_ring_active;
static int mouse_report_ready;

/* Mouse state */
int mouse_buttons;
int mouse_dx;
int mouse_dy;
int mouse_wheel;

static uint64_t v2p(const void *v) {
    uint64_t x = (uint64_t)(uintptr_t)v;
    if(x >= 0xFFFF800000000000ULL)
        return x - 0xFFFF800000000000ULL + kernel_phys_base;
    return x;
}

static void delay_us(uint64_t us) {
    uint64_t start = rdtsc();
    uint64_t ticks = us * 4000ULL;
    while(rdtsc() - start < ticks) __asm__ volatile("");
}

static void write_trb(volatile struct xhci_trb *ring, int idx,
                       uint64_t param, uint32_t status, uint32_t control)
{
    ring[idx].dwords[0] = (uint32_t)(param);
    ring[idx].dwords[1] = (uint32_t)(param >> 32);
    ring[idx].dwords[2] = status;
    ring[idx].dwords[3] = control;
}

static void write_link_trb(struct xhci_trb *ring, int n, int pcs)
{
    write_trb(ring, n - 1, v2p(ring), 0,
              TRB_LINK | TRB_TC | (pcs));
}

static void mmio_write64(volatile uint32_t *base, uint64_t val)
{
    base[0] = (uint32_t)val;
    base[1] = (uint32_t)(val >> 32);
}

static uint32_t pci_read(uint8_t bus, uint8_t dev, uint8_t func, uint8_t off)
{
    uint32_t addr = (1U << 31) | ((uint32_t)bus << 16) |
                    ((uint32_t)(dev & 0x1F) << 11) |
                    ((uint32_t)(func & 7) << 8) | (off & 0xFC);
    outl(PCI_CONF_ADDR, addr);
    return inl(PCI_CONF_DATA);
}

static void pci_write(uint8_t bus, uint8_t dev, uint8_t func, uint8_t off, uint32_t val)
{
    uint32_t addr = (1U << 31) | ((uint32_t)bus << 16) |
                    ((uint32_t)(dev & 0x1F) << 11) |
                    ((uint32_t)(func & 7) << 8) | (off & 0xFC);
    outl(PCI_CONF_ADDR, addr);
    outl(PCI_CONF_DATA, val);
}

static uint64_t pci_find_xhci(void)
{
    for(uint16_t bus = 0; bus < 256; bus++) {
        for(uint8_t dev = 0; dev < 32; dev++) {
            for(uint8_t func = 0; func < 8; func++) {
                uint32_t id = pci_read((uint8_t)bus, dev, func, 0);
                if(id == 0xFFFFFFFF || id == 0) continue;
                uint32_t cr = pci_read((uint8_t)bus, dev, func, 8);
                uint8_t base = (uint8_t)(cr >> 24);  /* byte at 0x0B = Base Class */
                uint8_t sub  = (uint8_t)(cr >> 16);  /* byte at 0x0A = Subclass */
                uint8_t prog = (uint8_t)(cr >> 8);   /* byte at 0x09 = Prog IF */
                if(base == 0x0C && sub == 0x03 && prog == 0x30) {
                    uint32_t bar0 = pci_read((uint8_t)bus, dev, func, 0x10);
                    uint64_t mbar = bar0 & ~0xFULL;
                    if(bar0 & 4) {
                        uint32_t bar1 = pci_read((uint8_t)bus, dev, func, 0x14);
                        mbar |= ((uint64_t)bar1 << 32);
                    }
                    if(mbar == 0) continue;
                    dbg_printf("[xHCI] found at %d:%d.%d BAR=0x%llx class=%02x/%02x/%02x\n",
                        bus, dev, func, (unsigned long long)mbar, base, sub, prog);
                    uint32_t cmd = pci_read((uint8_t)bus, dev, func, 4);
                    cmd |= 6;
                    pci_write((uint8_t)bus, dev, func, 4, cmd);
                    return mbar;
                }
            }
        }
    }
    return 0;
}

static const char *xhci_speed_str(int s)
{
    switch(s) {
        case 0: return "?";
        case XHCI_PORT_SPEED_FS: return "FS";
        case XHCI_PORT_SPEED_LS: return "LS";
        case XHCI_PORT_SPEED_HS: return "HS";
        case XHCI_PORT_SPEED_SS: return "SS";
        default: return "??";
    }
}

static void xhci_reset(void)
{
    ops->usbcmd = XHCI_CMD_HCRST;
    delay_us(10000);
    uint64_t deadline = rdtsc() + 4000000000ULL;
    while(rdtsc() < deadline) {
        if(!(ops->usbsts & XHCI_STS_CNR)) return;
    }
    dbg_printf("[xhci] RESET timeout! usbsts=0x%08x\n", ops->usbsts);
}

static void xhci_start(void)
{
    ops->usbcmd = XHCI_CMD_RUN;
    uint64_t deadline = rdtsc() + 2000000000ULL;
    while(rdtsc() < deadline) {
        if(!(ops->usbsts & XHCI_STS_HCH)) return;
    }
    dbg_printf("[xhci] START timeout! usbsts=0x%08x\n", ops->usbsts);
}

static int evt_available(void)
{
    return (evt_ring[evt_deq_idx].dwords[3] & 1) == (uint32_t)evt_ccs;
}

static int evt_get_trb_type(const volatile struct xhci_trb *ev)
{
    return (ev->dwords[3] >> TRB_TYPE_SHIFT) & 0x3F;
}

static int evt_get_completion(const volatile struct xhci_trb *ev)
{
    return (ev->dwords[2] >> COMP_CODE_SHIFT) & COMP_CODE_MASK;
}

static uint32_t evt_get_slot_id(const volatile struct xhci_trb *ev)
{
    return (ev->dwords[3] >> 24) & 0xFF;
}

static uint32_t evt_get_port_id(const volatile struct xhci_trb *ev)
{
    return (ev->dwords[0] >> 24) & 0xFF;
}

static void evt_advance(void)
{
    evt_deq_idx = (evt_deq_idx + 1) & RING_MASK;
    if(evt_deq_idx == 0)
        evt_ccs ^= 1;
    uint64_t erdp_phys = v2p(&evt_ring[evt_deq_idx]);
    mmio_write64((volatile uint32_t*)&ir->erdp, (erdp_phys & ~0xFULL) | (uint64_t)evt_ccs);
}

static void evt_drain_stale(void)
{
    for(int i = 0; i < 64; i++) {
        if(!evt_available()) break;
        evt_advance();
    }
}

static int wait_for_event(int type)
{
    uint64_t deadline = rdtsc() + 20000000000ULL; /* ~5s at 4GHz */
    while(rdtsc() < deadline) {
        if(evt_available()) {
            volatile struct xhci_trb *ev = &evt_ring[evt_deq_idx];
            int etype = evt_get_trb_type(ev);
            int comp = evt_get_completion(ev);
            if(etype == type || type == -1) {
                uint32_t slot = evt_get_slot_id(ev);
                evt_advance();
                return (comp == CMP_SUCCESS) ? (int)slot : -(int)comp;
            }
            evt_advance();
        }
    }
    dbg_printf("[xhci] wait_for_event(%d) timeout\n", type);
    return -99;
}

static void cmd_ring_doorbell(void)
{
    sfence();
    dbreg[0] = 0;
}

static int cmd_enable_slot(void)
{
    write_trb(cmd_ring, 0, 0, 0, TRB_ENABLE_SLOT | (cmd_pcs));
    write_link_trb(cmd_ring, 2, cmd_pcs);
    cmd_pcs ^= 1;
    cmd_ring_doorbell();
    int r = wait_for_event(EVT_CMD_COMPLETE);
    dbg_printf("[xhci] enable_slot -> %d\n", r);
    return r;
}

static int cmd_disable_slot(int slot_id)
{
    write_trb(cmd_ring, 0, 0, 0, TRB_DISABLE_SLOT | (cmd_pcs) | ((uint32_t)slot_id << 24));
    write_link_trb(cmd_ring, 2, cmd_pcs);
    cmd_pcs ^= 1;
    cmd_ring_doorbell();
    return wait_for_event(EVT_CMD_COMPLETE);
}

static int cmd_address_device(int slot_id, int use_bsr)
{
    memset(in_ctx_buf, 0, ctx_stride * 33);

    IN_CTRL(in_ctx_buf)[0] = 0;
    IN_CTRL(in_ctx_buf)[1] = (1 << 0) | (1 << 1);

    int mps = ep0_mps;
    if(!mps)
        mps = (device_speed == XHCI_PORT_SPEED_HS) ? 64 : 8;

    IN_SLOT(in_ctx_buf)->word0 = (device_speed << 20) | (1 << 27);
    IN_SLOT(in_ctx_buf)->word1 = (device_port << 16);

    IN_EP(in_ctx_buf, 0)->word0 = mps;
    IN_EP(in_ctx_buf, 0)->word1 = (4 << 3) | (3 << 6);
    IN_EP(in_ctx_buf, 0)->dequeue = v2p(ep0_ring) | 1;

    uint32_t ctrl = TRB_ADDRESS_DEV | (cmd_pcs) | TRB_IOC | ((uint32_t)slot_id << 24);
    if(use_bsr) ctrl |= TRB_BSR;
    write_trb(cmd_ring, 0, v2p(in_ctx_buf), 0, ctrl);
    write_link_trb(cmd_ring, 2, cmd_pcs);
    cmd_pcs ^= 1;
    cmd_ring_doorbell();
    int r = wait_for_event(EVT_CMD_COMPLETE);
    dbg_printf("[xhci] address_device(slot=%d, port=%d, speed=%d, BSR=%d) -> %d\n",
               slot_id, device_port, device_speed, use_bsr, r);
    return r;
}

static int cmd_reset_endpoint(int slot_id, int dci)
{
    write_trb(cmd_ring, 0, 0, (uint32_t)dci,
              TRB_RESET_EP | (cmd_pcs) | TRB_IOC | ((uint32_t)slot_id << 24));
    write_link_trb(cmd_ring, 2, cmd_pcs);
    cmd_pcs ^= 1;
    cmd_ring_doorbell();
    return wait_for_event(EVT_CMD_COMPLETE);
}

static int ep0_do_control(const struct usb_setup_pkt *setup, uint8_t *data)
{
    int idx = ep0_write_pos;
    uint32_t setup_dwords[2];
    setup_dwords[0] = (uint32_t)setup->bmReqType |
                      ((uint32_t)setup->bRequest << 8) |
                      ((uint32_t)setup->wValue << 16);
    setup_dwords[1] = (uint32_t)setup->wIndex |
                      ((uint32_t)setup->wLength << 16);

    uint32_t trt, status_dir;
    if(setup->wLength == 0) {
        trt = TRB_TRT_NO_DATA;
        status_dir = TRB_DIR_IN;
    } else if(setup->bmReqType & 0x80) {
        trt = TRB_TRT_IN;
        status_dir = TRB_DIR_OUT;
    } else {
        trt = TRB_TRT_OUT;
        status_dir = TRB_DIR_IN;
    }

    write_trb(ep0_ring, idx++,
              ((uint64_t)setup_dwords[1] << 32) | setup_dwords[0],
              (uint32_t)8 | ((uint32_t)trt << 16),
              TRB_SETUP_STAGE | (ep0_pcs) | TRB_CHAIN | TRB_IDT);

    if(setup->wLength > 0 && data) {
        uint32_t data_dir = (setup->bmReqType & 0x80) ? TRB_DIR_IN : TRB_DIR_OUT;
        write_trb(ep0_ring, idx++,
                  v2p(data), setup->wLength,
                  TRB_DATA | (ep0_pcs) | TRB_CHAIN | data_dir);
    }

    write_trb(ep0_ring, idx++, 0, 0,
              TRB_STATUS_STAGE | (ep0_pcs) | status_dir | TRB_IOC);

    ep0_write_pos = idx;

    sfence();
    dbreg[device_slot_id] = 1;

    return wait_for_event(EVT_TRANSFER);
}

static int usb_get_device_descriptor(struct usb_desc_device *desc, int len)
{
    struct usb_setup_pkt s = {
        .bmReqType = 0x80,
        .bRequest = USB_REQ_GET_DESCRIPTOR,
        .wValue = USB_DESC_DEVICE << 8,
        .wIndex = 0,
        .wLength = len,
    };
    memset(ctrl_buf, 0, len);
    int r = ep0_do_control(&s, ctrl_buf);
    if(r < 0) return r;
    memcpy(desc, ctrl_buf, len);
    return r;
}

static int usb_get_config_descriptor_raw(uint8_t *buf, int len)
{
    struct usb_setup_pkt s = {
        .bmReqType = 0x80,
        .bRequest = USB_REQ_GET_DESCRIPTOR,
        .wValue = USB_DESC_CONFIG << 8,
        .wIndex = 0,
        .wLength = len,
    };
    return ep0_do_control(&s, buf);
}

static int usb_get_hid_descriptor(int itf_num, struct usb_hid_descriptor *desc)
{
    struct usb_setup_pkt s = {
        .bmReqType = 0x81,
        .bRequest = USB_REQ_GET_DESCRIPTOR,
        .wValue = USB_DESC_HID << 8,
        .wIndex = (uint16_t)itf_num,
        .wLength = sizeof(*desc),
    };
    memset(ctrl_buf, 0, sizeof(*desc));
    int r = ep0_do_control(&s, ctrl_buf);
    if(r < 0) return r;
    memcpy(desc, ctrl_buf, sizeof(*desc));
    return r;
}

static int usb_get_hid_report_descriptor(int itf_num, uint8_t *buf, int len)
{
    struct usb_setup_pkt s = {
        .bmReqType = 0x81,
        .bRequest = USB_REQ_GET_DESCRIPTOR,
        .wValue = USB_DESC_HID_REPORT << 8,
        .wIndex = (uint16_t)itf_num,
        .wLength = len,
    };
    return ep0_do_control(&s, buf);
}

static int usb_set_configuration(uint8_t config)
{
    struct usb_setup_pkt s = {
        .bmReqType = 0x00,
        .bRequest = USB_REQ_SET_CONFIG,
        .wValue = config,
        .wIndex = 0,
        .wLength = 0,
    };
    return ep0_do_control(&s, NULL);
}

static int parse_hid_report_input_size(const uint8_t *desc, int len)
{
    int total_bits = 0;
    int i = 0;
    int global_rsize = 8;
    int global_rcount = 1;

    while(i + 1 < len) {
        uint8_t item = desc[i];
        int tag = (item >> 4) & 0xF;
        int type = (item >> 2) & 3;
        int size = item & 3;

        int data_size;
        switch(size) {
            case 0: data_size = 0; break;
            case 1: data_size = 1; break;
            case 2: data_size = 2; break;
            case 3: data_size = 4; break;
            default: data_size = 0; break;
        }

        if(i + 1 + data_size > len) break;

        switch(type) {
        case 0:
            if(tag == 8)
                total_bits += global_rsize * global_rcount;
            break;
        case 1:
            if(tag == 7) {
                if(data_size == 0) global_rsize = 0;
                else if(data_size == 1) global_rsize = desc[i + 1];
                else if(data_size == 2) global_rsize = desc[i + 1] | (desc[i + 2] << 8);
                else if(data_size == 4)
                    global_rsize = desc[i + 1] | ((int)desc[i + 2] << 8) |
                                   ((int)desc[i + 3] << 16) | ((int)desc[i + 4] << 24);
            } else if(tag == 9) {
                if(data_size == 0) global_rcount = 0;
                else if(data_size == 1) global_rcount = desc[i + 1];
                else if(data_size == 2) global_rcount = desc[i + 1] | (desc[i + 2] << 8);
                else if(data_size == 4)
                    global_rcount = desc[i + 1] | ((int)desc[i + 2] << 8) |
                                    ((int)desc[i + 3] << 16) | ((int)desc[i + 4] << 24);
            }
            break;
        case 2:
            break;
        case 3:
            if(tag == 15 && i + 2 < len) {
                int long_len = desc[i + 1];
                if(i + 1 + long_len >= len) break;
                i += 1 + long_len;
                continue;
            }
            break;
        }

        i += 1 + data_size;
    }

    int total_bytes = (total_bits + 7) / 8;
    if(total_bytes < 8) total_bytes = 8;
    if(total_bytes > MAX_REPORT_SIZE) total_bytes = MAX_REPORT_SIZE;
    return total_bytes;
}

static int setup_rings(void)
{
    memset((void*)cmd_ring, 0, sizeof(cmd_ring));
    cmd_pcs = 1;
    write_link_trb(cmd_ring, RING_SIZE, cmd_pcs);

    memset((void*)evt_ring, 0, sizeof(evt_ring));
    evt_ccs = 1;
    evt_deq_idx = 0;

    memset(&erst, 0, sizeof(erst));
    erst.seg_addr = v2p(evt_ring);
    erst.seg_size = RING_SIZE;

    memset((void*)ep0_ring, 0, sizeof(ep0_ring));
    ep0_pcs = 1;
    ep0_write_pos = 0;

    memset((void*)kbd_ring, 0, sizeof(kbd_ring));
    kbd_ring_pcs = 1;
    kbd_ring_write_pos = 0;
    kbd_ring_active = 0;

    return 1;
}

static void reset_ep0_ring(void)
{
    memset((void*)ep0_ring, 0, sizeof(ep0_ring));
    ep0_pcs = 1;
    ep0_write_pos = 0;
}

static int xhci_hw_init(void)
{
    setup_rings();

    {
        uint32_t hsp2 = caps->hcsparams2;
        int sp_bufs = (((hsp2 >> 27) & 0x1F) << 5) | (hsp2 & 0x1F);
        if(sp_bufs > 256) sp_bufs = 256;
        if(sp_bufs > 0) {
            for(int i = 0; i < sp_bufs; i++)
                scratchpad_ptr_array[i] = v2p(&scratchpad_pages[i]);
            dcbaap[0] = v2p(scratchpad_ptr_array);
        } else {
            dcbaap[0] = 0;
        }
    }

    memset(dev_ctx_buf, 0, sizeof(dev_ctx_buf));
    xhci_reset();

    mmio_write64((volatile uint32_t*)&ops->dcbaap, v2p(dcbaap));
    mmio_write64((volatile uint32_t*)&ops->crcr, v2p(cmd_ring) | (uint64_t)cmd_pcs);

    ir->erstsz = 1;
    mmio_write64((volatile uint32_t*)&ir->erstba, v2p(&erst));
    uint64_t erdp_v = v2p(evt_ring);
    mmio_write64((volatile uint32_t*)&ir->erdp, erdp_v | 1);
    ir->iman = 2;

    ops->config = (caps->hcsparams1 & 0xFF);

    xhci_start();

    ir->imod = IMOD_INTERVAL;
    ir->iman = 2;

    evt_drain_stale();

    for(int p = 0; p < xhci_num_ports; p++)
        ops->portregs[p].portsc |= PORTSC_PP;

    delay_us(20000); /* 20ms after PP per spec */

    return 1;
}

static int xhci_init_port_ctx(int probe_port, uint8_t *ctx)
{
    device_port = probe_port;
    device_slot_id = 0;

    uint32_t ps = ops->portregs[device_port - 1].portsc;
    device_speed = (ps >> PORTSC_SPEED_SH) & 0xF;
    dbg_printf("[xhci] port %d init: sc=0x%08x speed=%d PED=%d CCS=%d\n",
               probe_port, ps, device_speed, (ps >> 2) & 1, ps & 1);

    /* Clear any pending status change bits before reset */
    {
        uint32_t chg = ps & (PORTSC_CSC | PORTSC_PRC | PORTSC_WRC);
        if(chg)
            ops->portregs[device_port - 1].portsc = ps | chg;
        delay_us(100);
    }

    ops->portregs[device_port - 1].portsc |= PORTSC_PR;
    delay_us(1000);
    int reset_ok = 0;
    {
        uint64_t deadline = rdtsc() + 2000000000ULL;
        while(rdtsc() < deadline) {
            if(!(ops->portregs[device_port - 1].portsc & PORTSC_PR)) {
                reset_ok = 1;
                break;
            }
        }
    }
    if(!reset_ok) {
        ps = ops->portregs[device_port - 1].portsc;
        int speed = (ps >> PORTSC_SPEED_SH) & 0xF;
        if(speed == XHCI_PORT_SPEED_SS) {
            ops->portregs[device_port - 1].portsc |= PORTSC_WPR;
            delay_us(1000);
            uint64_t wdead = rdtsc() + 2000000000ULL;
            while(rdtsc() < wdead) {
                if(!(ops->portregs[device_port - 1].portsc & PORTSC_WPR)) {
                    reset_ok = 1;
                    break;
                }
            }
        }
    }
    if(!reset_ok) {
        dbg_printf("[xhci] port %d: reset FAILED, sc=0x%08x\n",
                   probe_port, ops->portregs[device_port - 1].portsc);
        return 0;
    }

    {
        uint32_t sc = ops->portregs[device_port - 1].portsc;
        ops->portregs[device_port - 1].portsc = sc | (PORTSC_CSC | PORTSC_PRC | PORTSC_WRC);
    }

    {
        uint32_t ps_r = ops->portregs[device_port - 1].portsc;
        device_speed = (ps_r >> PORTSC_SPEED_SH) & 0xF;
    }

    int slot = cmd_enable_slot();
    if(slot <= 0) return 0;

    dcbaap[slot] = v2p(ctx);
    device_slot_id = slot;
    ep0_mps = (device_speed == XHCI_PORT_SPEED_HS) ? 64 : 8;
    DEV_EP(ctx, 0)->word0 = ep0_mps;

    int r = cmd_address_device(slot, 0);
    if(r < 0) {
        cmd_disable_slot(slot);
        reset_ep0_ring();
        memset(ctx, 0, ctx_stride * 32);
        device_slot_id = 0;
        return 0;
    }

    reset_ep0_ring();
    return 1;
}

static int xhci_init_port(int probe_port)
{
    return xhci_init_port_ctx(probe_port, dev_ctx_buf);
}

static int xhci_enumerate_kbd(void)
{
    kbd_report_size = 8;
    kbd_use_boot_protocol = 0;
    kbd_hid_itf_num = 0;

    struct usb_desc_device dev_desc;
    memset(&dev_desc, 0, sizeof(dev_desc));
    int r2 = usb_get_device_descriptor(&dev_desc, 8);
    if(r2 < 0 || dev_desc.bMaxPacketSize0 == 0) return 0;

    int new_mps = dev_desc.bMaxPacketSize0;
    if(new_mps != 8 && new_mps != 16 && new_mps != 32 && new_mps != 64)
        new_mps = 8;
    ep0_mps = new_mps;
    DEV_EP(dev_ctx_buf, 0)->word0 = ep0_mps;

    int r = cmd_address_device(device_slot_id, 0);
    if(r < 0) return 0;

    reset_ep0_ring();

    memset(&dev_desc, 0, sizeof(dev_desc));
    r2 = usb_get_device_descriptor(&dev_desc, 18);
    if(r2 < 0) return 0;

    int is_kbd = 0;
    int kbd_ep_mps = 8;

    {
        uint8_t cfg_buf[256];
        memset(cfg_buf, 0, sizeof(cfg_buf));
        r2 = usb_get_config_descriptor_raw(cfg_buf, 9);
        if(r2 < 0) return 0;

        uint16_t total_len = cfg_buf[2] | ((uint16_t)cfg_buf[3] << 8);
        if(total_len > 256) total_len = 256;
        r2 = usb_get_config_descriptor_raw(cfg_buf, total_len);
        if(r2 < 0) return 0;

        int cur_class = 0, cur_sub = 0, cur_proto = 0;
        int found_hid_itf = -1;
        int offset = 0;
        while(offset + 1 < total_len) {
            int dlen = cfg_buf[offset];
            int dtype = cfg_buf[offset + 1];
            if(dlen < 2) break;
            if(dtype == USB_DESC_INTERFACE) {
                struct usb_desc_interface *iface =
                    (struct usb_desc_interface*)(cfg_buf + offset);
                cur_class = iface->bInterfaceClass;
                cur_sub = iface->bInterfaceSubClass;
                cur_proto = iface->bInterfaceProtocol;
                if(cur_class == USB_CLASS_HID)
                    found_hid_itf = iface->bInterfaceNumber;
                else
                    found_hid_itf = -1;
            }
            if(dtype == USB_DESC_ENDPOINT && cur_class == USB_CLASS_HID) {
                struct usb_desc_endpoint *ep =
                    (struct usb_desc_endpoint*)(cfg_buf + offset);
                if((ep->bmAttributes & 0x03) == 3 &&
                   (ep->bEndpointAddress & 0x80)) {
                    kbd_ep_addr = ep->bEndpointAddress;
                    kbd_ep_interval = ep->bInterval;
                    kbd_ep_mps = ep->wMaxPacketSize;
                    kbd_hid_itf_num = found_hid_itf;
                    is_kbd = 1;
                    if(cur_sub == USB_SUBCLASS_BOOT &&
                       cur_proto == USB_PROTOCOL_KBD)
                        kbd_use_boot_protocol = 1;
                }
            }
            offset += dlen;
        }
    }

    if(!is_kbd) return 0;

    kbd_report_size = kbd_ep_mps;
    if(kbd_report_size < 8) kbd_report_size = 8;
    if(kbd_report_size > MAX_REPORT_SIZE) kbd_report_size = MAX_REPORT_SIZE;

    /*
     * For non-boot-protocol keyboards, read the HID descriptor
     * and report descriptor to determine the exact report size.
     * If the device doesn't support this, the EP wMaxPacketSize
     * (set above) is used as a fallback.
     */
    if(!kbd_use_boot_protocol) {
        struct usb_hid_descriptor hdesc;
        int r3 = usb_get_hid_descriptor(kbd_hid_itf_num, &hdesc);
        if(r3 >= 0 && hdesc.bDescriptorType == USB_DESC_HID &&
           hdesc.wDescriptorLength > 0 &&
           hdesc.wDescriptorLength <= (int)sizeof(hid_report_desc_buf)) {
            int rdesc_len = hdesc.wDescriptorLength;
            int r4 = usb_get_hid_report_descriptor(kbd_hid_itf_num,
                                                    hid_report_desc_buf,
                                                    rdesc_len);
            if(r4 < 0) {
                cmd_reset_endpoint(device_slot_id, 1);
                reset_ep0_ring();
            }
            if(r4 >= 0) {
                int parsed_size = parse_hid_report_input_size(
                    hid_report_desc_buf, rdesc_len);
                if(parsed_size >= 8 && parsed_size <= MAX_REPORT_SIZE)
                    kbd_report_size = parsed_size;
            }
        } else {
            int dci = 1;
            if(r3 < 0) {
                cmd_reset_endpoint(device_slot_id, dci);
                reset_ep0_ring();
            }
        }
    }

    {
        int r3 = usb_set_configuration(1);
        if(r3 < 0) {
            fb_printf("usbkbd: SET_CONFIG failed (%d)\n", r3);
            return 0;
        }
    }

    if(kbd_use_boot_protocol) {
        struct usb_setup_pkt sp = {
            .bmReqType = 0x21,
            .bRequest = USB_REQ_SET_PROTOCOL,
            .wValue = 0,
            .wIndex = (uint16_t)kbd_hid_itf_num,
            .wLength = 0,
        };
        int r3 = ep0_do_control(&sp, NULL);
        if(r3 >= 0) {
            fb_puts("usbkbd: boot protocol enabled\n");
        } else {
            fb_puts("usbkbd: boot proto N/A, using report protocol\n");
            kbd_use_boot_protocol = 0;
        }
    }

    {
        struct usb_setup_pkt sp = {
            .bmReqType = 0x21,
            .bRequest = USB_REQ_SET_IDLE,
            .wValue = 0,
            .wIndex = (uint16_t)kbd_hid_itf_num,
            .wLength = 0,
        };
        ep0_do_control(&sp, NULL);
    }

    {
        int ep_num = kbd_ep_addr & 0x0F;
        int is_in = (kbd_ep_addr & 0x80) ? 1 : 0;
        int dci = ep_num * 2 + is_in;
        int eps_idx = dci - 1;
        kbd_eps_idx = eps_idx;

        if(dci < 2 || dci > 31 || eps_idx < 0 || eps_idx > 30) {
            fb_printf("usbkbd: bad EP address %02x\n", kbd_ep_addr);
            return 0;
        }

        memset(in_ctx_buf, 0, ctx_stride * 33);
        IN_CTRL(in_ctx_buf)[0] = 0;
        IN_CTRL(in_ctx_buf)[1] = (1 << 0) | (1 << dci);

        IN_SLOT(in_ctx_buf)->word0 = (device_speed << 20) | ((eps_idx + 1) << 27);
        IN_SLOT(in_ctx_buf)->word1 = (device_port << 16);

        IN_EP(in_ctx_buf, 0)->word0 = ep0_mps;
        IN_EP(in_ctx_buf, 0)->word1 = (4 << 3) | (3 << 6);
        IN_EP(in_ctx_buf, 0)->dequeue = v2p(ep0_ring) | 1;

        int ep_type = is_in ? 7 : 3;
        IN_EP(in_ctx_buf, eps_idx)->word0 = kbd_report_size;
        IN_EP(in_ctx_buf, eps_idx)->word1 = (ep_type << 3) | (3 << 6);
        IN_EP(in_ctx_buf, eps_idx)->dequeue = v2p(kbd_ring) | 1;

        write_trb(cmd_ring, 0, v2p(in_ctx_buf), 0,
                  TRB_CONFIG_EP | (cmd_pcs) | TRB_IOC | ((uint32_t)device_slot_id << 24));
        write_link_trb(cmd_ring, 2, cmd_pcs);
        cmd_pcs ^= 1;
        cmd_ring_doorbell();
        r = wait_for_event(EVT_CMD_COMPLETE);
        if(r < 0) {
            fb_printf("usbkbd: CONFIG_EP failed (%d)\n", r);
            return 0;
        }
    }

    memset(kbd_report_buf, 0, kbd_report_size);
    memset(prev_kbd_report, 0, kbd_report_size);
    memset(last_processed_report, 0, kbd_report_size);

    write_link_trb(kbd_ring, RING_SIZE, 1);
    kbd_ring_pcs = 1;
    kbd_ring_write_pos = 1;

    write_trb(kbd_ring, 0, v2p(kbd_report_buf), kbd_report_size,
              TRB_NORMAL | (kbd_ring_pcs) | TRB_IOC);
    kbd_ring_active = 1;
    kbd_device_slot_id = device_slot_id;

    {
        int dci_kbd = (kbd_ep_addr & 0x0F) * 2 + ((kbd_ep_addr >> 7) & 1);
        sfence();
        dbreg[kbd_device_slot_id] = dci_kbd;
    }

    fb_printf("usbkbd: init OK (report=%d)\n", kbd_report_size);
    return 1;
}

/* ══════════════════════════════════════════════
   USB Mouse enumeration (HID Boot Protocol)
   ══════════════════════════════════════════════ */

static int xhci_enumerate_mouse(void)
{
    mouse_report_size = MOUSE_REPORT_SIZE;

    struct usb_desc_device dev_desc;
    memset(&dev_desc, 0, sizeof(dev_desc));
    int r2 = usb_get_device_descriptor(&dev_desc, 8);
    if(r2 < 0 || dev_desc.bMaxPacketSize0 == 0) return 0;

    int new_mps = dev_desc.bMaxPacketSize0;
    if(new_mps != 8 && new_mps != 16 && new_mps != 32 && new_mps != 64)
        new_mps = 8;
    ep0_mps = new_mps;
    DEV_EP(mouse_dev_ctx_buf, 0)->word0 = ep0_mps;

    int r = cmd_address_device(device_slot_id, 0);
    if(r < 0) return 0;
    reset_ep0_ring();

    memset(&dev_desc, 0, sizeof(dev_desc));
    r2 = usb_get_device_descriptor(&dev_desc, 18);
    if(r2 < 0) return 0;

    int is_mouse = 0;
    int mouse_ep_mps = 8;

    {
        uint8_t cfg_buf[256];
        memset(cfg_buf, 0, sizeof(cfg_buf));
        r2 = usb_get_config_descriptor_raw(cfg_buf, 9);
        if(r2 < 0) return 0;

        uint16_t total_len = cfg_buf[2] | ((uint16_t)cfg_buf[3] << 8);
        if(total_len > 256) total_len = 256;
        r2 = usb_get_config_descriptor_raw(cfg_buf, total_len);
        if(r2 < 0) return 0;

        int cur_class = 0, cur_proto = 0;
        int offset = 0;
        while(offset + 1 < total_len) {
            int dlen = cfg_buf[offset];
            int dtype = cfg_buf[offset + 1];
            if(dlen < 2) break;
            if(dtype == USB_DESC_INTERFACE) {
                struct usb_desc_interface *iface =
                    (struct usb_desc_interface*)(cfg_buf + offset);
                cur_class = iface->bInterfaceClass;
                cur_proto = iface->bInterfaceProtocol;
            }
            if(dtype == USB_DESC_ENDPOINT && cur_class == USB_CLASS_HID &&
               cur_proto == USB_PROTOCOL_MOUSE) {
                struct usb_desc_endpoint *ep =
                    (struct usb_desc_endpoint*)(cfg_buf + offset);
                if((ep->bmAttributes & 0x03) == 3 &&
                   (ep->bEndpointAddress & 0x80)) {
                    mouse_ep_addr = ep->bEndpointAddress;
                    mouse_ep_interval = ep->bInterval;
                    mouse_ep_mps = ep->wMaxPacketSize;
                    is_mouse = 1;
                }
            }
            offset += dlen;
        }
    }

    if(!is_mouse) return 0;

    mouse_report_size = mouse_ep_mps;
    if(mouse_report_size < 4) mouse_report_size = 4;
    if(mouse_report_size > MOUSE_REPORT_SIZE) mouse_report_size = MOUSE_REPORT_SIZE;

    {
        int r3 = usb_set_configuration(1);
        if(r3 < 0) return 0;
    }

    /* SET_IDLE(0) for mouse */
    {
        struct usb_setup_pkt sp = {
            .bmReqType = 0x21,
            .bRequest = USB_REQ_SET_IDLE,
            .wValue = 0,
            .wIndex = 0,
            .wLength = 0,
        };
        ep0_do_control(&sp, NULL);
    }

    {
        int ep_num = mouse_ep_addr & 0x0F;
        int is_in = (mouse_ep_addr & 0x80) ? 1 : 0;
        int dci = ep_num * 2 + is_in;
        int eps_idx = dci - 1;
        mouse_eps_idx = eps_idx;

        if(dci < 2 || dci > 31 || eps_idx < 0 || eps_idx > 30)
            return 0;

        memset(in_ctx_buf, 0, ctx_stride * 33);
        IN_CTRL(in_ctx_buf)[0] = 0;
        IN_CTRL(in_ctx_buf)[1] = (1 << 0) | (1 << dci);

        IN_SLOT(in_ctx_buf)->word0 = (device_speed << 20) | ((eps_idx + 1) << 27);
        IN_SLOT(in_ctx_buf)->word1 = (device_port << 16);

        IN_EP(in_ctx_buf, 0)->word0 = ep0_mps;
        IN_EP(in_ctx_buf, 0)->word1 = (4 << 3) | (3 << 6);
        IN_EP(in_ctx_buf, 0)->dequeue = v2p(ep0_ring) | 1;

        int ep_type = is_in ? 7 : 3;
        IN_EP(in_ctx_buf, eps_idx)->word0 = mouse_report_size;
        IN_EP(in_ctx_buf, eps_idx)->word1 = (ep_type << 3) | (3 << 6);
        IN_EP(in_ctx_buf, eps_idx)->dequeue = v2p(mouse_ep_ring) | 1;

        write_trb(cmd_ring, 0, v2p(in_ctx_buf), 0,
                  TRB_CONFIG_EP | (cmd_pcs) | TRB_IOC | ((uint32_t)device_slot_id << 24));
        write_link_trb(cmd_ring, 2, cmd_pcs);
        cmd_pcs ^= 1;
        cmd_ring_doorbell();
        r = wait_for_event(EVT_CMD_COMPLETE);
        if(r < 0) return 0;
    }

    memset(mouse_report_buf, 0, mouse_report_size);
    memset(prev_mouse_report, 0, mouse_report_size);
    memset(last_mouse_processed, 0, mouse_report_size);

    memset((void*)mouse_ep_ring, 0, sizeof(mouse_ep_ring));
    write_link_trb(mouse_ep_ring, MOUSE_EP_RING_SIZE, 1);
    mouse_ring_pcs = 1;
    mouse_ring_write_pos = 1;

    write_trb(mouse_ep_ring, 0, v2p(mouse_report_buf), mouse_report_size,
              TRB_NORMAL | (mouse_ring_pcs) | TRB_IOC);
    mouse_ring_active = 1;
    mouse_slot_id = device_slot_id;

    {
        int dci_mouse = (mouse_ep_addr & 0x0F) * 2 + ((mouse_ep_addr >> 7) & 1);
        sfence();
        dbreg[mouse_slot_id] = dci_mouse;
    }

    fb_printf("usbmouse: init OK (report=%d)\n", mouse_report_size);
    return 1;
}

int usb_kbd_init(void)
{
    extern char __text_start[];
    kernel_phys_base = (uint64_t)(uintptr_t)__text_start & ~0x1FFFFFULL;

    xhci_initialized = 0;
    kbd_found = 0;
    device_slot_id = 0;
    report_ready = 0;

    uint64_t xhci_base = pci_find_xhci();
    if(!xhci_base) {
        fb_puts("usbkbd: no xHCI controller found\n");
        return 0;
    }

    paging_map_mmio_range(xhci_base, 0x10000);
    uint64_t xhci_virt = xhci_base + 0xFFFF800000000000ULL;

    caps = (volatile struct xhci_caps*)(uintptr_t)xhci_virt;
    uint8_t caplen = caps->caplength;
    ops  = (volatile struct xhci_ops*)(xhci_virt + caplen);
    uint32_t dboff_val = caps->dboff;
    uint32_t rtsoff_val = caps->rtsoff;
    dbreg = (volatile uint32_t*)(uintptr_t)(xhci_virt + dboff_val);
    ir  = (volatile struct xhci_intr*)(uintptr_t)(xhci_virt + rtsoff_val + 0x20);

    uint32_t hsp1 = caps->hcsparams1;
    xhci_num_ports = (hsp1 >> 16) & 0xFF;
    if(!xhci_num_ports) xhci_num_ports = (hsp1 >> 24) & 0xFF;
    if(!xhci_num_ports) xhci_num_ports = 1;
    if(xhci_num_ports > MAX_PORTS) xhci_num_ports = MAX_PORTS;

    {
        uint32_t hccp1 = caps->hccparams1;
        uint16_t xecp = (hccp1 >> 16) & 0xFFFF;
        while(xecp) {
            uint32_t cap = *(volatile uint32_t*)(uintptr_t)(xhci_virt + (xecp << 2));
            if((cap & 0xFF) == 1) {
                if(cap & (1U << 16)) {
                    volatile uint32_t *leg = (volatile uint32_t*)(uintptr_t)(xhci_virt + (xecp << 2));
                    *leg = cap | (1U << 24);
                    uint64_t dead = rdtsc() + 20000000000ULL;
                    while(rdtsc() < dead) {
                        if(!(*leg & (1U << 16))) break;
                    }
                    dbg_printf("[xhci] legacy handoff done\n");
                }
                break;
            }
            xecp = (cap >> 8) & 0xFF;
        }
    }

    {
        uint16_t ver = caps->hciversion;
        uint32_t hsp1 = caps->hcsparams1;
        uint32_t hccp1 = caps->hccparams1;
        ctx_stride = (hccp1 & 4) ? 64 : 32;
        dbg_printf("HCIVERSION=%u.%u HCSPARAMS1=0x%08x HCCPARAMS1=0x%08x ctx_stride=%d\n",
            ver >> 8, ver & 0xFF, hsp1, hccp1, ctx_stride);
    }

    xhci_hw_init();

    for(int probe = 1; probe <= xhci_num_ports; probe++) {
        uint32_t ps = ops->portregs[probe - 1].portsc;
        if(!(ps & PORTSC_CCS)) continue;

        int ok = xhci_init_port(probe);
        if(!ok) continue;

        int r = xhci_enumerate_kbd();
        if(r) {
            xhci_initialized = 1;
            kbd_found = 1;
            kbd_port = probe;
            return 1;
        }

        int sid = device_slot_id;
        if(sid > 0) {
            cmd_disable_slot(sid);
            memset(dev_ctx_buf, 0, sizeof(dev_ctx_buf));
            device_slot_id = 0;
            reset_ep0_ring();
        }
    }

    return 0;
}

static void rearm_kbd(void)
{
    if(!kbd_ring_active) return;
    int idx;
    if(kbd_ring_write_pos >= RING_SIZE - 1) {
        write_link_trb(kbd_ring, RING_SIZE, kbd_ring_pcs);
        kbd_ring_pcs ^= 1;
        kbd_ring_write_pos = 1;
        idx = 0;
    } else {
        idx = kbd_ring_write_pos;
        kbd_ring_write_pos++;
        if(kbd_ring_write_pos == RING_SIZE - 1)
            write_link_trb(kbd_ring, RING_SIZE, kbd_ring_pcs);
    }
    write_trb(kbd_ring, idx, v2p(kbd_report_buf), kbd_report_size,
              TRB_NORMAL | (kbd_ring_pcs) | TRB_IOC);

    int dci_kbd = (kbd_ep_addr & 0x0F) * 2 + ((kbd_ep_addr >> 7) & 1);
    sfence();
    dbreg[kbd_device_slot_id] = dci_kbd;
}

static int usb_poll_interrupt(void)
{
    if(!evt_available()) return 0;

    volatile struct xhci_trb *ev = &evt_ring[evt_deq_idx];
    int etype = evt_get_trb_type(ev);
    int was_kbd = 0;

    if(etype == EVT_PORT_STATUS) {
        int port_id = (int)evt_get_port_id(ev);
        if(port_id >= 1 && port_id <= xhci_num_ports) {
            uint32_t ps = ops->portregs[port_id - 1].portsc;
            uint32_t chg = ps & (PORTSC_CSC | PORTSC_PRC | PORTSC_WRC);
            if(chg)
                ops->portregs[port_id - 1].portsc = ps | chg;
        }
        evt_advance();
        return 0;
    }

    if(etype == EVT_TRANSFER) {
        uint32_t slot = evt_get_slot_id(ev);
        int comp = evt_get_completion(ev);
        if(slot == (uint32_t)kbd_device_slot_id && comp == CMP_SUCCESS) {
            memcpy(prev_kbd_report, kbd_report_buf, kbd_report_size);
            report_ready = 1;
            was_kbd = 1;
            rearm_kbd();
        }
        /* Consume only our events; leave mouse events for mouse_poll_interrupt */
        if(slot == (uint32_t)kbd_device_slot_id)
            evt_advance();
        return was_kbd;
    }

    /* Unknown event type; leave in ring for mouse handler */
    return 0;
}

static const char usb_kbd_ascii[256][2] = {
    [0x00] = {0,0}, [0x01] = {0,0}, [0x02] = {0,0}, [0x03] = {0,0},
    [0x04] = {'a','A'}, [0x05] = {'b','B'}, [0x06] = {'c','C'}, [0x07] = {'d','D'},
    [0x08] = {'e','E'}, [0x09] = {'f','F'}, [0x0A] = {'g','G'}, [0x0B] = {'h','H'},
    [0x0C] = {'i','I'}, [0x0D] = {'j','J'}, [0x0E] = {'k','K'}, [0x0F] = {'l','L'},
    [0x10] = {'m','M'}, [0x11] = {'n','N'}, [0x12] = {'o','O'}, [0x13] = {'p','P'},
    [0x14] = {'q','Q'}, [0x15] = {'r','R'}, [0x16] = {'s','S'}, [0x17] = {'t','T'},
    [0x18] = {'u','U'}, [0x19] = {'v','V'}, [0x1A] = {'w','W'}, [0x1B] = {'x','X'},
    [0x1C] = {'y','Y'}, [0x1D] = {'z','Z'},
    [0x1E] = {'1','!'}, [0x1F] = {'2','@'}, [0x20] = {'3','#'}, [0x21] = {'4','$'},
    [0x22] = {'5','%'}, [0x23] = {'6','^'}, [0x24] = {'7','&'}, [0x25] = {'8','*'},
    [0x26] = {'9','('}, [0x27] = {'0',')'},
    [0x28] = {'\n','\n'},
    [0x29] = {0,0},
    [0x2A] = {'\b','\b'},
    [0x2B] = {'\t','\t'},
    [0x2C] = {' ',' '},
    [0x2D] = {'-','_'}, [0x2E] = {'=','+'}, [0x2F] = {'[','{'}, [0x30] = {']','}'},
    [0x31] = {'\\','|'}, [0x32] = {0,0},
    [0x33] = {';',':'}, [0x34] = {'\'','"'}, [0x35] = {'`','~'},
    [0x36] = {',','<'}, [0x37] = {'.','>'}, [0x38] = {'/','?'},
    [0x39] = {0,0},
    [0x4C] = {0,0},
    [0x4F] = {0,0},
    [0x50] = {0,0},
    [0x51] = {0,0},
    [0x52] = {0,0},
};

static const char usb_kbd_ru[256][2] = {
    [0x04] = {0xF4,0xD4}, [0x05] = {0xE8,0xC8}, [0x06] = {0xF1,0xD1}, [0x07] = {0xE2,0xC2},
    [0x08] = {0xF3,0xD3}, [0x09] = {0xE0,0xC0}, [0x0A] = {0xEF,0xCF}, [0x0B] = {0xF0,0xD0},
    [0x0C] = {0xF8,0xD8}, [0x0D] = {0xEE,0xCE}, [0x0E] = {0xEB,0xCB}, [0x0F] = {0xE4,0xC4},
    [0x10] = {0xEC,0xCC}, [0x11] = {0xF2,0xD2}, [0x12] = {0xF9,0xD9}, [0x13] = {0xE7,0xC7},
    [0x14] = {0xE9,0xC9}, [0x15] = {0xEA,0xCA}, [0x16] = {0xFB,0xDB}, [0x17] = {0xE5,0xC5},
    [0x18] = {0xE3,0xC3}, [0x19] = {0xEC,0xCC}, [0x1A] = {0xF6,0xD6}, [0x1B] = {0xF7,0xD7},
    [0x1C] = {0xED,0xCD}, [0x1D] = {0xFF,0xDF},
    [0x1E] = {'1','!'}, [0x1F] = {'2','@'}, [0x20] = {'3','#'}, [0x21] = {'4','$'},
    [0x22] = {'5','%'}, [0x23] = {'6','^'}, [0x24] = {'7','&'}, [0x25] = {'8','*'},
    [0x26] = {'9','('}, [0x27] = {'0',')'},
    [0x28] = {'\n','\n'},
    [0x29] = {0,0},
    [0x2A] = {'\b','\b'},
    [0x2B] = {'\t','\t'},
    [0x2C] = {' ',' '},
    [0x2D] = {'-','_'}, [0x2E] = {'=','+'},
    [0x2F] = {0xF5,0xD5}, [0x30] = {0xFA,0xDA}, [0x31] = {'\\','|'},
    [0x32] = {0,0},
    [0x33] = {0xE6,0xC6}, [0x34] = {0xFD,0xDD}, [0x35] = {0xB8,0xA8},
    [0x36] = {0xE1,0xC1}, [0x37] = {0xFE,0xDE}, [0x38] = {0x2E,0x2C},
};

static int kbd_key_is_pressed(const uint8_t *report, int report_size, uint8_t key)
{
    if(report_size < 2) return 0;
    if(key == 0) return 0;
    for(int i = 2; i < report_size; i++) {
        if(report[i] == key) return 1;
    }
    return 0;
}

static uint8_t kbd_first_key(const uint8_t *report, int report_size)
{
    if(report_size < 2) return 0;
    for(int i = 2; i < report_size; i++) {
        if(report[i] >= 4 && report[i] <= 0xE7)
            return report[i];
    }
    return 0;
}

static char usb_kbd_decode(void)
{
    if(!report_ready) return 0;

    uint8_t *cur = prev_kbd_report;
    uint8_t *prev = last_processed_report;
    int sz = kbd_report_size;

    uint8_t cur_mod = cur[0];
    uint8_t prev_mod = prev[0];
    int shift = (cur_mod & 0x22) != 0;
    int alt = (cur_mod & 0x44) != 0;
    int alt_was = (prev_mod & 0x44) != 0;
    int shift_was = (prev_mod & 0x22) != 0;

    if(shift && alt && (!shift_was || !alt_was)) {
        kbd_layout ^= 1;
        memcpy(last_processed_report, cur, sz);
        report_ready = 0;
        return 0;
    }

    uint8_t cur_key = kbd_first_key(cur, sz);
    uint8_t prev_key = kbd_first_key(prev, sz);

    int is_new = 0;
    if(cur_key && cur_key != prev_key && !kbd_key_is_pressed(prev, sz, cur_key))
        is_new = 1;
    if(cur_key && cur_key == prev_key && cur_mod != prev_mod)
        is_new = 1;

    memcpy(last_processed_report, cur, sz);
    report_ready = 0;

    if(!is_new) return 0;

    const char (*tbl)[2] = kbd_layout ? usb_kbd_ru : usb_kbd_ascii;
    if(tbl[cur_key][0])
        return tbl[cur_key][shift ? 1 : 0];

    return 0;
}

char usb_kbd_getchar(void)
{
    if(!xhci_initialized || !kbd_found) return 0;

    usb_poll_interrupt();

    if(report_ready) {
        char c = usb_kbd_decode();
        return c;
    }
    return 0;
}

void usb_kbd_scan(void)
{
    if(!xhci_initialized) {
        usb_kbd_init();
        return;
    }

    kbd_found = 0;
    report_ready = 0;

    for(int probe = 1; probe <= xhci_num_ports; probe++) {
        uint32_t ps = ops->portregs[probe - 1].portsc;
        if(!(ps & PORTSC_CCS)) continue;

        int ok = xhci_init_port(probe);
        if(!ok) continue;

        int r = xhci_enumerate_kbd();
        if(r) {
            kbd_found = 1;
            return;
        }

        int sid = device_slot_id;
        if(sid > 0) {
            cmd_disable_slot(sid);
            memset(dev_ctx_buf, 0, sizeof(dev_ctx_buf));
            device_slot_id = 0;
            reset_ep0_ring();
        }
    }
}

/* ══════════════════════════════════════════════
   Mouse: poll + decode
   ══════════════════════════════════════════════ */

static void rearm_mouse(void)
{
    if(!mouse_ring_active) return;
    int idx;
    if(mouse_ring_write_pos >= MOUSE_EP_RING_SIZE - 1) {
        write_link_trb(mouse_ep_ring, MOUSE_EP_RING_SIZE, mouse_ring_pcs);
        mouse_ring_pcs ^= 1;
        mouse_ring_write_pos = 1;
        idx = 0;
    } else {
        idx = mouse_ring_write_pos;
        mouse_ring_write_pos++;
        if(mouse_ring_write_pos == MOUSE_EP_RING_SIZE - 1)
            write_link_trb(mouse_ep_ring, MOUSE_EP_RING_SIZE, mouse_ring_pcs);
    }
    write_trb(mouse_ep_ring, idx, v2p(mouse_report_buf), mouse_report_size,
              TRB_NORMAL | (mouse_ring_pcs) | TRB_IOC);

    int dci_mouse = (mouse_ep_addr & 0x0F) * 2 + ((mouse_ep_addr >> 7) & 1);
    sfence();
    dbreg[mouse_slot_id] = dci_mouse;
}

static int mouse_poll_interrupt(void)
{
    if(!evt_available()) return 0;

    volatile struct xhci_trb *ev = &evt_ring[evt_deq_idx];
    int etype = evt_get_trb_type(ev);
    int was_mouse = 0;

    if(etype == EVT_PORT_STATUS) {
        int port_id = (int)evt_get_port_id(ev);
        if(port_id >= 1 && port_id <= xhci_num_ports) {
            uint32_t ps = ops->portregs[port_id - 1].portsc;
            uint32_t chg = ps & (PORTSC_CSC | PORTSC_PRC | PORTSC_WRC);
            if(chg)
                ops->portregs[port_id - 1].portsc = ps | chg;
        }
        evt_advance();
        return 0;
    }

    if(etype == EVT_TRANSFER) {
        uint32_t slot = evt_get_slot_id(ev);
        int comp = evt_get_completion(ev);
        if(slot == (uint32_t)mouse_slot_id && comp == CMP_SUCCESS) {
            memcpy(prev_mouse_report, mouse_report_buf, mouse_report_size);
            mouse_report_ready = 1;
            was_mouse = 1;
            rearm_mouse();
        }
        if(slot == (uint32_t)mouse_slot_id)
            evt_advance();
        return was_mouse;
    }

    return 0;
}

int usb_mouse_get_event(int *buttons, int *dx, int *dy, int *wheel)
{
    if(!xhci_initialized || !mouse_found) return 0;

    mouse_poll_interrupt();

    if(!mouse_report_ready) return 0;

    uint8_t *cur = prev_mouse_report;
    uint8_t *prev = last_mouse_processed;

    int old_buttons = prev[0] & 0x07;
    int new_buttons = cur[0] & 0x07;

    int8_t raw_dx = (int8_t)cur[1];
    int8_t raw_dy = (int8_t)cur[2];
    int8_t raw_wheel = (mouse_report_size >= 4) ? -(int8_t)cur[3] : 0;

    int moved = (raw_dx != 0 || raw_dy != 0);
    int clicked = (new_buttons != old_buttons);
    int scrolled = (raw_wheel != 0);

    memcpy(last_mouse_processed, cur, mouse_report_size);
    mouse_report_ready = 0;

    if(!moved && !clicked && !scrolled) return 0;

    if(buttons) *buttons = new_buttons;
    if(dx) *dx = raw_dx;
    if(dy) *dy = raw_dy;
    if(wheel) *wheel = raw_wheel;

    return 1;
}

/* ══════════════════════════════════════════════
   Mouse: scan ports for mouse after keyboard
   ══════════════════════════════════════════════ */

void usb_mouse_init(void)
{
    if(!xhci_initialized) {
        dbg_printf("[mouse] xHCI not initialized\n");
        return;
    }
    if(mouse_found) return;

    int saved_slot = device_slot_id;
    int saved_port = device_port;

    for(int probe = 1; probe <= xhci_num_ports; probe++) {
        if(probe == kbd_port) continue; /* keyboard already occupies this port */
        uint32_t ps = ops->portregs[probe - 1].portsc;
        dbg_printf("[mouse] probe port %d: sc=0x%08x\n", probe, ps);
        if(!(ps & PORTSC_CCS)) continue;

        int ok = xhci_init_port_ctx(probe, mouse_dev_ctx_buf);
        if(!ok) {
            dbg_printf("[mouse] xhci_init_port_ctx(%d) FAILED\n", probe);
            continue;
        }
        dbg_printf("[mouse] xhci_init_port_ctx(%d) OK, slot=%d\n", probe, device_slot_id);

        int r = xhci_enumerate_mouse();
        dbg_printf("[mouse] xhci_enumerate_mouse()=%d, slot=%d\n", r, device_slot_id);
        if(r) {
            mouse_found = 1;
            mouse_slot_id = device_slot_id;
            mouse_port = probe;
            dbg_printf("[mouse] FOUND on port %d, slot %d, ep=0x%02x, size=%d\n",
                       probe, mouse_slot_id, mouse_ep_addr, mouse_report_size);
            return;
        }

        int sid = device_slot_id;
        if(sid > 0) {
            cmd_disable_slot(sid);
            memset(mouse_dev_ctx_buf, 0, sizeof(mouse_dev_ctx_buf));
            device_slot_id = 0;
            reset_ep0_ring();
        }
    }

    device_slot_id = saved_slot;
    device_port = saved_port;
    dbg_printf("[mouse] NOT FOUND\n");
}
