// ---------------------------------------------------------------------------
// USB host controller driver — xHCI (USB 2.0/3.x) + HID mouse support.
//
// **Polled** driver: it runs the xHCI event ring from the cursor task loop
// instead of wiring MSI/INTx into the PIC, matching how the kernel already
// polls the PS/2 mouse (see `cursor`).
//
// Minimal, mouse-only scope:
//   • take the controller out of reset, set up the command + event rings
//   • reset ports, run the slot/address/configure firmware sequence
//   • find the HID mouse boot interface, open its interrupt-IN endpoint,
//     and translate the standard report into `ps2::MouseEvent` (reusing
//     the cursor task's consumption path).
//
// DMA / MMIO: the kernel is identity-mapped (physical == virtual), so bus
// addresses are plain kernel pointers and every heap allocation is DMA-able.
// Rings/contexts are allocated once at boot and handed to the controller
// permanently.
//
// Register/TRB/context layouts follow the Linux xHCI driver definitions
// (drivers/usb/host/xhci*.h) and the xHCI 1.2 spec.  All exact field
// positions were cross-checked against the Linux sources:
//   * TRB type field is bits 10-15 (TRB_TYPE(p) = p << 10)
//   * transfer/command completion code is bits 24-31 of event dword2
//   * slot id = bits 24-31, ep id = bits 16-20 of event dword3
//   * endpoint context: type at bits 3-5, interval at 16-23,
//     max packet size at bits 16-31 of dword1, TR dequeue is a 64-bit
//     field at bytes 8-15 with DCS in bit 0
//   * slot context: route/LAST_CTX in dword0, root-hub-port at 16-23 of
//     dword1, speed at bits 20-23
// ---------------------------------------------------------------------------
#![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[23:16]
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; // PORTSC port 1 (port set stride 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; // bits 10..15
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_DIR_IN: u8 = 0x80;
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;

// ---------------------------------------------------------------------------
// 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;

struct Ring {
    raw: *mut u32,
    count: usize, // TRBs
    idx: usize,   // producer (command/transfer) or consumer (event) index
    cyc: u32,     // software cycle
}
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 EP0: Ring = ri_empty(8);
static mut EPIN: Ring = ri_empty(8);
static mut EV: Ring = ri_empty(32);

static mut DCBAA: *mut u64 = core::ptr::null_mut();
static mut DEVCTX: *mut u8 = core::ptr::null_mut();
static mut INCTX: *mut u8 = core::ptr::null_mut();
static mut ERST: *mut u8 = core::ptr::null_mut();

static SLOT: AtomicU32 = AtomicU32::new(0);
static HID_EPID: AtomicU32 = AtomicU32::new(0); // device-context index of the IN endpoint
static mut HID_BUF: [u8; 8] = [0; 8];
static mut HID_LEN: usize = 4;
static mut MPS0: u32 = 8;

static mut GOT_REPORT: bool = false;
static mut QUEUED: bool = false;
static mut RBTN: u8 = 0;
static mut RDX: i8 = 0;
static mut RDY: i8 = 0;

static INITOK: AtomicBool = AtomicBool::new(false);
static MOUSEPRESENT: AtomicBool = AtomicBool::new(false);

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 }
}
fn db(slot: u32) -> usize {
    unsafe { DBB }.wrapping_add(slot as usize * 4)
}
fn rt(off: usize) -> usize {
    unsafe { RTB + off }
}

fn dma_alloc(bytes: usize, align: usize) -> *mut u8 {
    unsafe {
        let layout = alloc::alloc::Layout::from_size_align(bytes.max(1), align).unwrap();
        let p = alloc::alloc::alloc(layout);
        if !p.is_null() {
            core::ptr::write_bytes(p, 0, bytes.max(1));
        }
        p
    }
}

// ---------------------------------------------------------------------------
// 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;
        }
    }
}

