// ---------------------------------------------------------------------------
// rsdos UHCI (USB 1.1) companion driver — HID boot keyboard + mouse.
//
// Real Intel PCH (e.g. OptiPlex 7010) USB2 ports are handled by an EHCI
// (high-speed devices) plus its companion UHCI pair (full/low-speed
// devices).  A real USB keyboard/mouse is full- or low-speed, so it lives
// on the UHCI — the EHCI cannot talk to it without a root-hub TT (Intel
// PCH EHCI has none; the ports belong to the companion UHCI).  QEMU's
// virtual usb-mouse/usb-kbd are high-speed when attached to usb-ehci, so
// this class of real hardware was never exercised by the xHCI/EHCI paths.
//
// This driver is polled (no IRQ), exactly like ehci.rs: the frame list is
// walked every 1 ms by the controller; we keep one armed interrupt-IN TD
// per device and drain reports from the shell task (keyboard) and cursor
// task (mouse).  Both devices share one control pipeline, serialized by a
// spin lock so the shell and cursor tasks cannot corrupt each other's
// transfers.
//
// Data structures follow the UHCI spec as implemented by QEMU's hcd-uhci
// and Linux's uhci-hcd (they agree field-for-field):
//   TD:   link(0x00) ctrl(0x04) token(0x08) buffer(0x0C)
//   QH:   link(0x00) el_link(0x04)
//   ctrl: ACTIVE=1<<23 IOC=1<<24 IOS=1<<25 LS=1<<26 C_ERR=3<<27 SPD=1<<29
//         ACTLEN=bits 10:0 (n-1, written by the HC)
//   token: PID=bits 7:0 (full byte: SETUP=0x2D IN=0x69 OUT=0xE1)
//         DEVADDR=bits 14:8, EP=bits 18:15, TOGGLE=bit 19,
//         EXPLEN=bits 31:21 (n-1)
//   link ptrs: bit0=terminate, bit1=QH type, bit2=depth-first
// ---------------------------------------------------------------------------

use crate::ps2::MouseEvent;
use core::sync::atomic::{AtomicBool, AtomicU32, Ordering};

// ---- Registers (I/O space; 16-bit unless noted) ----
const USBCMD: u16 = 0x00;
const USBSTS: u16 = 0x02;
const USBINTR: u16 = 0x04;
const FRNUM: u16 = 0x06;
const FRBASEADD: u16 = 0x08; // 32-bit
const SOFMOD: u16 = 0x0C;
const PORTSC_BASE: u16 = 0x10; // PORTSC1, PORTSC2 (16-bit each)

const CMD_RUN: u16 = 1 << 0;
const CMD_HCRESET: u16 = 1 << 1;
const CMD_CF: u16 = 1 << 6;

const STS_HCHALTED: u16 = 1 << 5;

const PSC_CCS: u16 = 1 << 0; // current connect status
const PSC_CSC: u16 = 1 << 1; // connect change (write 1 to clear)
const PSC_PE: u16 = 1 << 2; // port enabled
const PSC_PEDC: u16 = 1 << 3; // enable change (write 1 to clear)
const PSC_LSDA: u16 = 1 << 8; // low-speed device attached
const PSC_R: u16 = 1 << 9; // reset
const PSC_SUSP: u16 = 1 << 12; // suspend

// ---- TD / QH fields ----
const TD_ACTIVE: u32 = 1 << 23;
const TD_IOC: u32 = 1 << 24;
const TD_LS: u32 = 1 << 26;
const TD_ERR: u32 = 3 << 27; // C_ERR = 3 retries
const TD_STALLED: u32 = 1 << 22;
const TD_BABBLE: u32 = 1 << 20;
const TD_TIMEO: u32 = 1 << 18;
const TD_ACTLEN: u32 = 0x7FF; // actual length, n-1

const TOK_DEV_SHIFT: u32 = 8;
const TOK_EP_SHIFT: u32 = 15;
const TOK_TOGGLE: u32 = 1 << 19;
const TOK_EXPLEN_SHIFT: u32 = 21;

const PID_SETUP: u32 = 0x2D;
const PID_IN: u32 = 0x69;
const PID_OUT: u32 = 0xE1;

const PTR_TERM: u32 = 0x1;
const PTR_QH: u32 = 0x2;

// ---- HID boot requests ----
const RT_CLASS_IFACE: u8 = 0x21;
const REQ_SET_IDLE: u8 = 0x0A;
const REQ_SET_PROTOCOL: u8 = 0x0B;

// ---- Static state ----
static IO_BASE: AtomicU32 = AtomicU32::new(0);
static N_PORTS: AtomicU32 = AtomicU32::new(0);
static CTRL_UP: AtomicBool = AtomicBool::new(false);
static MOUSE_PRES: AtomicBool = AtomicBool::new(false);
static KBD_ARMED: AtomicBool = AtomicBool::new(false);
static MOUSE_ARMED: AtomicBool = AtomicBool::new(false);
// Port index each device was found on (usize::MAX = none).  Used so the
// keyboard scan skips the mouse's port and vice versa.
static KBD_PORT: AtomicU32 = AtomicU32::new(u32::MAX);
static MOUSE_PORT: AtomicU32 = AtomicU32::new(u32::MAX);
static DEV_KBD: AtomicU32 = AtomicU32::new(0);
static DEV_MOUSE: AtomicU32 = AtomicU32::new(0);
static KBD_EP: AtomicU32 = AtomicU32::new(1);
static KBD_MPS: AtomicU32 = AtomicU32::new(8);
static MOUSE_EP: AtomicU32 = AtomicU32::new(1);
static MOUSE_MPS: AtomicU32 = AtomicU32::new(4);
// Low-speed flag per device (TD_LS needed for 1.5 Mb/s).
static KBD_LS: AtomicBool = AtomicBool::new(false);
static MOUSE_LS: AtomicBool = AtomicBool::new(false);
// Next USB device address to assign (1, 2, ...).
static DEV_SEQ: AtomicU32 = AtomicU32::new(1);
// Throttles for the lazy re-enumeration paths (ms).
static LAST_KBD_ENUM: AtomicU32 = AtomicU32::new(0);
static LAST_MOUSE_ENUM: AtomicU32 = AtomicU32::new(0);
// Serializes control transfers between the shell task (keyboard) and the
// cursor task (mouse).  Bounded spin; on timeout the transfer is skipped
// and retried on the next throttled pass.
static CTRL_LOCK: AtomicBool = AtomicBool::new(false);
// Toggle bit the device expects next (alternates after each report).
static KBD_TOGGLE: AtomicU32 = AtomicU32::new(0);
static MOUSE_TOGGLE: AtomicU32 = AtomicU32::new(0);

