// ---------------------------------------------------------------------------
// USB host controller driver — xHCI (USB 2.0/3.x), HID mouse + keyboard.
//
// **Polled** driver: the event ring is drained from the cursor task loop (and
// the WM input pump for keyboards) instead of wiring MSI/INTx into the PIC.
//
// Scope:
//   • controller reset, command + event rings, scratchpad buffers (§5.4.5)
//   • one slot per HID device (up to MAX_XDEV), address/configure sequence
//   • boot-protocol HID mouse and keyboard interfaces
//   • reports decoded into ps2::MouseEvent / ps2::feed_scancode
//
// DMA: rings/contexts come from the page allocator (physical == kernel
// address via the identity map), never the heap, and are 64-byte aligned as
// the xHCI requires.  The controller is 64-bit capable, so no 4 GiB bound.
//
// Register/TRB/context layouts follow the Linux xHCI driver definitions and
// the xHCI 1.2 spec (TRB type bits 10-15, completion code bits 24-31 of event
// dword2, slot id bits 24-31 / ep id bits 16-20 of event dword3).
// ---------------------------------------------------------------------------
#![allow(non_snake_case)] // macro-style helpers mirror the Linux xHCI headers
use core::sync::atomic::{AtomicBool, AtomicU32, Ordering};

use crate::ps2::MouseEvent;

// ---- capability registers (offset from mmio base) ----
const CAP_CAPLENGTH: usize = 0x00;
const HCSPARAMS1: usize = 0x04; // MaxSlots[7:0] MaxIntrs[15:8] MaxPorts[31:24]
const HCSPARAMS2: usize = 0x08; // scratchpad buffers, ERST max
const HCCPARAMS1: usize = 0x10; // bit0=64-bit addressing; bit2=64-byte contexts
const DBOFF: usize = 0x14;
const RTSOFF: usize = 0x18;

// ---- operational registers (relative to op = mmio + CAPLENGTH) ----
const OP_USBCMD: usize = 0x00;
const OP_USBSTS: usize = 0x04;
const OP_CRCR: usize = 0x18; // 64-bit
const OP_DCBAAP: usize = 0x30; // 64-bit
const OP_CONFIG: usize = 0x38;
const OP_PORTSC: usize = 0x400;
const PORT_STRIDE: usize = 0x10;

// ---- runtime registers (relative to mmio + RTSOFF), interrupter 0 ----
const RT_IMAN: usize = 0x20;
const RT_IMOD: usize = 0x24;
const RT_ERSTSZ: usize = 0x28;
const RT_ERSTBA: usize = 0x30; // 64-bit
const RT_ERDP: usize = 0x38; // 64-bit

const CMD_RUN: u32 = 1 << 0;
const CMD_RESET: u32 = 1 << 1;
const STS_HALT: u32 = 1 << 0;
const STS_CNR: u32 = 1 << 11;

const CMD_RING_CYCLE: u32 = 1 << 0;

// ---- TRB control-word bit layout ----
const TRB_CYCLE: u32 = 1 << 0;
const TRB_IOC: u32 = 1 << 5;
const TRB_IDT: u32 = 1 << 6;
const TRB_TYPE_BITMASK: u32 = 0xfc00;
fn TYPE(t: u32) -> u32 {
    t << 10
}
fn SLOT_ID(p: u32) -> u32 {
    (p & 0xff) << 24
}
fn TRB_LEN(p: u32) -> u32 {
    p & 0x1fffff
}
const TRB_DIR_IN: u32 = 1 << 16;
fn TRB_TX_TYPE(p: u32) -> u32 {
    p << 16
}
const TRB_DATA_OUT: u32 = 2;
const TRB_DATA_IN: u32 = 3;

const TRB_NORMAL: u32 = 1;
const TRB_SETUP: u32 = 2;
const TRB_DATA: u32 = 3;
const TRB_STATUS: u32 = 4;
const TRB_ENABLE_SLOT: u32 = 9;
const TRB_ADDR_DEV: u32 = 11;
const TRB_CONFIG_EP: u32 = 12;
const TRB_EV_TRANSFER: u32 = 0x20;
const TRB_EV_CMD_COMPLETE: u32 = 0x21;

const COMP_SUCCESS: u32 = 1;
const COMP_SHORT: u32 = 13;

fn GET_COMP_CODE(d2: u32) -> u32 {
    (d2 >> 24) & 0xff
}
fn TRB_TO_SLOT_ID(d3: u32) -> u32 {
    (d3 >> 24) & 0xff
}
fn TRB_TO_EP_ID(d3: u32) -> u32 {
    (d3 >> 16) & 0x1f
}
fn TRB_FIELD_TO_TYPE(d3: u32) -> u32 {
    (d3 & TRB_TYPE_BITMASK) >> 10
}

// ---- context fields ----
fn LAST_CTX(p: u32) -> u32 {
    p << 27
}
fn ROOTPORT(p: u32) -> u32 {
    (p & 0xff) << 16
}
fn SLOT_SPEED_X(p: u32) -> u32 {
    (p & 0xf) << 20
}
fn EP_TYPE(p: u32) -> u32 {
    p << 3
}
fn MAX_PACKET(p: u32) -> u32 {
    (p & 0xffff) << 16
}
fn ERROR_COUNT(p: u32) -> u32 {
    (p & 0x3) << 1
}
fn EP_INTERVAL(p: u32) -> u32 {
    (p & 0xff) << 16
}
fn EP_AVG(p: u32) -> u32 {
    p & 0xffff
}
const CTRL_EP: u32 = 4;
const INT_IN_EP: u32 = 7;