// ---------------------------------------------------------------------------
// Event ring: drain produced events, calling handle(d0,d1,d2,d3) per event,
// then advance ERDP to the consumption point so the controller reuses it.
// ---------------------------------------------------------------------------
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;
    let mut raw: [u32; 4] = [0; 4];
    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);
                    raw = [d0, d1, d2, d3];
                    done = true;
                }
            });
        }
        core::hint::spin_loop();
    }
    if code != COMP_SUCCESS {
        unsafe {
            let evidx = core::ptr::addr_of!(EV).cast::<Ring>().read().idx;
            crate::serial_println!(
                "usb: cmd-complete code={} slot={} seen={} evidx={}",
                code, slot, seen, evidx
            );
            for i in 0..4 {
                let s = EV.raw.add(i * 4);
                crate::serial_println!(
                    "  ev[{i}] {:08x} {:08x} {:08x} {:08x}",
                    s.add(0).read(), s.add(1).read(), s.add(2).read(), s.add(3).read()
                );
            }
        }
    }
    (code, slot)
}

/// Enqueue a command TRB and ring the host command doorbell.
fn issue_cmd(d3: u32) -> (u32, u32) {
    issue_cmd_ptr(d3, 0)
}

/// Enqueue a command TRB that carries a 64-bit parameter (used by
/// Address-Device / Configure-Endpoint for the input-context pointer).
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()
}

/// Wait for a transfer-complete event on the given device-context endpoint id.
fn wait_xfer(epid: u32) -> u32 {
    let mut code = 0u32;
    let mut done = false;
    let sl = SLOT.load(Ordering::Relaxed);
    for _ in 0..4_000_000 {
        if done {
            break;
        }
        unsafe {
            drain_events(&mut |_, _, d2, d3| {
                if TRB_FIELD_TO_TYPE(d3) == TRB_EV_TRANSFER
                    && TRB_TO_SLOT_ID(d3) == sl
                    && TRB_TO_EP_ID(d3) == epid
                {
                    code = GET_COMP_CODE(d2);
                    done = true;
                }
            });
        }
        core::hint::spin_loop();
    }
    code
}

// ---------------------------------------------------------------------------
// Control transfer on EP0 (setup + optional data + status), per Linux
// xhci_queue_ctrl_tx.  Returns OK(len) or the failure completion code.
// ---------------------------------------------------------------------------
fn ctrl_xfer(
    bmreq: u8,
    breq: u8,
    wval: u16,
    widx: u16,
    buf: *mut u8,
    buflen: u32,
    dir_in: bool,
) -> Result<usize, u32> {
    let slot = SLOT.load(Ordering::Relaxed);
    let e0 = unsafe { &mut *core::ptr::addr_of_mut!(EP0) };
    unsafe {
        // Setup stage (immediate data, 8-byte setup packet in d0/d1).
        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 },
            );
        }
        // Status stage.
        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); // ring EP0 (doorbell target value 1)
        let _ = rd(db(slot));
    }
    let code = wait_xfer(1);
    if code == COMP_SUCCESS || code == COMP_SHORT {
        Ok(buflen as usize)
    } else {
        Err(code)
    }
}

// ---------------------------------------------------------------------------
// Context construction helpers
// ---------------------------------------------------------------------------
fn build_addr_ictx(port_speed: u32, port_num: 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.raw 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) * 0x10;
        // Ensure powered.
        wr(psc, rd(psc) | PORT_POWER);
        for _ in 0..1000 {
            core::hint::spin_loop();
        }
        // Assert reset, wait for the hardware to clear it (reset done).
        wr(psc, rd(psc) | PORT_RESET);
        for _ in 0..4_000_000 {
            if rd(psc) & PORT_RESET == 0 {
                break;
            }
            core::hint::spin_loop();
        }
        // Wait for the port to be enabled.
        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)))
    }
}

// ---------------------------------------------------------------------------
// Enumeration + HID setup
// ---------------------------------------------------------------------------
fn read_device_descriptor() -> Result<[u8; 18], u32> {
    let mut d = [0u8; 18];
    ctrl_xfer(
        USB_DIR_IN,
        USB_REQ_GET_DESCRIPTOR,
        (DT_DEVICE as u16) << 8,
        0,
        d.as_mut_ptr(),
        18,
        true,
    )?;
    Ok(d)
}

