// ---------------------------------------------------------------------------
// AHCI — SATA host-controller driver (polling, DMA, no interrupts).
//
// Finds the chipset's AHCI controller (class 1 / subclass 6 / prog-if 1,
// e.g. the ICH9 on q35) on PCI bus 0, maps its ABAR (BAR5), starts every
// implemented port that has a device and exposes sector-level reads/writes
// for the FAT32 filesystem.  Multiple disks are supported; the filesystem
// selects one via `select()` and tries each in turn until one mounts.
//
// Commands are issued through command-slot 0 with one 4 KiB DMA buffer per
// port; completion is polled via PxCI (no ISR, no MSI).  Everything is
// identity-mapped, so "physical" DMA addresses are plain kernel pointers.
// The display is never involved, so this is a straightforward polling
// loop like the rest of the single-threaded kernel.
// ---------------------------------------------------------------------------

use crate::pci;

const SECTOR: usize = 512;
const MAX_DISKS: usize = 8;

// ABAR (HBA memory space) offsets.
const CAP: usize = 0x00;
const GHC: usize = 0x04;
const PI: usize = 0x0C;
const PORT_BASE: usize = 0x100;
const PORT_STRIDE: usize = 0x80;

// Port register offsets (relative to the port's base).
const P_CLB: usize = 0x00;
const P_CLBU: usize = 0x04;
const P_FB: usize = 0x08;
const P_FBU: usize = 0x0C;
const P_IS: usize = 0x10;
const P_CMD: usize = 0x18;
const P_TFD: usize = 0x20;
const P_SSTS: usize = 0x28;
const P_SCTL: usize = 0x2C;
const P_CI: usize = 0x38;

const CMD_ST: u32 = 1 << 0; // start
const CMD_SUD: u32 = 1 << 1; // spin-up device
const CMD_FRE: u32 = 1 << 4; // FIS receive enable
const CMD_DMAE: u32 = 1 << 23; // DMA enable (QEMU/ICH9) — required or DMA commands silently transfer nothing
const CMD_FR: u32 = 1 << 14; // FIS receive running
const CMD_CR: u32 = 1 << 15; // command list running

const ATA_IDENTIFY: u8 = 0xEC;
const ATA_READ_DMA_EXT: u8 = 0x25;
const ATA_WRITE_DMA_EXT: u8 = 0x35;

/// One attached disk: register map + DMA structures.
#[derive(Clone, Copy)]
struct Disk {
    port: usize,   // port register base = ABAR + 0x100
    port_idx: u32, // implemented-port index this disk lives on (rescan bookkeeping)
    clb: usize,    // command list, 1024 B (32 × 32 B), 256 B aligned
    ct: usize,     // command table, 256 B, 128 B aligned
    buf: usize,    // 4 KiB DMA data buffer (page aligned)
    sectors: u64,  // total sectors (LBA48 from IDENTIFY)
}

static mut DISKS: [Option<Disk>; MAX_DISKS] = [None; MAX_DISKS];
static mut DISK_COUNT: usize = 0;
static mut ACTIVE: usize = 0;

// Controller state kept for runtime rescan: ABAR + the implemented-port
// bitmap, plus the last-known device presence (PxSSTS.DET == 3) per port.
// `poll_dirty` compares live DET against this, so a hotplug is noticed
// with a handful of MMIO reads; the full re-scan only runs on a change.
const MAX_PORTS: usize = 32;
static mut ABAR: usize = 0;
static mut NP: u32 = 0;
static mut IMPL_PORTS: u32 = 0;
static mut PORT_UP: [bool; MAX_PORTS] = [false; MAX_PORTS];

fn reg(a: usize) -> u32 {
    unsafe { core::ptr::read_volatile(a as *const u32) }
}

fn wreg(a: usize, v: u32) {
    unsafe { core::ptr::write_volatile(a as *mut u32, v) }
}