// ---- PORTSC ----
const PORT_CONNECT: u32 = 1 << 0;
const PORT_PE: u32 = 1 << 1;
const PORT_RESET: u32 = 1 << 4;
const PORT_POWER: u32 = 1 << 9;
fn DEV_SPEED(p: u32) -> u32 {
    (p >> 10) & 0xf
}

// ---- USB meta protocol ----
const USB_REQ_GET_DESCRIPTOR: u8 = 0x06;
const USB_REQ_SET_CONFIGURATION: u8 = 0x09;
const DT_DEVICE: u8 = 0x01;
const DT_CONFIG: u8 = 0x02;
const DT_INTERFACE: u8 = 0x04;
const DT_ENDPOINT: u8 = 0x05;
const HID_CLASS: u8 = 0x03;
const HID_SUBCLASS_BOOT: u8 = 0x01;
const HID_PROTO_KEYBOARD: u8 = 1;
const HID_PROTO_MOUSE: u8 = 2;
const REQ_SET_PROTOCOL: u8 = 0x0B;
const REQ_SET_IDLE: u8 = 0x0A;

const MAX_XDEV: usize = 4;

#[derive(Clone, Copy, PartialEq)]
enum Kind {
    Mouse,
    Keyboard,
}

// ---------------------------------------------------------------------------
// State (cooperative single-task kernel → unsafe statics)
// ---------------------------------------------------------------------------
static mut M: usize = 0; // mmio base
static mut OPB: usize = 0; // operational base
static mut DBB: usize = 0; // doorbell base
static mut RTB: usize = 0; // runtime base
static mut CTRL_CS: usize = 32; // context size
static mut MAXPORTS: u32 = 0;
static mut N_SLOTS: u32 = 1; // MaxSlotsEn we program

#[derive(Clone, Copy)]
struct Ring {
    raw: *mut u32,
    count: usize,
    idx: usize,
    cyc: u32,
}
const fn ri_empty(count: usize) -> Ring {
    Ring {
        raw: core::ptr::null_mut(),
        count,
        idx: 0,
        cyc: 1,
    }
}
static mut CMD: Ring = ri_empty(64);
static mut EV: Ring = ri_empty(32);

static mut DCBAA: *mut u64 = core::ptr::null_mut();
static mut INCTX: *mut u8 = core::ptr::null_mut();
static mut ERST: *mut u8 = core::ptr::null_mut();
static mut SPB_ARRAY: *mut u64 = core::ptr::null_mut();
static mut NUM_SPB: u32 = 0;
/// Report buffers (16 bytes per device) in one DMA page.
static mut REPORTS: *mut u8 = core::ptr::null_mut();

/// One enumerated HID device (slot + interrupt-IN endpoint).
#[derive(Clone, Copy)]
struct XDev {
    slot: u32,
    dci: u32, // endpoint context index of the IN endpoint
    ep0: Ring,
    epin: Ring,
    devctx: *mut u8,
    kind: Kind,
    len: u32,
    kbd: crate::hid::Keyboard,
    armed: bool,
    ready: bool,
    code: u32,
    reports: u32,
    failures: u32,
}
impl XDev {
    const EMPTY: XDev = XDev {
        slot: 0,
        dci: 0,
        ep0: ri_empty(8),
        epin: ri_empty(8),
        devctx: core::ptr::null_mut(),
        kind: Kind::Mouse,
        len: 0,
        kbd: crate::hid::Keyboard::new(),
        armed: false,
        ready: false,
        code: 0,
        reports: 0,
        failures: 0,
    };
}
static mut XDEVS: [XDev; MAX_XDEV] = [XDev::EMPTY; MAX_XDEV];
static NXDEV: AtomicU32 = AtomicU32::new(0);

static INITOK: AtomicBool = AtomicBool::new(false);
static DISCOVERED_TRIED: AtomicU32 = AtomicU32::new(0);
static LAST_SCAN: AtomicU32 = AtomicU32::new(0);

unsafe fn dev_mut(i: usize) -> &'static mut XDev {
    unsafe { &mut *core::ptr::addr_of_mut!(XDEVS[i]) }
}

fn cs() -> usize {
    unsafe { CTRL_CS }
}

// ---------------------------------------------------------------------------
// MMIO accessors
// ---------------------------------------------------------------------------
unsafe fn rd(a: usize) -> u32 {
    unsafe { (a as *const u32).read_volatile() }
}
unsafe fn wr(a: usize, v: u32) {
    unsafe { (a as *mut u32).write_volatile(v) }
}
unsafe fn wr64(a: usize, v: u64) {
    unsafe {
        (a as *mut u32).write_volatile(v as u32);
        (a as *mut u32).add(1).write_volatile((v >> 32) as u32);
        let _ = rd(a);
    }
}
fn op(off: usize) -> usize {
    unsafe { OPB + off }
}

/// 16-bit I/O read (UHCI port status/control registers live in I/O space).
unsafe fn io_iow16(port: u16) -> u16 {
    let v: u16;
    unsafe {
        core::arch::asm!("in ax, dx", out("ax") v, in("dx") port, options(nomem, nostack, preserves_flags));
    }
    v
}
fn db(slot: u32) -> usize {
    unsafe { DBB }.wrapping_add(slot as usize * 4)
}
fn rt(off: usize) -> usize {
    unsafe { RTB + off }
}

/// DMA page(s) from the page allocator (physical == CPU address).  Always
/// 4 KiB aligned, which satisfies every xHCI ring/context alignment rule.
fn dma_alloc(bytes: usize, _align: usize) -> *mut u8 {
    let pages = bytes.div_ceil(4096).max(1);
    match crate::memory::alloc_pages_above(pages, 0x100_0000) {
        Some(p) => {
            let ptr = p as *mut u8;
            unsafe { core::ptr::write_bytes(ptr, 0, pages * 4096) };
            ptr
        }
        None => core::ptr::null_mut(),
    }
}