// ---- DMA: frame list + one struct page ----
static mut FRAMELIST: *mut u32 = core::ptr::null_mut();
static mut P: *mut u32 = core::ptr::null_mut(); // QHs, TDs, buffers

// Offsets within the struct page (16-byte-aligned QHs/TDs).
const OFF_CTRL_QH: usize = 0x000;
const OFF_KBD_QH: usize = 0x010;
const OFF_MOUSE_QH: usize = 0x020;
const OFF_SETUP_TD: usize = 0x030;
const OFF_DATA_TD: usize = 0x040;
const OFF_STAT_TD: usize = 0x050;
const OFF_KBD_TD: usize = 0x060;
const OFF_MOUSE_TD: usize = 0x070;
const OFF_REQ: usize = 0x100; // 8-byte setup packet
const OFF_DESC: usize = 0x200; // device descriptor (64 B)
const OFF_CFG: usize = 0x300; // config descriptor (256 B)
const OFF_KBD_REPORT: usize = 0x500; // 8 B
const OFF_MOUSE_REPORT: usize = 0x520; // 8 B

// Keyboard: previous report + pending scancode FIFO (fed to ps2 decoder).
static mut KBD_PREV: [u8; 8] = [0; 8];
const KBDQ_SIZE: usize = 64;
static mut KBDQ: [u8; KBDQ_SIZE] = [0; KBDQ_SIZE];
static KBDQ_HEAD: AtomicU32 = AtomicU32::new(0);
static KBDQ_LEN: AtomicU32 = AtomicU32::new(0);

// ---- port I/O ----
#[inline]
unsafe fn inw(port: u16) -> u16 {
    let v: u16;
    unsafe {
        core::arch::asm!("in ax, dx", out("ax") v, in("dx") port, options(nostack, preserves_flags));
    }
    v
}
#[inline]
unsafe fn outw(port: u16, v: u16) {
    unsafe {
        core::arch::asm!("out dx, ax", in("dx") port, in("ax") v, options(nostack, preserves_flags));
    }
}
#[inline]
unsafe fn outl(port: u16, v: u32) {
    unsafe {
        core::arch::asm!("out dx, eax", in("dx") port, in("eax") v, options(nostack, preserves_flags));
    }
}

fn io() -> u16 {
    IO_BASE.load(Ordering::Relaxed) as u16
}
/// Absolute I/O port of root port `p`'s PORTSC register.
fn prtsc(p: usize) -> u16 {
    io() + PORTSC_BASE + (p as u16) * 2
}

// ---- DMA helpers (mirror ehci.rs) ----
unsafe fn dma_alloc(bytes: usize) -> *mut u8 {
    unsafe {
        let pages = bytes.div_ceil(4096);
        let Some(phys) = crate::memory::alloc_pages(pages) else {
            return core::ptr::null_mut();
        };
        let p = phys as *mut u8;
        core::ptr::write_bytes(p, 0, pages * 4096);
        p
    }
}

unsafe fn rd(addr: *const u32) -> u32 {
    unsafe { addr.read_volatile() }
}
unsafe fn wr(addr: *mut u32, v: u32) {
    unsafe { addr.write_volatile(v) }
}

fn td_addr(o: usize) -> *mut u32 {
    unsafe { P.add(o / 4) }
}
fn buf_addr(o: usize) -> *mut u8 {
    unsafe { P.add(o / 4) as *mut u8 }
}

#[inline]
fn ms_now() -> u64 {
    crate::interrupts::uptime_ms()
}

/// Busy-wait for `ms` milliseconds (interrupts stay enabled; the tick keeps
/// the scheduler alive).
fn ms_spin(ms: u64) {
    let t0 = ms_now();
    while ms_now().wrapping_sub(t0) < ms {
        core::hint::spin_loop();
    }
}

/// Build an 8-byte USB setup packet.
fn ctrl_req(bm: u8, breq: u8, val: u16, idx: u16, len: u16) -> [u8; 8] {
    let mut r = [0u8; 8];
    r[0] = bm;
    r[1] = breq;
    r[2] = val as u8;
    r[3] = (val >> 8) as u8;
    r[4] = idx as u8;
    r[5] = (idx >> 8) as u8;
    r[6] = len as u8;
    r[7] = (len >> 8) as u8;
    r
}

// ---------------------------------------------------------------------------
// Control transfer pipeline (shared, lock-protected)
// ---------------------------------------------------------------------------

fn lock_ctrl() -> bool {
    for _ in 0..200_000 {
        if !CTRL_LOCK.swap(true, Ordering::Acquire) {
            return true;
        }
        core::hint::spin_loop();
    }
    false
}
fn unlock_ctrl() {
    CTRL_LOCK.store(false, Ordering::Release);
}