fn read_config_descriptor(buf: &mut [u8]) -> Result<usize, u32> {
    let total = ctrl_xfer(
        USB_DIR_IN,
        USB_REQ_GET_DESCRIPTOR,
        (DT_CONFIG as u16) << 8,
        0,
        buf.as_mut_ptr(),
        buf.len() as u32,
        true,
    )?;
    Ok(total)
}

/// Walk a config descriptor tree looking for a HID interrupt IN endpoint
/// with an interrupt (transfer type 3).  Returns (mps, interval, eepid).
fn find_hid_mouse(cfg: &[u8]) -> Option<(u32, u32, u32)> {
    let mut i = 0usize;
    let mut in_iface = false;
    while i + 2 <= cfg.len() {
        let bLen = cfg[i] as usize;
        if bLen == 0 {
            break;
        }
        let btype = cfg[i + 1];
        if btype == DT_INTERFACE && i + 9 <= cfg.len() {
            // interface descriptor: +5 = bInterfaceClass
            in_iface = cfg[i + 5] == HID_CLASS;
        } else if btype == DT_ENDPOINT && in_iface && i + 7 <= cfg.len() {
            let addr = cfg[i + 2];
            let attrs = cfg[i + 3];
            let dir_in = addr & 0x80 != 0;
            let xfer = attrs & 0x03;
            if dir_in && xfer == 3 {
                let mps = (cfg[i + 4] as u32) | ((cfg[i + 5] as u32) << 8);
                let interval = cfg[i + 6] as u32;
                let epnum = (addr & 0x0f) as u32;
                let euid = 2 * epnum + 1; // device-context index (IN)
                return Some((mps, interval, euid));
            }
        }
        i += bLen;
    }
    None
}

// ---------------------------------------------------------------------------
// Bring one port's device through address + configure to an HID mouse
// ---------------------------------------------------------------------------
fn enumerate(port: u32, speed: u32) -> bool {
    let _ = DEV_SPEED;
    // 1. Enable slot.
    let (code, slotid) = issue_cmd(TYPE(TRB_ENABLE_SLOT));
    if code != COMP_SUCCESS || slotid == 0 {
        crate::serial_println!("usb: enable-slot failed code={} slot={}", code, slotid);
        return false;
    }
    SLOT.store(slotid, Ordering::Relaxed);
    unsafe {
        *DCBAA.add(slotid as usize) = DEVCTX as u64;
    }
    // 2. Address device.
    build_addr_ictx(speed, port);
    let (code, _) = issue_cmd_ptr(TYPE(TRB_ADDR_DEV) | SLOT_ID(slotid), unsafe { INCTX } as usize);
    if code != COMP_SUCCESS {
        crate::serial_println!("usb: address-device failed code={}", code);
        return false;
    }
    // 3. Read the device descriptor (picks up bMaxPacketSize0).
    match read_device_descriptor() {
        Ok(d) => {
            let mps = d[7];
            if mps != 0 {
                unsafe {
                    MPS0 = mps as u32;
                }
            }
            crate::serial_println!(
                "usb: dev {:04x}:{:04x} class={:#x} mps0={}",
                u16::from_le_bytes([d[8], d[9]]),
                u16::from_le_bytes([d[10], d[11]]),
                d[4],
                d[7]
            );
        }
        Err(c) => crate::serial_println!("usb: get-device-desc failed code={}", c),
    }
    // 4. Read the whole config descriptor and find the HID mouse IN ep.
    let mut cfg = [0u8; 256];
    let got = match read_config_descriptor(&mut cfg) {
        Ok(n) => n,
        Err(c) => {
            crate::serial_println!("usb: get-config failed code={}", c);
            return false;
        }
    };
    let _ = got;
    match find_hid_mouse(&cfg) {
        Some((mps, interval, euid)) => {
            crate::serial_println!(
                "usb: HID mouse IN ep mps={} interval={} euid={}",
                mps,
                interval,
                euid
            );
            // 5. Set configuration = 1 (first config value present).
            let cfgval = cfg[5];
            if cfgval == 0 {
                return false;
            }
            if read_configuration(cfgval).is_err() {
                return false;
            }
            // 6. Configure the IN endpoint.
            if !configure_hid_endpoint(euid, mps, interval) {
                return false;
            }
            HID_EPID.store(euid, Ordering::Relaxed);
            unsafe {
                HID_LEN = mps as usize;
                QUEUED = false;
            }
            // 7. Arm the first IN read.
            queue_report();
            MOUSEPRESENT.store(true, Ordering::Relaxed);
            crate::serial_println!(
                "usb: HID boot mouse armed (slot {})",
                SLOT.load(Ordering::Relaxed)
            );
            true
        }
        None => {
            crate::serial_println!("usb: no HID mouse endpoint in config");
            false
        }
    }
}