// ---------------------------------------------------------------------------
// Ring primitives
// ---------------------------------------------------------------------------
unsafe fn trb_push(r: &mut Ring, d0: u32, d1: u32, d2: u32, d3: u32) {
    unsafe {
        let s = r.raw.add((r.idx % r.count) * 4);
        *s.add(0) = d0;
        *s.add(1) = d1;
        *s.add(2) = d2;
        *s.add(3) = (d3 & !TRB_CYCLE) | r.cyc;
        r.idx += 1;
        if r.idx.is_multiple_of(r.count) {
            r.cyc ^= 1;
        }
    }
}

/// Drain produced events, calling `handle(d0,d1,d2,d3)` per event, then
/// advance ERDP so the controller can reuse the segment entries.
unsafe fn drain_events(handle: &mut dyn FnMut(u32, u32, u32, u32)) {
    unsafe {
        let ev = &mut *core::ptr::addr_of_mut!(EV);
        let mut consumed = 0usize;
        loop {
            let idx = ev.idx % ev.count;
            let s = ev.raw.add(idx * 4);
            let d3 = s.add(3).read();
            if (d3 & 1) != ev.cyc {
                break;
            }
            let (d0, d1, d2) = (s.add(0).read(), s.add(1).read(), s.add(2).read());
            handle(d0, d1, d2, d3);
            ev.idx += 1;
            consumed += 1;
            if ev.idx.is_multiple_of(ev.count) {
                ev.cyc ^= 1;
            }
        }
        if consumed > 0 {
            let idx = ev.idx % ev.count;
            wr64(rt(RT_ERDP), ev.raw.add(idx * 4) as u64);
        }
    }
}

/// Wait for a command-completion event; returns (code, slot).
fn wait_cmd() -> (u32, u32) {
    let mut code = 0u32;
    let mut slot = 0u32;
    let mut done = false;
    let mut seen = 0u32;
    for _ in 0..4_000_000 {
        if done {
            break;
        }
        unsafe {
            drain_events(&mut |_d0, _d1, d2, d3| {
                seen += 1;
                if TRB_FIELD_TO_TYPE(d3) == TRB_EV_CMD_COMPLETE {
                    code = GET_COMP_CODE(d2);
                    slot = TRB_TO_SLOT_ID(d3);
                    done = true;
                }
            });
        }
        core::hint::spin_loop();
    }
    if code != COMP_SUCCESS {
        crate::serial_println!(
            "usb: cmd-complete code={} slot={} seen={}",
            code,
            slot,
            seen
        );
    }
    (code, slot)
}

fn issue_cmd(d3: u32) -> (u32, u32) {
    issue_cmd_ptr(d3, 0)
}

fn issue_cmd_ptr(d3: u32, ptr: usize) -> (u32, u32) {
    unsafe {
        trb_push(
            &mut *core::ptr::addr_of_mut!(CMD),
            ptr as u32,
            (ptr as u64 >> 32) as u32,
            0,
            d3,
        );
        wr(db(0), 0);
        let _ = rd(db(0));
    }
    wait_cmd()
}

/// Record a transfer event against the device it belongs to.  Called for
/// *every* transfer event, including ones drained while waiting for a
/// control transfer: dropping them would lose an armed device's report and
/// stall that device forever.
fn route_event(d2: u32, d3: u32) {
    if TRB_FIELD_TO_TYPE(d3) != TRB_EV_TRANSFER {
        return;
    }
    let slot = TRB_TO_SLOT_ID(d3);
    let epid = TRB_TO_EP_ID(d3);
    let n = NXDEV.load(Ordering::Relaxed) as usize;
    for i in 0..n {
        let d = unsafe { dev_mut(i) };
        if d.slot == slot && d.dci == epid {
            d.ready = true;
            d.code = GET_COMP_CODE(d2);
        }
    }
}

/// Wait for a transfer-complete event on (slot, endpoint id).
fn wait_xfer(slot: u32, epid: u32) -> u32 {
    let mut code = u32::MAX;
    let mut done = false;
    for _ in 0..4_000_000 {
        if done {
            break;
        }
        unsafe {
            drain_events(&mut |_d0, _d1, d2, d3| {
                route_event(d2, d3);
                if TRB_FIELD_TO_TYPE(d3) == TRB_EV_TRANSFER
                    && TRB_TO_SLOT_ID(d3) == slot
                    && TRB_TO_EP_ID(d3) == epid
                {
                    code = GET_COMP_CODE(d2);
                    done = true;
                }
            });
        }
        core::hint::spin_loop();
    }
    if !done { 0 } else { code }
}

// ---------------------------------------------------------------------------
// Control transfer on a device's EP0 (setup + optional data + status).
// ---------------------------------------------------------------------------
fn ctrl_xfer(
    di: usize,
    bmreq: u8,
    breq: u8,
    wval: u16,
    widx: u16,
    buf: *mut u8,
    buflen: u32,
    dir_in: bool,
) -> Result<usize, u32> {
    unsafe {
        let d = dev_mut(di);
        let slot = d.slot;
        let e0 = &mut d.ep0;
        let tx = if buflen > 0 {
            TRB_TX_TYPE(if dir_in { TRB_DATA_IN } else { TRB_DATA_OUT })
        } else {
            0
        };
        trb_push(
            e0,
            (wval as u32) << 16 | (breq as u32) << 8 | bmreq as u32,
            (widx as u32) | buflen << 16,
            TRB_LEN(8),
            TYPE(TRB_SETUP) | TRB_IDT | tx,
        );
        if buflen > 0 {
            trb_push(
                e0,
                buf as u32,
                (buf as usize >> 32) as u32,
                TRB_LEN(buflen),
                TYPE(TRB_DATA) | if dir_in { TRB_DIR_IN } else { 0 },
            );
        }
        let stat_dir = if buflen > 0 && dir_in { 0 } else { TRB_DIR_IN };
        trb_push(e0, 0, 0, 0, TYPE(TRB_STATUS) | TRB_IOC | stat_dir);
        wr(db(slot), 1);
        let _ = rd(db(slot));
    }
    let code = wait_xfer(unsafe { XDEVS[di].slot }, 1);
    if code == COMP_SUCCESS || code == COMP_SHORT {
        Ok(buflen as usize)
    } else {
        Err(code)
    }
}