/// Scan bus 0 for the AHCI controller (SATA, AHCI prog-if).
fn find_ahci() -> Option<(u8, u8)> {
    for dev in 0..32u8 {
        for func in 0..8u8 {
            let id = pci::config_read(0, dev, func, 0);
            if id == 0xFFFF_FFFF {
                if func == 0 {
                    break;
                }
                continue;
            }
            let cc = pci::config_read(0, dev, func, 0x08);
            let class = (cc >> 24) as u8;
            let subclass = (cc >> 16) as u8;
            let prog_if = (cc >> 8) as u8;
            if class == 0x01 && subclass == 0x06 && prog_if == 0x01 {
                return Some((dev, func));
            }
        }
    }
    None
}

/// Bring up one port: stop any firmware state, wire up its DMA structures,
/// wait for the link, start the command engine and run IDENTIFY.  Returns
/// the disk if a usable device answered.
fn start_port(abar: usize, port_idx: u32) -> Option<Disk> {
    let port = abar + PORT_BASE + port_idx as usize * PORT_STRIDE;

    // DMA structures: one 8 KiB allocation → page 0 = cmd list + FIS + cmd
    // table, page 1 = data buffer.
    let mem = crate::memory::alloc_pages(2)?;
    let clb = mem;
    let fis = mem + 1024;
    let ct = mem + 1280;
    let buf = mem + 4096;

    // Stop the port if firmware left it running.
    let mut cmd = reg(port + P_CMD);
    if cmd & CMD_ST != 0 {
        wreg(port + P_CMD, cmd & !CMD_ST);
        let mut i = 0;
        while reg(port + P_CMD) & CMD_CR != 0 && i < 1_000_000 {
            i += 1;
        }
        cmd = reg(port + P_CMD);
    }
    if cmd & CMD_FRE != 0 {
        wreg(port + P_CMD, cmd & !CMD_FRE);
        let mut i = 0;
        while reg(port + P_CMD) & CMD_FR != 0 && i < 1_000_000 {
            i += 1;
        }
    }
    unsafe {
        core::ptr::write_bytes(clb as *mut u8, 0, 1024);
        core::ptr::write_bytes(fis as *mut u8, 0, 256);
        core::ptr::write_bytes(ct as *mut u8, 0, 256);
    }

    wreg(port + P_CLB, clb as u32);
    wreg(port + P_CLBU, 0);
    wreg(port + P_FB, fis as u32);
    wreg(port + P_FBU, 0);

    // Allow device detection (clear PxSCTL.DET/IPM), then wait for the
    // link to come up (PxSSTS.DET == 3, communication established).
    // A port with no device settles at DET == 0 quickly — give up early.
    wreg(port + P_SCTL, 0);
    let mut link_ok = false;
    for i in 0..10_000_000u32 {
        let det = reg(port + P_SSTS) & 0xF;
        if det == 3 {
            link_ok = true;
            break;
        }
        if det == 0 && i > 100_000 {
            break; // settled with no device
        }
    }
    if !link_ok {
        crate::serial_println!(
            "ahci: no device on port {} (ssts={:#x})",
            port_idx,
            reg(port + P_SSTS)
        );
        return None;
    }

    // Start the command engine (DMAE required for DMA commands).
    wreg(port + P_CMD, CMD_FRE | CMD_SUD | CMD_DMAE);
    wreg(port + P_CMD, CMD_ST | CMD_FRE | CMD_SUD | CMD_DMAE);

    // ATAPI (optical) devices reject ATA IDENTIFY — skip them so the
    // scan doesn't waste a timeout on a CD-ROM.
    if reg(port + P_CMD) & (1 << 24) != 0 {
        crate::serial_println!("ahci: port {}: ATAPI (optical) device, skipping", port_idx);
        return None;
    }

    let mut d = Disk {
        port,
        port_idx,
        clb,
        ct,
        buf,
        sectors: 0,
    };
    // IDENTIFY DEVICE as a sanity check + to learn the capacity.
    if issue(&d, ATA_IDENTIFY, 0, 1, false).is_err() {
        crate::serial_println!(
            "ahci: IDENTIFY failed on port {} (not an ATA disk?)",
            port_idx
        );
        return None;
    }
    let id = unsafe { core::slice::from_raw_parts(d.buf as *const u16, 256) };
    d.sectors = (id[100] as u64)
        | ((id[101] as u64) << 16)
        | ((id[102] as u64) << 32)
        | ((id[103] as u64) << 48);
    if d.sectors == 0 {
        d.sectors = id[60] as u64; // LBA28 fallback
    }
    let mut model = [0u8; 40];
    for i in 0..20 {
        model[i * 2] = (id[27 + i] >> 8) as u8;
        model[i * 2 + 1] = id[27 + i] as u8;
    }
    let model = core::str::from_utf8(&model).unwrap_or("?").trim();
    crate::serial_println!(
        "ahci: port {} up — {} sectors (~{} MB) \"{}\"",
        port_idx,
        d.sectors,
        d.sectors / 2048,
        model
    );
    Some(d)
}