fn read_configuration(val: u8) -> Result<(), u32> {
    ctrl_xfer(
        0,
        USB_REQ_SET_CONFIGURATION,
        val.into(),
        0,
        core::ptr::null_mut(),
        0,
        false,
    )
    .map(|_| ())
}

fn configure_hid_endpoint(euid: u32, mps: u32, interval: u32) -> bool {
    unsafe {
        let slot = SLOT.load(Ordering::Relaxed);
        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 << euid); // Add slot + this EP
        // Input slot context: copy the output (post-address) slot context so
        // route/speed/rootport are preserved, then bump LAST_CTX to the ep.
        let osc = DEVCTX.offset(0).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);
        // Endpoint IN context.
        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 = EPIN.raw as u64 | 1;
        *ep.offset(2).cast::<u64>() = deq;
        let (code, _) = issue_cmd_ptr(TYPE(TRB_CONFIG_EP) | SLOT_ID(slot), INCTX as usize);
        code == COMP_SUCCESS
    }
}

fn queue_report() {
    let slot = SLOT.load(Ordering::Relaxed);
    let epid = HID_EPID.load(Ordering::Relaxed);
    let epin = unsafe { &mut *core::ptr::addr_of_mut!(EPIN) };
    unsafe {
        let b = core::ptr::addr_of_mut!(HID_BUF) as *mut u8;
        trb_push(
            epin,
            b as u32,
            (b as usize >> 32) as u32,
            TRB_LEN(HID_LEN as u32),
            TYPE(TRB_NORMAL) | TRB_IOC,
        );
        wr(db(slot), epid);
        let _ = rd(db(slot));
        QUEUED = true;
    }
}