// ---------------------------------------------------------------------------
// Context construction
// ---------------------------------------------------------------------------
fn build_addr_ictx(port_speed: u32, port_num: u32, ep0: *mut u32, mps0: u32) {
    unsafe {
        core::ptr::write_bytes(INCTX, 0, 32 + 8 * cs());
        let cc = INCTX as *mut u32;
        *cc.add(0) = 0; // Drop flags
        *cc.add(1) = (1 << 0) | (1 << 1); // Add slot + EP0
        let sc = INCTX.offset(32) as *mut u32;
        *sc.add(0) = LAST_CTX(1) | SLOT_SPEED_X(port_speed);
        *sc.add(1) = ROOTPORT(port_num + 1);
        let e0 = INCTX.add(32 + cs()) as *mut u32;
        *e0.add(0) = 0;
        *e0.add(1) = EP_TYPE(CTRL_EP) | MAX_PACKET(mps0) | ERROR_COUNT(3);
        (INCTX.add(32 + cs() + 8) as *mut u64).write_volatile((ep0 as u64) | 1);
        (INCTX.add(32 + cs() + 16) as *mut u32).write_volatile(EP_AVG(8));
    }
}

// ---------------------------------------------------------------------------
// Port reset
// ---------------------------------------------------------------------------
unsafe fn reset_port(port: u32) -> Result<u32, ()> {
    unsafe {
        let psc = op(OP_PORTSC) + (port as usize) * PORT_STRIDE;
        wr(psc, rd(psc) | PORT_POWER);
        for _ in 0..1000 {
            core::hint::spin_loop();
        }
        wr(psc, rd(psc) | PORT_RESET);
        for _ in 0..4_000_000 {
            if rd(psc) & PORT_RESET == 0 {
                break;
            }
            core::hint::spin_loop();
        }
        for _ in 0..4_000_000 {
            if rd(psc) & PORT_PE != 0 {
                break;
            }
            core::hint::spin_loop();
        }
        if rd(psc) & PORT_PE == 0 {
            return Err(());
        }
        Ok(DEV_SPEED(rd(psc)))
    }
}

// ---------------------------------------------------------------------------
// Descriptor parsing: first boot HID mouse/keyboard interface
// ---------------------------------------------------------------------------
/// Returns (kind, interface number, ep address, mps, interval).
fn find_hid(cfg: &[u8]) -> Option<(Kind, u8, u8, u32, u32)> {
    let mut i = 0usize;
    let mut cand: Option<(Kind, u8)> = None;
    while i + 2 <= cfg.len() {
        let blen = cfg[i] as usize;
        if blen < 2 {
            break;
        }
        let bt = cfg[i + 1];
        if bt == DT_INTERFACE && blen >= 9 && i + 9 <= cfg.len() {
            let iface = cfg[i + 2];
            let class = cfg[i + 5];
            let sub = cfg[i + 6];
            let proto = cfg[i + 7];
            cand = if class == HID_CLASS && sub == HID_SUBCLASS_BOOT {
                match proto {
                    HID_PROTO_KEYBOARD => Some((Kind::Keyboard, iface)),
                    HID_PROTO_MOUSE => Some((Kind::Mouse, iface)),
                    _ => None,
                }
            } else {
                None
            };
        } else if bt == DT_ENDPOINT && blen >= 7 && i + 7 <= cfg.len() {
            if let Some((kind, iface)) = cand {
                let addr = cfg[i + 2];
                let attrs = cfg[i + 3];
                let mps = (cfg[i + 4] as u32) | ((cfg[i + 5] as u32) << 8);
                let interval = cfg[i + 6] as u32;
                if (addr & 0x80) != 0 && (attrs & 0x03) == 0x03 {
                    return Some((kind, iface, addr, mps, interval));
                }
            }
        }
        i += blen;
    }
    None
}