/// Issue one control transfer on EP0 of `dev`.  `mps` is the device's EP0
/// max packet, `low_speed` selects the TD_CTRL_LS bit, `req` the setup
/// packet, `dir_in` a device->host data phase, `buf/len` the data buffer.
/// Returns bytes received (IN) or 0 (OUT/setup-only), or a negative error.
/// Completion is polled (STATUS TD active bit) with a 5 s timeout.
unsafe fn ctl_xfer(
    dev: u32,
    _mps: u32,
    low_speed: bool,
    req: &[u8; 8],
    dir_in: bool,
    buf: *mut u8,
    len: usize,
) -> i32 {
    unsafe {
        // The controller DMA's the setup packet straight from the request
        // buffer — copy it there first.
        core::ptr::copy_nonoverlapping(req.as_ptr(), buf_addr(OFF_REQ), 8);
        core::sync::atomic::fence(Ordering::SeqCst);
        let ls = if low_speed { TD_LS } else { 0 };

        // SETUP TD (toggle 0, 8 bytes, always from the request buffer).
        let st = td_addr(OFF_SETUP_TD);
        wr(st.add(0), 0); // link: filled below
        wr(st.add(1), TD_ACTIVE | TD_ERR | ls);
        wr(
            st.add(2),
            PID_SETUP | (dev << TOK_DEV_SHIFT) | (7 << TOK_EXPLEN_SHIFT), // 8 bytes
        );
        wr(st.add(3), buf_addr(OFF_REQ) as u32);

        // STATUS TD (toggle 1, opposite direction of the data phase).
        let sp = td_addr(OFF_STAT_TD);
        let status_pid = if dir_in { PID_OUT } else { PID_IN };
        wr(sp.add(0), PTR_TERM);
        wr(sp.add(1), TD_ACTIVE | TD_ERR | ls);
        wr(
            sp.add(2),
            status_pid | (dev << TOK_DEV_SHIFT) | TOK_TOGGLE | (0 << TOK_EXPLEN_SHIFT),
        );
        wr(sp.add(3), buf_addr(OFF_REQ) as u32);

        if len == 0 {
            // Setup-only transfer: SETUP -> STATUS.
            wr(st.add(0), td_addr(OFF_STAT_TD) as u32);
        } else {
            // DATA TD (toggle 1).
            let dt = td_addr(OFF_DATA_TD);
            let dpid = if dir_in { PID_IN } else { PID_OUT };
            wr(dt.add(0), td_addr(OFF_STAT_TD) as u32);
            wr(dt.add(1), TD_ACTIVE | TD_ERR | ls);
            wr(
                dt.add(2),
                dpid | (dev << TOK_DEV_SHIFT)
                    | TOK_TOGGLE
                    | (((len as u32).saturating_sub(1) & 0x7FF) << TOK_EXPLEN_SHIFT),
            );
            wr(dt.add(3), buf as u32);
            wr(st.add(0), dt as u32);
        }

        // Point the control QH at the new chain and let the HC run it.
        core::sync::atomic::fence(Ordering::SeqCst);
        wr(td_addr(OFF_CTRL_QH).add(1), td_addr(OFF_SETUP_TD) as u32);
        core::sync::atomic::fence(Ordering::SeqCst);

        // Wait for completion (STATUS TD clears ACTIVE).
        let t0 = ms_now();
        loop {
            let c = rd(sp.add(1));
            if c & TD_ACTIVE == 0 {
                // Error checks: the status TD mirrors the chain's fate.
                if c & (TD_STALLED | TD_BABBLE | TD_TIMEO) != 0 {
                    return -2;
                }
                let got = if len > 0 && dir_in {
                    ((rd(td_addr(OFF_DATA_TD).add(1)) & TD_ACTLEN) + 1) as i32
                } else {
                    0
                };
                return got.min(len as i32);
            }
            if ms_now().wrapping_sub(t0) >= 5000 {
                return -1; // timeout
            }
            core::hint::spin_loop();
        }
    }
}

// ---------------------------------------------------------------------------
// Port management
// ---------------------------------------------------------------------------

/// Reset a connected root port and enable it.  Returns the low-speed flag
/// (false = full-speed) on success.
unsafe fn reset_port(p: usize) -> Option<bool> {
    unsafe {
        let r = prtsc(p);
        // Clear connect/enable change bits (write 1s).
        outw(r, PSC_CSC | PSC_PEDC);
        // Start reset.
        outw(r, PSC_R);
        ms_spin(50);
        // End reset.
        outw(r, 0);
        // Enable the port (device must be connected).
        outw(r, PSC_PE);
        let t0 = ms_now();
        while ms_now().wrapping_sub(t0) < 100 {
            let psc = inw(r);
            if psc & PSC_PE != 0 {
                return Some(psc & PSC_LSDA != 0);
            }
            ms_spin(1);
        }
        None
    }
}

// ---------------------------------------------------------------------------
// Config-descriptor parsing: find a HID boot interface with the wanted
// protocol and its first interrupt IN endpoint.
// ---------------------------------------------------------------------------

struct HidEp {
    epnum: u8,
    mps: u16,
    iface: u8,
}