/// Probe, map and start every implemented port that has a device.  Returns
/// true if at least one disk is ready for sector I/O.  Called once at boot
/// before the filesystem init; the filesystem then calls `select()` to pick
/// which disk to read.
pub fn init() -> bool {
    let Some((dev, func)) = find_ahci() else {
        return false;
    };
    // Enable memory space + bus mastering so MMIO and DMA work.
    let cmd_reg = pci::config_read(0, dev, func, 0x04);
    pci::config_write(0, dev, func, 0x04, cmd_reg | 0x7);

    let Some(abar) = pci::bar64(0, dev, func, 5) else {
        return false;
    };
    crate::serial_println!(
        "ahci: controller at {:02x}:{:02x}.{} abar={:#x}",
        0,
        dev,
        func,
        abar
    );

    // AHCI enable (GHC.AE) before touching port registers.
    wreg(abar as usize + GHC, reg(abar as usize + GHC) | (1 << 31));

    // Scan implemented ports (CAP.NP + 1 of them, filtered by PI).
    let cap = reg(abar as usize + CAP);
    let np = ((cap & 0x1F) + 1).min(MAX_PORTS as u32);
    let pi = reg(abar as usize + PI);
    unsafe {
        ABAR = abar as usize;
        NP = np;
        IMPL_PORTS = pi;
    }
    for port_idx in 0..np {
        if (pi >> port_idx) & 1 == 0 {
            continue;
        }
        let det = reg(abar as usize + PORT_BASE + port_idx as usize * PORT_STRIDE + P_SSTS) & 0xF;
        unsafe {
            PORT_UP[port_idx as usize] = det == 3;
        }
        if let Some(d) = start_port(abar as usize, port_idx) {
            // The device answered IDENTIFY — it's definitely present.
            unsafe {
                PORT_UP[port_idx as usize] = true;
            }
            if !place_disk(d) {
                crate::serial_println!("ahci: port {} up but disk table full", port_idx);
            }
        }
    }
    unsafe {
        ACTIVE = 0;
    }
    unsafe { DISK_COUNT > 0 }
}

/// Record a started disk in the table, reusing a tombstoned (None) slot so
/// indices stay stable across plug/unplug cycles.  Returns false if full.
fn place_disk(d: Disk) -> bool {
    unsafe {
        for i in 0..MAX_DISKS {
            if DISKS[i].is_none() {
                DISKS[i] = Some(d);
                DISK_COUNT = DISK_COUNT.max(i + 1);
                return true;
            }
        }
    }
    false
}

/// Number of disk slots in use (high-water mark; tombstones included).
pub fn disk_count() -> usize {
    unsafe { DISK_COUNT }
}

/// Cheap runtime hotplug probe: true if any implemented port's device
/// presence (PxSSTS.DET == 3) differs from the last known state.  A few
/// MMIO reads — safe to call on every shell prompt.
pub fn poll_dirty() -> bool {
    let abar = unsafe { ABAR };
    if abar == 0 {
        return false;
    }
    let (np, pi) = (unsafe { NP }, unsafe { IMPL_PORTS });
    for port_idx in 0..np {
        if (pi >> port_idx) & 1 == 0 {
            continue;
        }
        let det = reg(abar + PORT_BASE + port_idx as usize * PORT_STRIDE + P_SSTS) & 0xF;
        if (det == 3) != unsafe { PORT_UP[port_idx as usize] } {
            return true;
        }
    }
    false
}