// ---------------------------------------------------------------------------
// Enumeration: slot + address + descriptors + configure IN endpoint
// ---------------------------------------------------------------------------
fn enumerate(port: u32, speed: u32) -> bool {
    let idx = NXDEV.load(Ordering::Relaxed) as usize;
    if idx >= MAX_XDEV {
        return false;
    }
    // Per-device rings + output device context.
    let ep0_raw = dma_alloc(8 * 16, 64) as *mut u32;
    let epin_raw = dma_alloc(8 * 16, 64) as *mut u32;
    let devctx = dma_alloc((1 + 8) * cs(), 64);
    let report = unsafe { REPORTS.add(idx * 16) };
    if ep0_raw.is_null() || epin_raw.is_null() || devctx.is_null() || report.is_null() {
        crate::serial_println!("usb: DMA allocation failed");
        return false;
    }

    // Enable slot.
    let (code, slot) = issue_cmd(TYPE(TRB_ENABLE_SLOT));
    if code != COMP_SUCCESS || slot == 0 {
        crate::serial_println!("usb: enable-slot failed code={}", code);
        return false;
    }
    unsafe {
        *DCBAA.add(slot as usize) = devctx as u64;
        let d = dev_mut(idx);
        d.slot = slot;
        d.ep0 = Ring {
            raw: ep0_raw,
            count: 8,
            idx: 0,
            cyc: 1,
        };
        d.epin = Ring {
            raw: epin_raw,
            count: 8,
            idx: 0,
            cyc: 1,
        };
        d.devctx = devctx;
        d.armed = false;
        d.ready = false;
        d.code = 0;
        d.reports = 0;
        d.failures = 0;
        d.kbd = crate::hid::Keyboard::new();
    }

    // Address device with EP0 max packet 8 (first transactions are tiny).
    build_addr_ictx(speed, port, ep0_raw, 8);
    let (code, _) = issue_cmd_ptr(TYPE(TRB_ADDR_DEV) | SLOT_ID(slot), unsafe {
        INCTX as usize
    });
    if code != COMP_SUCCESS {
        crate::serial_println!("usb: address-device failed code={}", code);
        return false;
    }

    // Device descriptor.
    match read_device_descriptor(idx) {
        Ok(d) => {
            crate::serial_println!(
                "usb: dev {:04x}:{:04x} class={:#x} mps0={} slot={}",
                u16::from_le_bytes([d[8], d[9]]),
                u16::from_le_bytes([d[10], d[11]]),
                d[4],
                d[7],
                slot
            );
        }
        Err(c) => {
            crate::serial_println!("usb: get-device-desc failed code={}", c);
            return false;
        }
    }

    // Config descriptor -> HID interface.
    let mut cfg = [0u8; 512];
    let got = match read_config_descriptor(idx, &mut cfg) {
        Ok(n) => n,
        Err(c) => {
            crate::serial_println!("usb: get-config failed code={}", c);
            return false;
        }
    };
    let Some((kind, iface, _epaddr, mps, interval)) = find_hid(&cfg[..got]) else {
        crate::serial_println!("usb: no HID mouse/keyboard endpoint in config");
        return false;
    };
    let cfgval = cfg[5];
    if cfgval == 0 {
        return false;
    }
    if ctrl_xfer(idx, 0x00, USB_REQ_SET_CONFIGURATION, cfgval as u16, 0, core::ptr::null_mut(), 0, false)
        .is_err()
    {
        return false;
    }
    // Boot protocol: fixed 3-byte mouse / 8-byte keyboard reports.
    if ctrl_xfer(idx, 0x21, REQ_SET_PROTOCOL, 0, iface as u16, core::ptr::null_mut(), 0, false)
        .is_err()
    {
        crate::serial_println!("usb: set-protocol(boot) stalled (continuing)");
    }
    let _ = ctrl_xfer(idx, 0x21, REQ_SET_IDLE, 0, iface as u16, core::ptr::null_mut(), 0, false);

    if !configure_hid_endpoint(idx, _epaddr, mps, interval) {
        return false;
    }
    unsafe {
        let d = dev_mut(idx);
        d.kind = kind;
        d.len = mps.min(8).max(1);
        d.armed = true;
    }
    NXDEV.fetch_add(1, Ordering::Relaxed);
    queue_report(idx);
    crate::serial_println!(
        "usb: HID {} armed (slot {} mps {} interval {})",
        if kind == Kind::Mouse { "mouse" } else { "keyboard" },
        slot,
        mps,
        interval
    );
    true
}

fn read_device_descriptor(di: usize) -> Result<[u8; 18], u32> {
    let mut d = [0u8; 18];
    let b = d.as_mut_ptr();
    ctrl_xfer(
        di,
        0x80,
        USB_REQ_GET_DESCRIPTOR,
        (DT_DEVICE as u16) << 8,
        0,
        b,
        18,
        true,
    )
    .map(|_| d)
}

fn read_config_descriptor(di: usize, buf: &mut [u8; 512]) -> Result<usize, u32> {
    ctrl_xfer(
        di,
        0x80,
        USB_REQ_GET_DESCRIPTOR,
        (DT_CONFIG as u16) << 8,
        0,
        buf.as_mut_ptr(),
        buf.len() as u32,
        true,
    )
}

/// Commit the interrupt-IN endpoint context and ring its doorbell.
fn configure_hid_endpoint(di: usize, epaddr: u8, mps: u32, interval: u32) -> bool {
    unsafe {
        let d = dev_mut(di);
        core::ptr::write_bytes(INCTX, 0, 32 + 8 * cs());
        let cc = INCTX as *mut u32;
        // Add Context flags are DCI bits: bit0 = slot, bit(2*ep+1) = IN ep.
        // Endpoint DCI for an IN endpoint is 2 * ep_number + 1 (EP1 IN -> 3).
        let euid = ((epaddr & 0x0f) as u32) * 2 + 1;
        *cc.add(0) = 0;
        *cc.add(1) = (1 << 0) | (1 << euid);
        d.dci = euid;
        let osc = d.devctx.cast::<u32>();
        let sc = INCTX.offset(32).cast::<u32>();
        *sc.add(0) = (*osc.add(0) & !LAST_CTX(31)) | LAST_CTX(euid);
        *sc.add(1) = *osc.add(1);
        let ep = INCTX.add(32 + (euid as usize) * cs()).cast::<u32>();
        *ep.add(0) = EP_INTERVAL(interval);
        *ep.add(1) = EP_TYPE(INT_IN_EP) | MAX_PACKET(mps) | ERROR_COUNT(3);
        *ep.add(4) = EP_AVG(mps);
        let deq = d.epin.raw as u64 | 1;
        *ep.offset(2).cast::<u64>() = deq;
        let slot = d.slot;
        let (code, _) = issue_cmd_ptr(TYPE(TRB_CONFIG_EP) | SLOT_ID(slot), INCTX as usize);
        if code != COMP_SUCCESS {
            crate::serial_println!("usb: configure-endpoint failed code={}", code);
            return false;
        }
        true
    }
}