// ---------------------------------------------------------------------------
// 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);

    // Enable MMIO + bus-mastering on the PCI command register.
    unsafe {
        let cmd = crate::pci::config_read(b, dev, func, 0x04);
        crate::pci::config_write(
            b,
            dev,
            func,
            0x04,
            (cmd & 0xFFFF0000) | ((cmd & 0xFFFF) | 0x6),
        );
        let bar = crate::pci::bar64(b, dev, func, 0);
        let base = bar.unwrap() as usize;
        M = base;
        let caplen = rd(base + CAP_CAPLENGTH);
        let hcs1 = rd(base + HCSPARAMS1);
        let hcc = rd(base + HCCPARAMS1);
        crate::serial_println!(
            "usb: mmio={:#x} caplen={:#x} hcs1={:#x} hcc1={:#x}",
            base, caplen, hcs1, hcc
        );
        OPB = base + (caplen & 0xff) as usize;
        DBB = base + rd(base + DBOFF) as usize;
        RTB = base + rd(base + RTSOFF) as usize;
        MAXPORTS = (hcs1 >> 24) & 0xff; // MaxPorts is bits 31:24 (bits 23:16 are MaxIntrs)
        let _ = hcc;
        CTRL_CS = if hcc & (1 << 2) != 0 { 64 } else { 32 };
    }

    // Reset the controller (stop, then reset).  Wait for reset completion
    // on the CMD_RESET bit clearing (per Linux xhci_reset), then for the
    // Controller-Not-Ready flag to clear before any doorbell/op-register
    // activity.
    unsafe {
        wr(op(OP_USBCMD), rd(op(OP_USBCMD)) & !CMD_RUN);
        // Wait for halt.
        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();
        }
    }

    // Allocate rings / contexts.
    unsafe {
        CMD.raw = dma_alloc(64 * 16, 64) as *mut u32;
        EV.raw = dma_alloc(32 * 16, 64) as *mut u32;
        EP0.raw = dma_alloc(8 * 16, 64) as *mut u32;
        EPIN.raw = dma_alloc(8 * 16, 64) as *mut u32;
        ERST = dma_alloc(16, 64);
        DCBAA = dma_alloc(64 * 8, 64) as *mut u64;
        DEVCTX = dma_alloc((1 + 8) * cs(), 64);
        INCTX = dma_alloc(32 + 8 * cs(), 64);
        // ERST segment 0 -> event ring.
        let e = ERST as *mut u32;
        *e.add(0) = EV.raw as u32;
        *e.add(1) = (EV.raw as usize >> 32) as u32;
        *e.add(2) = 32; // segment size (TRBs)
    }

    // Program registers.
    unsafe {
        wr64(op(OP_DCBAAP), DCBAA as u64);
        wr64(
            op(OP_CRCR),
            ((CMD.raw as u64) & !0x3f) | CMD_RING_CYCLE as u64,
        );
        // MaxSlotsEn = 1 (we use a single slot).
        wr(op(OP_CONFIG), (rd(op(OP_CONFIG)) & !0xff) | 1);
        // Event ring.
        wr(rt(RT_ERSTSZ), 1);
        wr64(rt(RT_ERSTBA), ERST as u64);
        wr64(rt(RT_ERDP), EV.raw as u64);
        // Interrupter 0: clear pending, enable.
        wr(rt(RT_IMAN), rd(rt(RT_IMAN)) | 0x2);
        wr(rt(RT_IMOD), 0);
        // Run.
        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) — mouse discovery deferred to cursor task",
        unsafe { MAXPORTS }
    );
    true
}

/// Bitmask of root ports whose attempt already failed; skip them so we don't
/// re-reset the port on every poll.
static DISCOVERED_TRIED: AtomicU32 = AtomicU32::new(0);

/// Periodic, cheap scan called from the cursor task until a HID mouse is
/// found.  A USB device can finish connecting long after the controller
/// starts, so discovery is driven asynchronously rather than blocking boot.
fn maybe_discover() {
    if MOUSEPRESENT.load(Ordering::Relaxed) {
        return;
    }
    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) * 0x10;
            // Only touch ports that are both connected and not yet enabled.
            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)
                    && enumerate(p, speed) {
                    return;
                }
                DISCOVERED_TRIED.fetch_or(1 << p, Ordering::Relaxed);
            }
        }
    }
}

pub fn mouse_present() -> bool {
    MOUSEPRESENT.load(Ordering::Relaxed)
}

/// True once the xHCI controller is brought up (regardless of whether a
/// device has attached yet).  The cursor task is spawned as long as this is
/// set so async mouse discovery can run.
pub fn controller_up() -> bool {
    INITOK.load(Ordering::Relaxed)
}

/// Reset discovery state (used by the `usb` shell command to rescan).
pub fn rescan() {
    DISCOVERED_TRIED.store(0, Ordering::Relaxed);
}