/// Walk a config descriptor for a boot-protocol HID interface (class 3,
/// subclass 1) whose protocol matches `want` (1 = keyboard, 2 = mouse),
/// returning its first interrupt IN endpoint.
fn find_hid_ep(cfg: &[u8], want: u8) -> Option<HidEp> {
    let mut i = 0;
    while i + 1 < cfg.len() && cfg[i] != 0 {
        let blen = cfg[i] as usize;
        let btype = cfg[i + 1];
        if blen < 2 || i + blen > cfg.len() {
            return None;
        }
        if btype == 4 {
            // Interface descriptor.
            if blen >= 9
                && cfg[i + 5] == 3 // bInterfaceClass = HID
                && cfg[i + 6] == 1 // bInterfaceSubClass = boot
                && cfg[i + 7] == want
            // bInterfaceProtocol
            {
                let iface = cfg[i + 2];
                // Scan this interface's endpoints.
                let mut j = i + blen;
                while j + 1 < cfg.len() && cfg[j] != 0 {
                    let elen = cfg[j] as usize;
                    let etype = cfg[j + 1];
                    if etype == 4 {
                        break; // next interface
                    }
                    if etype == 5 && elen >= 7 {
                        let epa = cfg[j + 2];
                        let attrs = cfg[j + 3];
                        if epa & 0x80 != 0 && attrs & 0x03 == 0x03 {
                            let mps = (cfg[j + 4] as u16) | ((cfg[j + 5] as u16) << 8);
                            return Some(HidEp {
                                epnum: epa & 0x0F,
                                mps: mps.max(1),
                                iface,
                            });
                        }
                    }
                    if elen < 2 {
                        return None;
                    }
                    j += elen;
                }
            }
        }
        i += blen;
    }
    None
}

// ---------------------------------------------------------------------------
// Device enumeration (one root port)
// ---------------------------------------------------------------------------

/// Enumerate the device on root port `p` and arm the interrupt-IN endpoint
/// if it is a HID boot device of protocol `want` (1 = keyboard, 2 = mouse).
/// Returns true when armed.
unsafe fn enumerate_port(p: usize, want: u8) -> bool {
    unsafe {
        let Some(low_speed) = reset_port(p) else {
            return false;
        };

        // 1. GET_DESCRIPTOR(device) at address 0 — 8 bytes for EP0 MPS.
        let dd = buf_addr(OFF_DESC);
        core::ptr::write_bytes(dd, 0, 64);
        let req = ctrl_req(0x80, 6, 0x0100, 0, 8);
        if ctl_xfer(0, 8, low_speed, &req, true, dd, 8) < 0 {
            crate::serial_println!("uhci: port {} get-dev-desc failed (want {})", p, want);
            return false;
        }
        let mps0 = core::ptr::read_volatile(dd.add(7)).max(8) as u32;

        // 2. SET_ADDRESS.
        let dev = DEV_SEQ.fetch_add(1, Ordering::Relaxed);
        let req = ctrl_req(0x00, 5, dev as u16, 0, 0);
        if ctl_xfer(0, mps0, low_speed, &req, false, core::ptr::null_mut(), 0) < 0 {
            crate::serial_println!("uhci: set-address failed on port {}", p);
            return false;
        }

        // 3. Full device descriptor at the new address.
        core::ptr::write_bytes(dd, 0, 64);
        let req = ctrl_req(0x80, 6, 0x0100, 0, 18);
        let got = ctl_xfer(dev, mps0, low_speed, &req, true, dd, 18);
        if got < 0 {
            crate::serial_println!("uhci: re-get dev-desc failed on port {}", p);
            return false;
        }
        let class = core::ptr::read_volatile(dd.add(4));
        let _ = class; // HID devices have device class 0x00; the HID class
                       // lives in the interface descriptor (checked below).

        // 4. Full config descriptor.
        let cfg = buf_addr(OFF_CFG);
        core::ptr::write_bytes(cfg, 0, 256);
        let req = ctrl_req(0x80, 6, 0x0200, 0, 256);
        let got = ctl_xfer(dev, mps0, low_speed, &req, true, cfg, 256);
        if got < 0 {
            crate::serial_println!(
                "uhci: get-config failed on port {} code={} (data ctrl={:#x} stat ctrl={:#x})",
                p,
                got,
                rd(td_addr(OFF_DATA_TD).add(1)),
                rd(td_addr(OFF_STAT_TD).add(1))
            );
            return false;
        }
        let Some(ep) = find_hid_ep(core::slice::from_raw_parts(cfg, got as usize), want) else {
            return false; // not the device we're looking for
        };

        // 5. SET_CONFIGURATION(1).
        let req = ctrl_req(0x00, 9, 1, 0, 0);
        if ctl_xfer(dev, mps0, low_speed, &req, false, core::ptr::null_mut(), 0) < 0 {
            crate::serial_println!("uhci: set-config failed on port {}", p);
            return false;
        }

        // 6. Boot protocol + no idle.  Some devices stall these; tolerate.
        let req = ctrl_req(RT_CLASS_IFACE, REQ_SET_PROTOCOL, 1, ep.iface as u16, 0);
        if ctl_xfer(dev, mps0, low_speed, &req, false, core::ptr::null_mut(), 0) < 0 {
            crate::serial_println!("uhci: set-protocol stalled on port {}", p);
        }
        let req = ctrl_req(RT_CLASS_IFACE, REQ_SET_IDLE, 0, ep.iface as u16, 0);
        if ctl_xfer(dev, mps0, low_speed, &req, false, core::ptr::null_mut(), 0) < 0 {
            crate::serial_println!("uhci: set-idle stalled on port {}", p);
        }

        // 7. Arm the interrupt-IN TD.
        arm_in_td(want, dev, &ep, low_speed, p);
        crate::serial_println!(
            "uhci: HID boot {} armed (port {} dev {} ep{}({:#x}) mps={} {})",
            if want == 1 { "keyboard" } else { "mouse" },
            p,
            dev,
            ep.epnum,
            if want == 1 { 0x81 } else { 0x81 },
            ep.mps,
            if low_speed { "low-speed" } else { "full-speed" }
        );
        true
    }
}