/// Arm one interrupt-IN transfer on the device's endpoint ring.
fn queue_report(di: usize) {
    unsafe {
        let d = dev_mut(di);
        let slot = d.slot;
        let epid = d.dci;
        trb_push(
            &mut d.epin,
            REPORTS.add(di * 16) as u32,
            (REPORTS.add(di * 16) as usize >> 32) as u32,
            TRB_LEN(d.len),
            TYPE(TRB_NORMAL) | TRB_IOC,
        );
        d.ready = false;
        wr(db(slot), epid);
        let _ = rd(db(slot));
    }
}

// ---------------------------------------------------------------------------
// Controller bring-up
// ---------------------------------------------------------------------------
fn find_xhci() -> Option<(u8, u8, u8)> {
    for d in crate::pci::probe() {
        if d.class == 0x0C && d.subclass == 0x03 && d.prog_if == 0x30 {
            return Some((d.bus, d.device, d.func));
        }
    }
    None
}

pub fn init() -> bool {
    let Some((b, dev, func)) = find_xhci() else {
        crate::serial_println!("usb: no xHCI controller");
        return false;
    };
    crate::serial_println!("usb: xHCI at {:02x}:{:02x}.{}", b, dev, func);

    unsafe {
        let cmd = crate::pci::config_read(b, dev, func, 0x04);
        crate::pci::config_write(
            b,
            dev,
            func,
            0x04,
            (cmd & 0xFFFF_0000) | ((cmd & 0xFFFF) | 0x6 | 0x400),
        );
        let bar = crate::pci::bar64(b, dev, func, 0);
        let base = bar.unwrap_or(0) as usize;
        if base == 0 {
            crate::serial_println!("usb: bad BAR");
            return false;
        }
        M = base;
        let caplen = rd(base + CAP_CAPLENGTH);
        let hcs1 = rd(base + HCSPARAMS1);
        let hcs2 = rd(base + HCSPARAMS2);
        let hcc = rd(base + HCCPARAMS1);
        OPB = base + (caplen & 0xff) as usize;
        DBB = base + rd(base + DBOFF) as usize;
        RTB = base + rd(base + RTSOFF) as usize;
        MAXPORTS = (hcs1 >> 24) & 0xff;
        N_SLOTS = (hcs1 & 0xff).clamp(1, MAX_XDEV as u32);
        CTRL_CS = if hcc & (1 << 2) != 0 { 64 } else { 32 };
        // HCSPARAMS2: Max Scratchpad Buffers = Hi[31:27] << 5 | Lo[25:21].
        NUM_SPB = (((hcs2 >> 27) & 0x1f) << 5) | ((hcs2 >> 21) & 0x1f);
        crate::serial_println!(
            "usb: mmio={:#x} caplen={:#x} hcs1={:#x} hcs2={:#x} hcc1={:#x}",
            base,
            caplen,
            hcs1,
            hcs2,
            hcc
        );
        let (maxports, nslots, csz, nspb) = (MAXPORTS, N_SLOTS, CTRL_CS, NUM_SPB);
        crate::serial_println!(
            "usb: ports={} slots={} context={}B scratchpad={}",
            maxports,
            nslots,
            csz,
            nspb
        );

        match crate::pci::claim_usb_legacy(b, dev, func, base, 0x04, true) {
            0 => crate::serial_println!("usb: xHCI USB legacy: already released"),
            1 => crate::serial_println!("usb: xHCI USB legacy handed over by BIOS"),
            2 => crate::serial_println!("usb: xHCI USB legacy claim timed out"),
            _ => crate::serial_println!("usb: xHCI USB legacy: no capability"),
        }
    }

    // Reset: stop, then RESET, then wait for CNR to clear.
    unsafe {
        wr(op(OP_USBCMD), rd(op(OP_USBCMD)) & !CMD_RUN);
        let mut t = 0;
        while t < 1_000_000 && rd(op(OP_USBSTS)) & STS_HALT == 0 {
            t += 1;
            core::hint::spin_loop();
        }
        wr(op(OP_USBCMD), rd(op(OP_USBCMD)) | CMD_RESET);
        t = 0;
        while t < 10_000_000 && rd(op(OP_USBCMD)) & CMD_RESET != 0 {
            t += 1;
            core::hint::spin_loop();
        }
        t = 0;
        while t < 10_000_000 && rd(op(OP_USBSTS)) & STS_CNR != 0 {
            t += 1;
            core::hint::spin_loop();
        }
    }

    // Rings, contexts, scratchpads.
    unsafe {
        let cmd_ring = dma_alloc(64 * 16, 64) as *mut u32;
        let ev_ring = dma_alloc(32 * 16, 64) as *mut u32;
        let erst = dma_alloc(16, 64);
        let dcb = dma_alloc((N_SLOTS as usize + 1) * 8, 64) as *mut u64;
        let ictx = dma_alloc(32 + 8 * cs(), 64);
        let reports = dma_alloc(4096, 64);
        if cmd_ring.is_null()
            || ev_ring.is_null()
            || erst.is_null()
            || dcb.is_null()
            || ictx.is_null()
            || reports.is_null()
        {
            crate::serial_println!("usb: DMA allocation failed");
            return false;
        }
        CMD = Ring {
            raw: cmd_ring,
            count: 64,
            idx: 0,
            cyc: 1,
        };
        EV = Ring {
            raw: ev_ring,
            count: 32,
            idx: 0,
            cyc: 1,
        };
        ERST = erst;
        DCBAA = dcb;
        INCTX = ictx;
        REPORTS = reports;

        // Scratchpad buffers (§4.20): a 64-byte array of buffer addresses
        // whose entries the controller requires to exist when HCSPARAMS2
        // reports them.  DCBAAP[0] must point at that array.
        if NUM_SPB > 0 {
            let arr = dma_alloc(NUM_SPB as usize * 8 + 64, 64) as *mut u64;
            if !arr.is_null() {
                for i in 0..NUM_SPB as usize {
                    let p = dma_alloc(4096, 64);
                    if p.is_null() {
                        crate::serial_println!("usb: scratchpad buffer {} failed", i);
                        return false;
                    }
                    arr.add(i).write_volatile(p as u64);
                }
                SPB_ARRAY = arr;
                *DCBAA.add(0) = arr as u64;
                let nspb = NUM_SPB;
                crate::serial_println!("usb: allocated {} scratchpad buffer(s)", nspb);
            }
        }

        let e = ERST as *mut u32;
        *e.add(0) = ev_ring as u32;
        *e.add(1) = (ev_ring as usize >> 32) as u32;
        *e.add(2) = 32;
    }

    unsafe {
        wr64(op(OP_DCBAAP), DCBAA as u64);
        wr64(
            op(OP_CRCR),
            ((CMD.raw as u64) & !0x3f) | CMD_RING_CYCLE as u64,
        );
        wr(op(OP_CONFIG), (rd(op(OP_CONFIG)) & !0xff) | N_SLOTS);
        wr(rt(RT_ERSTSZ), 1);
        wr64(rt(RT_ERSTBA), ERST as u64);
        wr64(rt(RT_ERDP), EV.raw as u64);
        wr(rt(RT_IMAN), (rd(rt(RT_IMAN)) & !0x2) | 0x1);
        wr(rt(RT_IMOD), 0);
        wr(op(OP_USBCMD), rd(op(OP_USBCMD)) | CMD_RUN);
        let mut t = 0;
        while t < 4_000_000 && rd(op(OP_USBSTS)) & STS_HALT != 0 {
            t += 1;
            core::hint::spin_loop();
        }
        if rd(op(OP_USBSTS)) & STS_HALT != 0 {
            crate::serial_println!("usb: controller failed to start");
            return false;
        }
    }
    INITOK.store(true, Ordering::Relaxed);
    crate::serial_println!(
        "usb: controller up ({} ports) — HID discovery deferred to polling",
        unsafe { MAXPORTS }
    );
    true
}