/// Non-blocking poll: drain the event ring, then decode one mouse report if
/// one arrived (re-arming the IN transfer so reports keep flowing).
/// Returns `None` when no report is pending (or no mouse).
pub fn poll_mouse() -> Option<MouseEvent> {
    // Drive discovery while no device is attached, then drain reports.
    if !MOUSEPRESENT.load(Ordering::Relaxed) {
        if INITOK.load(Ordering::Relaxed) {
            maybe_discover();
        }
        return None;
    }
    let epid = HID_EPID.load(Ordering::Relaxed);
    let sl = SLOT.load(Ordering::Relaxed);
    unsafe {
        drain_events(&mut |_, _, d2, d3| {
            if TRB_FIELD_TO_TYPE(d3) == TRB_EV_TRANSFER
                && TRB_TO_SLOT_ID(d3) == sl
                && TRB_TO_EP_ID(d3) == epid
            {
                let c = GET_COMP_CODE(d2);
                if c == COMP_SUCCESS || c == COMP_SHORT {
                    // A fresh report sits in HID_BUF.
                    let p = core::ptr::addr_of!(HID_BUF) as *const u8;
                    RBTN = core::ptr::read_volatile(p);
                    RDX = core::ptr::read_volatile(p.add(1)) as i8;
                    RDY = core::ptr::read_volatile(p.add(2)) as i8;
                    GOT_REPORT = true;
                    QUEUED = false;
                }
            }
        });
        if GOT_REPORT {
            GOT_REPORT = false;
            let ev = Some(MouseEvent {
                dx: RDX as i16,
                dy: RDY as i16,
                left: RBTN & 0x01 != 0,
                right: RBTN & 0x02 != 0,
                middle: RBTN & 0x04 != 0,
            });
            queue_report(); // re-arm for the next report
            return ev;
        }
        if !QUEUED {
            queue_report(); // keep exactly one IN read outstanding
        }
    }
    None
}

/// One-line status for the shell's `usb` command.
pub fn diagnose() -> alloc::string::String {
    let mut s = alloc::string::String::new();
    s.push_str(&alloc::format!(
        "xhci mmio={:#x} mapped={} ports={} mouse={} slot={}\n",
        unsafe { M },
        if crate::paging::is_mapped(unsafe { M }) { "yes" } else { "no" },
        unsafe { MAXPORTS },
        if MOUSEPRESENT.load(Ordering::Relaxed) {
            "yes"
        } else {
            "no"
        },
        SLOT.load(Ordering::Relaxed)
    ));
    unsafe {
        // Page-table walk for the MMIO base (physical == virtual here).
        let phys = M;
        let pml4 = crate::paging::cr3() as *const u64;
        let l4 = (phys >> 39) & 0x1FF;
        let l3 = (phys >> 30) & 0x1FF;
        let l2 = (phys >> 21) & 0x1FF;
        let l1 = (phys >> 12) & 0x1FF;
        let pml4e = pml4.add(l4).read_volatile();
        let pdpte = if pml4e & 1 != 0 {
            ((pml4e & 0x000f_ffff_ffff_f000) as *const u64).add(l3).read_volatile()
        } else { 0 };
        let pde = if pdpte & 1 != 0 && pdpte & (1 << 7) == 0 {
            ((pdpte & 0x000f_ffff_ffff_f000) as *const u64).add(l2).read_volatile()
        } else { 0 };
        let pte = if pde & 1 != 0 && pde & (1 << 7) == 0 {
            ((pde & 0x000f_ffff_ffff_f000) as *const u64).add(l1).read_volatile()
        } else { 0 };
        s.push_str(&alloc::format!(
            "  pt: pml4e={:#018x} pdpte={:#018x} pde={:#018x} pte={:#018x}\n",
            pml4e, pdpte, pde, pte
        ));
        for p in 0..8 {
            let psc = op(OP_PORTSC) + (p as usize) * 0x10;
            let v = rd(psc);
            s.push_str(&alloc::format!("  port{}: {:#010x}\n", p, v));
        }
    }
    if MOUSEPRESENT.load(Ordering::Relaxed) {
        s.push_str(&alloc::format!(
            "  hid ep-index={} mps={} armed={}\n",
            HID_EPID.load(Ordering::Relaxed),
            unsafe { HID_LEN },
            unsafe { QUEUED }
        ));
    }
    s
}