/// Build the interrupt-IN TD for a device and point its QH element at it.
unsafe fn arm_in_td(want: u8, dev: u32, ep: &HidEp, low_speed: bool, port: usize) {
    unsafe {
        let (qho, tdo, buf_off, mps, is_kbd) = if want == 1 {
            (
                OFF_KBD_QH,
                OFF_KBD_TD,
                OFF_KBD_REPORT,
                (ep.mps as u32).max(8),
                true,
            )
        } else {
            (
                OFF_MOUSE_QH,
                OFF_MOUSE_TD,
                OFF_MOUSE_REPORT,
                (ep.mps as u32).max(4),
                false,
            )
        };
        let qh = td_addr(qho);
        let td = td_addr(tdo);
        let ls = if low_speed { TD_LS } else { 0 };
        let explen = if is_kbd { 8 } else { 3 };
        let toggle = 0; // endpoint toggles start at 0 after SET_CONFIG

        wr(td.add(0), PTR_TERM);
        wr(td.add(1), TD_ACTIVE | TD_ERR | TD_IOC | ls);
        wr(
            td.add(2),
            PID_IN
                | (dev << TOK_DEV_SHIFT)
                | ((ep.epnum as u32) << TOK_EP_SHIFT)
                | (toggle << 19)
                | (explen << TOK_EXPLEN_SHIFT),
        );
        wr(td.add(3), buf_addr(buf_off) as u32);
        // QH: el_link points at the TD (link/el_link written last).
        wr(qh.add(0), PTR_TERM); // no next QH (per-device QH at chain end)
        core::sync::atomic::fence(Ordering::SeqCst);
        wr(qh.add(1), td as u32);
        core::sync::atomic::fence(Ordering::SeqCst);

        match is_kbd {
            true => {
                KBD_ARMED.store(true, Ordering::Relaxed);
                KBD_PORT.store(port as u32, Ordering::Relaxed);
                DEV_KBD.store(dev, Ordering::Relaxed);
                KBD_EP.store(ep.epnum as u32, Ordering::Relaxed);
                KBD_MPS.store(mps, Ordering::Relaxed);
                KBD_LS.store(low_speed, Ordering::Relaxed);
                KBD_TOGGLE.store(0, Ordering::Relaxed);
            }
            false => {
                MOUSE_ARMED.store(true, Ordering::Relaxed);
                MOUSE_PORT.store(port as u32, Ordering::Relaxed);
                DEV_MOUSE.store(dev, Ordering::Relaxed);
                MOUSE_EP.store(ep.epnum as u32, Ordering::Relaxed);
                MOUSE_MPS.store(mps, Ordering::Relaxed);
                MOUSE_LS.store(low_speed, Ordering::Relaxed);
                MOUSE_TOGGLE.store(0, Ordering::Relaxed);
                MOUSE_PRES.store(true, Ordering::Relaxed);
            }
        }
    }
}

/// Scan all root ports and arm the first device of the wanted protocol.
unsafe fn arm_device(want: u8) -> bool {
    unsafe {
        if !lock_ctrl() {
            return false;
        }
        let np = N_PORTS.load(Ordering::Relaxed) as usize;
        let other = if want == 1 {
            MOUSE_PORT.load(Ordering::Relaxed)
        } else {
            KBD_PORT.load(Ordering::Relaxed)
        };
        for p in 0..np {
            if p as u32 == other {
                continue; // claimed by the other device
            }
            let psc = inw(prtsc(p));
            if psc & PSC_CCS == 0 {
                continue;
            }
            if enumerate_port(p, want) {
                unlock_ctrl();
                return true;
            }
        }
        unlock_ctrl();
        false
    }
}

// ---------------------------------------------------------------------------
// Keyboard: HID report -> scancode-set-1 make/break bytes -> ps2 decoder
// ---------------------------------------------------------------------------

/// HID usage -> scancode set 1 (make code).  High byte 0xE0 = E0-prefixed.
/// 0 = unmapped.  Table covers the common keyboard usages (0x04..0x63).
static HID_TO_SCAN: [u16; 0x64] = {
    let mut t = [0u16; 0x64];
    let mut i = 0;
    while i < 0x64 {
        t[i] = 0;
        i += 1;
    }
    // Letters a-z
    t[0x04] = 0x1E;
    t[0x05] = 0x30;
    t[0x06] = 0x2E;
    t[0x07] = 0x20;
    t[0x08] = 0x12;
    t[0x09] = 0x21;
    t[0x0A] = 0x22;
    t[0x0B] = 0x23;
    t[0x0C] = 0x17;
    t[0x0D] = 0x24;
    t[0x0E] = 0x25;
    t[0x0F] = 0x26;
    t[0x10] = 0x32;
    t[0x11] = 0x31;
    t[0x12] = 0x18;
    t[0x13] = 0x19;
    t[0x14] = 0x10;
    t[0x15] = 0x13;
    t[0x16] = 0x1F;
    t[0x17] = 0x14;
    t[0x18] = 0x16;
    t[0x19] = 0x2F;
    t[0x1A] = 0x11;
    t[0x1B] = 0x2D;
    t[0x1C] = 0x15;
    t[0x1D] = 0x2C;
    // Digits + punctuation
    t[0x1E] = 0x02;
    t[0x1F] = 0x03;
    t[0x20] = 0x04;
    t[0x21] = 0x05;
    t[0x22] = 0x06;
    t[0x23] = 0x07;
    t[0x24] = 0x08;
    t[0x25] = 0x09;
    t[0x26] = 0x0A;
    t[0x27] = 0x0B;
    t[0x28] = 0x1C; // Enter
    t[0x29] = 0x01; // Esc
    t[0x2A] = 0x0E; // Backspace
    t[0x2B] = 0x0F; // Tab
    t[0x2C] = 0x39; // Space
    t[0x2D] = 0x0C; // -
    t[0x2E] = 0x0D; // =
    t[0x2F] = 0x1A; // [
    t[0x30] = 0x1B; // ]
    t[0x31] = 0x2B; // backslash
    t[0x33] = 0x27; // ;
    t[0x34] = 0x28; // '
    t[0x35] = 0x29; // `
    t[0x36] = 0x33; // ,
    t[0x37] = 0x34; // .
    t[0x38] = 0x35; // /
    t[0x39] = 0x3A; // CapsLock
                    // F-keys
    t[0x3A] = 0x3B; // F1
    t[0x3B] = 0x3C;
    t[0x3C] = 0x3D;
    t[0x3D] = 0x3E;
    t[0x3E] = 0x3F;
    t[0x3F] = 0x40;
    t[0x40] = 0x41;
    t[0x41] = 0x42;
    t[0x42] = 0x43;
    t[0x43] = 0x44; // F10
    t[0x44] = 0xE057; // F11
    t[0x45] = 0xE058; // F12
                      // Nav cluster
    t[0x49] = 0xE052; // Insert
    t[0x4A] = 0xE047; // Home
    t[0x4B] = 0xE049; // PageUp
    t[0x4C] = 0xE053; // Delete
    t[0x4D] = 0xE04F; // End
    t[0x4E] = 0xE051; // PageDown
    t[0x4F] = 0xE04D; // Right
    t[0x50] = 0xE04B; // Left
    t[0x51] = 0xE050; // Down
    t[0x52] = 0xE048; // Up
                      // Keypad
    t[0x53] = 0x45; // NumLock
    t[0x54] = 0xE035; // /
    t[0x55] = 0x37; // *
    t[0x56] = 0x4A; // -
    t[0x57] = 0x4E; // +
    t[0x58] = 0xE01C; // Enter
    t[0x59] = 0x4F; // 1
    t[0x5A] = 0x50; // 2
    t[0x5B] = 0x51; // 3
    t[0x5C] = 0x4B; // 4
    t[0x5D] = 0x4C; // 5
    t[0x5E] = 0x4D; // 6
    t[0x5F] = 0x47; // 7
    t[0x60] = 0x48; // 8
    t[0x61] = 0x49; // 9
    t[0x62] = 0x52; // 0
    t[0x63] = 0x53; // .
    t
};