/// Scan ports until all slots are used or every port has been tried.
fn maybe_discover() {
    if !INITOK.load(Ordering::Relaxed) {
        return;
    }
    if NXDEV.load(Ordering::Relaxed) as usize >= MAX_XDEV {
        return;
    }
    let now = crate::interrupts::uptime_ms() as u32;
    if now.wrapping_sub(LAST_SCAN.load(Ordering::Relaxed)) < 200 {
        return;
    }
    LAST_SCAN.store(now, Ordering::Relaxed);
    let maxports = unsafe { MAXPORTS };
    for p in 0..maxports {
        if DISCOVERED_TRIED.load(Ordering::Relaxed) & (1 << p) != 0 {
            continue;
        }
        unsafe {
            let psc = op(OP_PORTSC) + (p as usize) * PORT_STRIDE;
            if rd(psc) & (PORT_CONNECT | PORT_PE) == PORT_CONNECT {
                crate::serial_println!("usb: port {} connect (speed {})", p, DEV_SPEED(rd(psc)));
                if let Ok(speed) = reset_port(p) {
                    if enumerate(p, speed) {
                        continue; // keep scanning for the other HID device
                    }
                }
                DISCOVERED_TRIED.fetch_or(1 << p, Ordering::Relaxed);
            }
        }
    }
}

fn has_kind(kind: Kind) -> bool {
    let n = NXDEV.load(Ordering::Relaxed) as usize;
    for i in 0..n {
        if unsafe { XDEVS[i].kind } == kind {
            return true;
        }
    }
    false
}

pub fn mouse_present() -> bool {
    has_kind(Kind::Mouse)
}

pub fn controller_up() -> bool {
    INITOK.load(Ordering::Relaxed)
}

pub fn rescan() {
    DISCOVERED_TRIED.store(0, Ordering::Relaxed);
    LAST_SCAN.store(0, Ordering::Relaxed);
}

/// Drain completion events into the device table.
fn pump_reports() {
    unsafe {
        drain_events(&mut |_d0, _d1, d2, d3| route_event(d2, d3));
    }
}

pub fn poll_mouse() -> Option<MouseEvent> {
    if !INITOK.load(Ordering::Relaxed) {
        return None;
    }
    if !has_kind(Kind::Mouse) {
        maybe_discover();
    }
    pump_reports();
    let n = NXDEV.load(Ordering::Relaxed) as usize;
    for i in 0..n {
        let d = unsafe { dev_mut(i) };
        if d.kind != Kind::Mouse || !d.armed || !d.ready {
            continue;
        }
        if d.code != COMP_SUCCESS && d.code != COMP_SHORT {
            d.armed = false;
            d.failures += 1;
            d.ready = false;
            continue;
        }
        let hb = unsafe { REPORTS.add(i * 16) };
        let (b0, b1, b2) = unsafe {
            (
                core::ptr::read_volatile(hb),
                core::ptr::read_volatile(hb.add(1)),
                core::ptr::read_volatile(hb.add(2)),
            )
        };
        let ev = MouseEvent {
            dx: (b1 as i8) as i16,
            dy: (b2 as i8) as i16,
            left: b0 & 0x01 != 0,
            right: b0 & 0x02 != 0,
            middle: b0 & 0x04 != 0,
        };
        d.reports += 1;
        if d.reports == 1 {
            crate::serial_println!("usb: mouse report dx={} dy={} b={:#x}", ev.dx, ev.dy, b0);
        }
        queue_report(i);
        return Some(ev);
    }
    None
}

