// ---------------------------------------------------------------------------
// 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 port 0
// and exposes sector-level reads/writes for the FAT32 filesystem.
//
// Commands are issued through command-slot 0 with one 4 KiB DMA buffer;
// 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;

// ABAR (HBA memory space) offsets.
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.
struct Disk {
    port: usize,  // port register base = ABAR + 0x100
    clb: usize,   // command list, 1024 B (32 × 32 B), 256 B aligned
    fis: usize,   // FIS receive area, 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 DISK: Option<Disk> = None;

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
}

/// Probe, map and start the first AHCI port.  Returns true if a disk is
/// ready for sector I/O.  Called once at boot before the filesystem init.
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));

    // Pick the first implemented port.
    let pi = reg(abar as usize + PI);
    let mut port_idx = 0u32;
    while port_idx < 32 && (pi >> port_idx) & 1 == 0 {
        port_idx += 1;
    }
    if port_idx >= 32 {
        crate::serial_println!("ahci: no ports implemented (pi={:#x})", pi);
        return false;
    }
    let port = abar as usize + 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 Some(mem) = crate::memory::alloc_pages(2) else {
        crate::serial_println!("ahci: cannot allocate DMA pages");
        return false;
    };
    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).
    wreg(port + P_SCTL, 0);
    let mut link_ok = false;
    for _ in 0..10_000_000 {
        if reg(port + P_SSTS) & 0xF == 3 {
            link_ok = true;
            break;
        }
    }
    if !link_ok {
        crate::serial_println!(
            "ahci: no device on port {} (ssts={:#x})",
            port_idx,
            reg(port + P_SSTS)
        );
        return false;
    }

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

    let mut d = Disk {
        port,
        clb,
        fis,
        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");
        return false;
    }
    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
    );
    unsafe {
        DISK = Some(d);
    }
    true
}

/// 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.
    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 field layout (SATA spec; verified against QEMU/Linux):
        // 4-6 LBA[23:0], 7 device|LBA[27:24], 8-10 LBA[31:48],
        // 11 features high, 12-13 sector count.
        *(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.
    wreg(d.port + P_CI, 1);
    let mut i = 0;
    while reg(d.port + P_CI) & 1 != 0 && i < 100_000_000 {
        i += 1;
    }
    if reg(d.port + P_CI) & 1 != 0 {
        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(());
    }
    // Error?  Clear any port error bits either way.
    if reg(d.port + P_TFD) & 1 != 0 {
        wreg(d.port + 0x30, 0xFFFF_FFFF); // PxSERR: clear
        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(())
}

/// DMA `count` sectors starting at `lba` into the internal buffer.
pub fn read_sectors(lba: u64, count: u32) -> Result<(), ()> {
    let d = unsafe { DISK.as_ref() }.ok_or(())?;
    if lba + count as u64 > d.sectors {
        return Err(());
    }
    issue(d, ATA_READ_DMA_EXT, lba, count, false)
}

/// DMA `count` sectors starting at `lba` out of the internal buffer.
pub fn write_sectors(lba: u64, count: u32) -> Result<(), ()> {
    let d = unsafe { DISK.as_ref() }.ok_or(())?;
    if lba + count as u64 > d.sectors {
        return Err(());
    }
    issue(d, ATA_WRITE_DMA_EXT, lba, count, true)
}

/// Pointer to the internal 4 KiB DMA buffer.
pub fn buf() -> *mut u8 {
    unsafe { DISK.as_ref() }
        .map(|d| d.buf as *mut u8)
        .unwrap_or(core::ptr::null_mut())
}

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