/// Modifier bit -> scancode (set 1, high byte 0xE0 = E0-prefixed).
static MOD_TO_SCAN: [u16; 8] = [
    0x1D,   // 0 LCtrl
    0x2A,   // 1 LShift
    0x38,   // 2 LAlt
    0xE05B, // 3 LGui
    0xE01D, // 4 RCtrl
    0x36,   // 5 RShift
    0xE038, // 6 RAlt
    0xE05C, // 7 RGui
];

fn kbdq_push(b: u8) {
    unsafe {
        let h = KBDQ_HEAD.load(Ordering::Relaxed) as usize;
        let l = KBDQ_LEN.load(Ordering::Relaxed) as usize;
        if l >= KBDQ_SIZE {
            return;
        }
        KBDQ[(h + l) % KBDQ_SIZE] = b;
        KBDQ_LEN.store((l + 1) as u32, Ordering::Relaxed);
    }
}
fn kbdq_pop() -> Option<u8> {
    unsafe {
        let l = KBDQ_LEN.load(Ordering::Relaxed) as usize;
        if l == 0 {
            return None;
        }
        let h = KBDQ_HEAD.load(Ordering::Relaxed) as usize;
        let b = KBDQ[h];
        KBDQ_HEAD.store(((h + 1) % KBDQ_SIZE) as u32, Ordering::Relaxed);
        KBDQ_LEN.store((l - 1) as u32, Ordering::Relaxed);
        Some(b)
    }
}

/// Push the make (press) or break (release) scancode sequence for a
/// scancode-map entry (0 = unmapped, 0x00XX = plain, 0xE0XX = E0-prefixed).
fn push_scancode(scan: u16, press: bool) {
    if scan == 0 {
        return;
    }
    if scan & 0xFF00 == 0xE000 {
        kbdq_push(0xE0);
        let code = (scan & 0xFF) as u8;
        kbdq_push(if press { code } else { code | 0x80 });
    } else {
        let code = scan as u8;
        kbdq_push(if press { code } else { code | 0x80 });
    }
}

/// Diff the completed HID report against the previous one and queue the
/// scancode transitions (make/break) for the shared PS/2 decoder.
unsafe fn diff_report(report: &[u8]) {
    unsafe {
        let prev = &mut *core::ptr::addr_of_mut!(KBD_PREV);
        let new_mods = report[0];
        let old_mods = prev[0];
        for bit in 0..8 {
            let new_set = new_mods & (1 << bit) != 0;
            let old_set = old_mods & (1 << bit) != 0;
            if new_set != old_set {
                push_scancode(MOD_TO_SCAN[bit], new_set);
            }
        }
        let new_keys = &report[2..8];
        let old_keys = &prev[2..8];
        // Releases first (old keys not in new set).
        for &k in old_keys {
            if k == 0 || new_keys.contains(&k) {
                continue;
            }
            if k as usize >= HID_TO_SCAN.len() {
                continue;
            }
            push_scancode(HID_TO_SCAN[k as usize], false);
        }
        // Then presses (new keys not in the old set).
        for &k in new_keys {
            if k == 0 || old_keys.contains(&k) {
                continue;
            }
            if k as usize >= HID_TO_SCAN.len() {
                continue;
            }
            push_scancode(HID_TO_SCAN[k as usize], true);
        }
        prev.copy_from_slice(report);
    }
}