pub fn try_poll_key() -> Option<crate::ps2::KeyEvent> {
    if !INITOK.load(Ordering::Relaxed) {
        return None;
    }
    if !has_kind(Kind::Keyboard) {
        maybe_discover();
    }
    if let Some(b) = crate::hid::pop_byte() {
        return crate::ps2::feed_scancode(b);
    }
    pump_reports();
    let n = NXDEV.load(Ordering::Relaxed) as usize;
    for i in 0..n {
        let d = unsafe { dev_mut(i) };
        if d.kind != Kind::Keyboard || !d.armed || !d.ready {
            continue;
        }
        if d.code != COMP_SUCCESS && d.code != COMP_SHORT {
            d.armed = false;
            d.failures += 1;
            d.ready = false;
            continue;
        }
        let hb = unsafe { REPORTS.add(i * 16) };
        let mut rep = [0u8; 8];
        let len = (d.len as usize).min(8);
        for (k, b) in rep.iter_mut().take(len).enumerate() {
            *b = unsafe { core::ptr::read_volatile(hb.add(k)) };
        }
        d.kbd.diff(&rep);
        d.reports += 1;
        if d.reports == 1 {
            crate::serial_println!(
                "usb: kbd report mods={:#x} keys={:02x} {:02x} {:02x} {:02x} {:02x} {:02x}",
                rep[0], rep[2], rep[3], rep[4], rep[5], rep[6], rep[7]
            );
        }
        queue_report(i);
        return crate::hid::pop_byte().and_then(crate::ps2::feed_scancode);
    }
    None
}

// ---------------------------------------------------------------------------
// Diagnostics (shell `usb`)
// ---------------------------------------------------------------------------
pub fn diagnose() -> alloc::string::String {
    let mut s = alloc::string::String::new();
    s.push_str(&alloc::format!(
        "xhci mmio={:#x} mapped={} ports={} slots={} scratchpad={}\n",
        unsafe { M },
        if crate::paging::is_mapped(unsafe { M }) {
            "yes"
        } else {
            "no"
        },
        unsafe { MAXPORTS },
        unsafe { N_SLOTS },
        unsafe { NUM_SPB }
    ));
    let n = NXDEV.load(Ordering::Relaxed) as usize;
    s.push_str(&alloc::format!("  {} device(s):\n", n));
    for i in 0..n {
        let d = unsafe { &*core::ptr::addr_of!(XDEVS[i]) };
        s.push_str(&alloc::format!(
            "    {} slot={} dci={} armed={} fails={} reports={} ring={:#x}\n",
            if d.kind == Kind::Mouse { "mouse" } else { "kbd" },
            d.slot,
            d.dci,
            d.armed,
            d.failures,
            d.reports,
            d.epin.raw as usize
        ));
    }
    unsafe {
        for p in 0..8 {
            let psc = op(OP_PORTSC) + (p as usize) * PORT_STRIDE;
            let v = rd(psc);
            s.push_str(&alloc::format!("  port{}: {:#010x}\n", p, v));
        }
    }

    // Cross-controller attach scan (read-only): locate devices on real hw.
    s.push_str("usbctl: scanning all USB controllers:\n");
    unsafe {
        for d in crate::pci::probe() {
            if d.class != 0x0C || d.subclass != 0x03 {
                continue;
            }
            let lbl = match d.prog_if {
                0x00 => "UHCI",
                0x10 => "OHCI",
                0x20 => "EHCI",
                0x30 => "xHCI",
                _ => "USB?",
            };
            let bar = crate::pci::bar64(d.bus, d.device, d.func, 0);
            let base = bar.unwrap_or(0);
            s.push_str(&alloc::format!(
                "  {} {:02x}:{:02x}.{} bar={:#x}\n",
                lbl,
                d.bus,
                d.device,
                d.func,
                base
            ));
            if base == 0 || base == u64::MAX {
                continue;
            }
            match d.prog_if {
                0x30 => {
                    let caplen = rd(base as usize);
                    let hcs1 = rd(base as usize + 0x04);
                    let np = ((hcs1 >> 24) & 0xff) as usize;
                    let opb = base as usize + (caplen & 0xff) as usize;
                    s.push_str("     xhci ports=");
                    for p in 0..np.min(16) {
                        let v = rd(opb + 0x400 + p * 0x10);
                        s.push_str(&alloc::format!(
                            "{}{:#x} ",
                            if v & 1 != 0 { "*" } else { "" },
                            v & 0x3ff
                        ));
                    }
                    s.push('\n');
                }
                0x20 => {
                    let cl = rd(base as usize) & 0xff;
                    let sp = rd(base as usize + 0x04);
                    let hc = rd(base as usize + 0x08);
                    let np = (sp & 0xf) as usize;
                    let opb = base as usize + cl as usize;
                    let cfg = rd(opb + 0x40);
                    s.push_str(&alloc::format!(
                        "     ehci nports={} hcc={:#x} cfgflag={:#x} ",
                        np, hc, cfg
                    ));
                    for p in 0..np.min(16) {
                        let v = rd(opb + 0x44 + p * 0x04);
                        s.push_str(&alloc::format!(
                            "{}{:#x},",
                            if v & 1 != 0 { "*" } else { "" },
                            v & 0xffff
                        ));
                    }
                    s.push('\n');
                }
                0x00 => {
                    let lo = crate::pci::config_read(d.bus, d.device, d.func, 0x10);
                    let io = (lo & 0xFFF0) as u16;
                    if io == 0 {
                        continue;
                    }
                    s.push_str("     uhci ports=");
                    for p in 0..2 {
                        let pr = io_iow16(io + 0x10 + (p as u16) * 2);
                        s.push_str(&alloc::format!(
                            "{}{:x},",
                            if pr & 1 != 0 { "*" } else { "" },
                            pr & 0xff
                        ));
                    }
                    s.push('\n');
                }
                _ => {}
            }
        }
    }
    s
}