/// Re-scan the controller's ports for arrivals/removals.  Starts ports
/// that now have a link, tombstones (None) the slots of ports whose device
/// vanished, and refreshes the poll state.  Slot indices never shift, so
/// the filesystem's disk handles stay valid.  Returns (added, removed).
pub fn rescan() -> (usize, usize) {
    let abar = unsafe { ABAR };
    if abar == 0 {
        return (0, 0);
    }
    let (np, pi) = (unsafe { NP }, unsafe { IMPL_PORTS });
    let mut added = 0usize;
    let mut removed = 0usize;
    for port_idx in 0..np {
        if (pi >> port_idx) & 1 == 0 {
            continue;
        }
        let port = abar + PORT_BASE + port_idx as usize * PORT_STRIDE;
        let det = reg(port + P_SSTS) & 0xF;
        let up = det == 3;
        let slot = disks()
            .iter()
            .position(|d| d.is_some_and(|d| d.port_idx == port_idx));
        let up_prev = unsafe { PORT_UP[port_idx as usize] };
        unsafe {
            PORT_UP[port_idx as usize] = up;
        }
        if up == up_prev {
            continue;
        }
        if up {
            // Arrival: full bring-up on this port.
            if let Some(d) = start_port(abar, port_idx) {
                if place_disk(d) {
                    added += 1;
                } else {
                    crate::serial_println!("ahci: port {} up but disk table full", port_idx);
                }
            }
        } else if let Some(i) = slot {
            // Removal: tombstone the slot; indices stay stable.
            unsafe {
                DISKS[i] = None;
            }
            crate::serial_println!("ahci: port {} down — device removed", port_idx);
            removed += 1;
        }
    }
    (added, removed)
}

/// Is the disk at index `idx` still attached?  (False for tombstones.)
pub fn is_up(idx: usize) -> bool {
    disks().get(idx).is_some_and(|d| d.is_some())
}

/// Make `idx` the disk the sector API reads/writes.  Returns false if the
/// index is out of range or its slot is a tombstone.
pub fn select(idx: usize) -> bool {
    unsafe {
        if idx < DISK_COUNT && DISKS[idx].is_some() {
            ACTIVE = idx;
            true
        } else {
            false
        }
    }
}