/// Re-arm a completed interrupt-IN TD (one-outstanding report pattern).
unsafe fn rearm_td(qho: usize, tdo: usize, toggle: u32, ls: bool, mps: usize, explen: u32) {
    unsafe {
        let qh = td_addr(qho);
        let td = td_addr(tdo);
        let t = rd(td.add(2));
        let dev = (t >> TOK_DEV_SHIFT) & 0x7F;
        let ep = (t >> TOK_EP_SHIFT) & 0xF;
        wr(
            td.add(1),
            TD_ACTIVE | TD_ERR | TD_IOC | if ls { TD_LS } else { 0 },
        );
        wr(
            td.add(2),
            PID_IN
                | (dev << TOK_DEV_SHIFT)
                | (ep << TOK_EP_SHIFT)
                | (toggle << 19)
                | (explen << TOK_EXPLEN_SHIFT),
        );
        core::sync::atomic::fence(Ordering::SeqCst);
        wr(qh.add(1), td as u32);
        core::sync::atomic::fence(Ordering::SeqCst);
        let _ = mps;
    }
}

// ---------------------------------------------------------------------------
// Public interface
// ---------------------------------------------------------------------------

pub fn controller_up() -> bool {
    CTRL_UP.load(Ordering::Relaxed)
}
pub fn mouse_present() -> bool {
    MOUSE_PRES.load(Ordering::Relaxed)
}

/// Find the first UHCI companion controller (class 0x0C, subclass 0x03,
/// prog_if 0x00).
fn find_uhci() -> Option<(u8, u8, u8)> {
    for d in crate::pci::probe() {
        if d.class == 0x0C && d.subclass == 0x03 && d.prog_if == 0x00 {
            return Some((d.bus, d.device, d.func));
        }
    }
    None
}

/// Bring up the first UHCI controller: claim it, reset it, run the frame
/// list, and enumerate any connected boot HID devices.
pub fn init() -> bool {
    let Some((b, d, f)) = find_uhci() else {
        crate::serial_println!("uhci: no UHCI controller");
        return false;
    };
    crate::serial_println!("uhci: USB1.1 UHCI at {:02x}:{:02x}.{}", b, d, f);

    unsafe {
        // I/O space + bus master (polled: no INTx assertion needed, but
        // clear it anyway so a stray legacy IRQ can never storm).
        let cmd = crate::pci::config_read(b, d, f, 0x04);
        crate::pci::config_write(
            b,
            d,
            f,
            0x04,
            (cmd & 0xFFFF_0000) | (cmd & 0xFFFF) | 0x5 | 0x400,
        );
        let mut base: u32 = 0;
        // BAR0 first (Intel ICH), then BAR4 (PIIX3/PIIX4-style).  Take the
        // first nonzero I/O BAR.
        for bar_idx in [0u8, 4] {
            let bar = crate::pci::config_read(b, d, f, 0x10 + bar_idx * 4);
            let cand = bar & !0x3;
            if cand != 0 {
                base = cand;
                break;
            }
        }
        if base == 0 {
            crate::serial_println!("uhci: no I/O BAR");
            return false;
        }
        IO_BASE.store(base, Ordering::Relaxed);

        // Controller reset: stop, then HCRESET (UHCI clears HCRESET itself
        // within microseconds).
        outw(io(), 0);
        outw(io() + USBCMD, CMD_HCRESET);
        ms_spin(5);
        let mut t = 0;
        while t < 1_000_000 && inw(io() + USBCMD) & CMD_HCRESET != 0 {
            t += 1;
            core::hint::spin_loop();
        }
        outw(io() + USBINTR, 0);

        // DMA: 4 KiB frame list (1024 entries) + struct page.
        FRAMELIST = dma_alloc(4096) as *mut u32;
        P = dma_alloc(4096) as *mut u32;
        if FRAMELIST.is_null() || P.is_null() {
            crate::serial_println!("uhci: DMA alloc failed");
            return false;
        }

        // Wire the schedule: every frame entry -> CTRL_QH -> KBD_QH ->
        // MOUSE_QH -> terminate.  Each QH's element is set when its device
        // arms; the control QH element is set per transfer.
        let fl = FRAMELIST;
        for i in 0..1024 {
            wr(fl.add(i), (td_addr(OFF_CTRL_QH) as u32) | PTR_QH);
        }
        wr(
            td_addr(OFF_CTRL_QH).add(0),
            (td_addr(OFF_KBD_QH) as u32) | PTR_QH,
        );
        wr(td_addr(OFF_CTRL_QH).add(1), PTR_TERM);
        wr(
            td_addr(OFF_KBD_QH).add(0),
            (td_addr(OFF_MOUSE_QH) as u32) | PTR_QH,
        );
        wr(td_addr(OFF_KBD_QH).add(1), PTR_TERM);
        wr(td_addr(OFF_MOUSE_QH).add(0), PTR_TERM);
        wr(td_addr(OFF_MOUSE_QH).add(1), PTR_TERM);
        core::sync::atomic::fence(Ordering::SeqCst);

        // Frame list base, configure, run.
        outl(io() + FRBASEADD, fl as u32);
        outw(io() + SOFMOD, 64);
        outw(io() + FRNUM, 0);
        outw(io() + USBCMD, CMD_CF);
        outw(io() + USBCMD, CMD_CF | CMD_RUN);
        ms_spin(2);
        if inw(io() + USBSTS) & STS_HCHALTED != 0 {
            crate::serial_println!("uhci: controller failed to run");
            return false;
        }

        // Port count: UHCI has 2 ports at 0x10/0x12; a few have more.
        let mut np = 0;
        while np < 4 {
            let psc = inw(prtsc(np));
            if psc == 0xFFFF || psc == 0xff7f {
                break;
            }
            np += 1;
        }
        N_PORTS.store(np as u32, Ordering::Relaxed);
        CTRL_UP.store(true, Ordering::Relaxed);
        crate::serial_println!("uhci: controller up (io={:#x} ports={})", base, np);

        // Enumerate a boot keyboard first (input priority), then a mouse.
        crate::serial_println!(
            "uhci: dbg P={:#x} FL={:#x} ctrlqh={:#x} el={:#x} fl0={:#x} setup_td={:#x} setup_link={:#x}",
            P as usize,
            FRAMELIST as usize,
            td_addr(OFF_CTRL_QH) as usize,
            rd(td_addr(OFF_CTRL_QH).add(1)),
            rd(FRAMELIST.add(0)),
            td_addr(OFF_SETUP_TD) as usize,
            rd(td_addr(OFF_SETUP_TD).add(0))
        );
        let _ = arm_device(1);
        let _ = arm_device(2);
        true
    }
}