/// Issue one command through slot 0 and poll for completion.
/// `count` is sectors (1-8); data flows through the disk's DMA buffer.
fn issue(d: &Disk, cmd: u8, lba: u64, count: u32, write: bool) -> Result<(), ()> {
    if count == 0 || count > 8 {
        return Err(());
    }
    // Wait for slot 0 to be free.
    let mut i = 0;
    while reg(d.port + P_CI) & 1 != 0 && i < 10_000_000 {
        i += 1;
    }
    if reg(d.port + P_CI) & 1 != 0 {
        return Err(());
    }

    unsafe {
        core::ptr::write_bytes(d.ct as *mut u8, 0, 256);
    }

    // Host-to-device FIS (20 bytes) at the command table start.
    // Layout (SATA spec, cross-checked against QEMU + Linux libahci):
    //   0 type 0x27, 1 flags, 2 command, 3 features low,
    //   4-6 LBA[23:0], 7 device|LBA[27:24], 8-10 LBA[31:48],
    //   11 features high, 12-13 sector count, 14 ICC, 15 control.
    // Flags byte: only the C bit (0x80).  QEMU rejects H2D FISes with
    // bits 6:4 set (reserved); the transfer direction is signalled by
    // the command-header W bit instead, which issue() sets below.
    let fis = d.ct as *mut u8;
    unsafe {
        *fis = 0x27; // FIS type H2D
        *(fis.add(1)) = 0x80; // C only — QEMU rejects flags bits 6:4
        *(fis.add(2)) = cmd;
        *(fis.add(4)) = lba as u8;
        *(fis.add(5)) = (lba >> 8) as u8;
        *(fis.add(6)) = (lba >> 16) as u8;
        *(fis.add(7)) = 0x40 | (((lba >> 24) & 0x0F) as u8);
        *(fis.add(8)) = (lba >> 24) as u8;
        *(fis.add(9)) = (lba >> 32) as u8;
        *(fis.add(10)) = (lba >> 40) as u8;
        *(fis.add(11)) = 0; // features high
        *(fis.add(12)) = count as u8;
        *(fis.add(13)) = (count >> 8) as u8;
        *(fis.add(14)) = 0; // ICC
        *(fis.add(15)) = 0; // control
    }

    // PRD entry at the command table +0x80.
    let prd = (d.ct + 0x80) as *mut u32;
    unsafe {
        *prd = d.buf as u32;
        *prd.add(1) = 0;
        *prd.add(2) = 0;
        *prd.add(3) = (count * SECTOR as u32) - 1; // DBC, no I bit
    }

    // Command header 0. Layout per Linux libahci / QEMU AHCICmdHdr:
    //   DW0 bits 4:0 = CFL (FIS length in dwords), bit 6 = W,
    //   bits 31:16 = PRDTL; DW1 status; DW2/3 = command table address.
    let hdr = d.clb as *mut u32;
    unsafe {
        *hdr = 5 | (1 << 16) | if write { 1 << 6 } else { 0 };
        *hdr.add(1) = 0;
        *hdr.add(2) = d.ct as u32;
        *hdr.add(3) = 0;
    }

    // Issue and poll.  Bail out as soon as the device reports an error
    // (QEMU leaves ATAPI aborts pending in PxCI rather than completing
    // the slot, so we must not wait for the timeout on those).
    wreg(d.port + P_CI, 1);
    let mut i = 0;
    while reg(d.port + P_CI) & 1 != 0 && i < 100_000_000 {
        if reg(d.port + P_TFD) & 1 != 0 {
            break;
        }
        i += 1;
    }
    let tfd = reg(d.port + P_TFD);
    let tfd_err = tfd & 1 != 0;
    if reg(d.port + P_CI) & 1 != 0 && !tfd_err {
        crate::serial_println!(
            "ahci: command {:#x} lba={} cnt={} timed out (tfd={:#x} ci={:#x} cmd={:#x} prdbc={:#x})",
            cmd,
            lba,
            count,
            reg(d.port + P_TFD),
            reg(d.port + P_CI),
            reg(d.port + P_CMD),
            unsafe { *(d.clb as *const u32).add(1) }
        );
        return Err(());
    }
    if tfd_err {
        // Clear the error bits and any wedged slot so the port stays usable.
        wreg(d.port + 0x30, 0xFFFF_FFFF); // PxSERR: clear
        if reg(d.port + P_CI) & 1 != 0 {
            wreg(d.port + P_CI, 1);
        }
        crate::serial_println!(
            "ahci: command {:#x} error (tfd={:#x})",
            cmd,
            reg(d.port + P_TFD)
        );
        return Err(());
    }
    let is = reg(d.port + P_IS);
    if is != 0 {
        wreg(d.port + P_IS, is); // write-1-to-clear
    }
    Ok(())
}

fn disks() -> &'static mut [Option<Disk>; MAX_DISKS] {
    unsafe { &mut *core::ptr::addr_of_mut!(DISKS) }
}

fn active() -> Option<&'static Disk> {
    disks().get(unsafe { ACTIVE }).and_then(|d| d.as_ref())
}

/// Bounds-checked DMA of `count` sectors starting at `lba` into/out of the
/// selected disk's buffer (single code path for both directions).
fn dma_sectors(lba: u64, count: u32, cmd: u8, write: bool) -> Result<(), ()> {
    let d = active().ok_or(())?;
    if lba + count as u64 > d.sectors {
        return Err(());
    }
    issue(d, cmd, lba, count, write)
}

/// DMA `count` sectors starting at `lba` into the selected disk's buffer.
pub fn read_sectors(lba: u64, count: u32) -> Result<(), ()> {
    dma_sectors(lba, count, ATA_READ_DMA_EXT, false)
}

/// DMA `count` sectors starting at `lba` out of the selected disk's buffer.
pub fn write_sectors(lba: u64, count: u32) -> Result<(), ()> {
    dma_sectors(lba, count, ATA_WRITE_DMA_EXT, true)
}

/// Pointer to the selected disk's 4 KiB DMA buffer.
pub fn buf() -> *mut u8 {
    active()
        .map(|d| d.buf as *mut u8)
        .unwrap_or(core::ptr::null_mut())
}

/// Total sector count of the selected disk.
pub fn sector_count() -> u64 {
    active().map(|d| d.sectors).unwrap_or(0)
}