/// Non-blocking keyboard poll: drains the scancode FIFO (fed by completed
/// HID reports) through the shared PS/2 decoder.  Lazy re-enumeration is
/// throttled to one attempt per second.
pub fn try_poll_key() -> Option<crate::ps2::KeyEvent> {
    unsafe {
        // Fast path: pending scancodes first.
        if let Some(b) = kbdq_pop() {
            return crate::ps2::feed_scancode(b);
        }
        if KBD_ARMED.load(Ordering::Relaxed) {
            let td = td_addr(OFF_KBD_TD);
            let c = rd(td.add(1));
            if c & TD_ACTIVE == 0 {
                if c & (TD_STALLED | TD_BABBLE | TD_TIMEO) != 0 {
                    // Transfer error: drop the arm; bounded re-enum follows.
                    KBD_ARMED.store(false, Ordering::Relaxed);
                    return None;
                }
                // Completed: diff the report, re-arm, then drain the FIFO.
                let rep = buf_addr(OFF_KBD_REPORT);
                let mut report = [0u8; 8];
                for (i, b) in report.iter_mut().enumerate() {
                    *b = core::ptr::read_volatile(rep.add(i));
                }
                diff_report(&report);
                let tog = KBD_TOGGLE.load(Ordering::Relaxed) ^ 1;
                KBD_TOGGLE.store(tog, Ordering::Relaxed);
                rearm_td(
                    OFF_KBD_QH,
                    OFF_KBD_TD,
                    tog,
                    KBD_LS.load(Ordering::Relaxed),
                    KBD_MPS.load(Ordering::Relaxed) as usize,
                    8,
                );
                return kbdq_pop().and_then(crate::ps2::feed_scancode);
            }
            return None;
        }
        // Not armed: bring the keyboard up lazily, never hot-loop.
        if CTRL_UP.load(Ordering::Relaxed) {
            let now = crate::interrupts::uptime_ms() as u32;
            if now.wrapping_sub(LAST_KBD_ENUM.load(Ordering::Relaxed)) >= 1000 {
                LAST_KBD_ENUM.store(now, Ordering::Relaxed);
                let _ = arm_device(1);
            }
        }
        kbdq_pop().and_then(crate::ps2::feed_scancode)
    }
}

/// Poll for a mouse report (mirrors ehci::poll_mouse): hot path reads the
/// armed TD's ACTIVE bit in DMA memory; completion decodes the report and
/// re-arms; a stalled/TD-error transfer drops the arm for a bounded
/// re-enumeration.
pub fn poll_mouse() -> Option<MouseEvent> {
    unsafe {
        if MOUSE_ARMED.load(Ordering::Relaxed) {
            let td = td_addr(OFF_MOUSE_TD);
            let c = rd(td.add(1));
            if c & TD_ACTIVE == 0 {
                if c & (TD_STALLED | TD_BABBLE | TD_TIMEO) != 0 {
                    MOUSE_ARMED.store(false, Ordering::Relaxed);
                    return None;
                }
                let hb = buf_addr(OFF_MOUSE_REPORT);
                let b0 = core::ptr::read_volatile(hb);
                let b1 = core::ptr::read_volatile(hb.add(1));
                let b2 = core::ptr::read_volatile(hb.add(2));
                let ev = Some(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,
                });
                let tog = MOUSE_TOGGLE.load(Ordering::Relaxed) ^ 1;
                MOUSE_TOGGLE.store(tog, Ordering::Relaxed);
                rearm_td(
                    OFF_MOUSE_QH,
                    OFF_MOUSE_TD,
                    tog,
                    MOUSE_LS.load(Ordering::Relaxed),
                    MOUSE_MPS.load(Ordering::Relaxed) as usize,
                    3,
                );
                return ev;
            }
            return None;
        }
        if CTRL_UP.load(Ordering::Relaxed) {
            let now = crate::interrupts::uptime_ms() as u32;
            if now.wrapping_sub(LAST_MOUSE_ENUM.load(Ordering::Relaxed)) >= 1000 {
                LAST_MOUSE_ENUM.store(now, Ordering::Relaxed);
                let _ = arm_device(2);
            }
        }
        None
    }
}

/// One-line status for the shell `usb` command.
pub fn diagnose() -> alloc::string::String {
    let mut s = alloc::string::String::new();
    s.push_str(&alloc::format!(
        "uhci io={:#x} ports={} kbd={} mouse={} kbd_dev={} mouse_dev={}\n",
        IO_BASE.load(Ordering::Relaxed),
        N_PORTS.load(Ordering::Relaxed),
        if KBD_ARMED.load(Ordering::Relaxed) {
            "armed"
        } else {
            "-"
        },
        if MOUSE_ARMED.load(Ordering::Relaxed) {
            "armed"
        } else {
            "-"
        },
        DEV_KBD.load(Ordering::Relaxed),
        DEV_MOUSE.load(Ordering::Relaxed)
    ));
    let np = N_PORTS.load(Ordering::Relaxed) as usize;
    for p in 0..np.min(4) {
        let v = unsafe { inw(prtsc(p)) };
        s.push_str(&alloc::format!(
            "  uport{}: {}{}{}{:#x}\n",
            p,
            if v & PSC_CCS != 0 { "conn " } else { "none " },
            if v & PSC_PE != 0 { "en " } else { "" },
            if v & PSC_LSDA != 0 { "low " } else { "" },
            v
        ));
    }
    s
}