veridian_kernel/drivers/

ahci.rs

//! AHCI (Advanced Host Controller Interface) / SATA Controller Driver
//!
//! Implements AHCI HBA (Host Bus Adapter) discovery via PCI enumeration,
//! port detection with device signature identification, and block-level
//! read/write operations using the AHCI command infrastructure (FIS,
//! Command List, PRDT). Includes NCQ (Native Command Queuing) stubs.
//!
//! PCI identification: class 0x01 (Mass Storage), subclass 0x06 (SATA),
//! prog-if 0x01 (AHCI 1.0).

#![allow(dead_code)]

#[cfg(feature = "alloc")]
extern crate alloc;

#[cfg(feature = "alloc")]
use alloc::{string::String, vec::Vec};
use core::sync::atomic::{AtomicBool, Ordering};

use crate::error::KernelError;
#[cfg(feature = "alloc")]
use crate::fs::blockdev::BlockDevice;

// ---------------------------------------------------------------------------
// PCI identification constants
// ---------------------------------------------------------------------------

/// PCI class code for mass storage controllers.
const PCI_CLASS_MASS_STORAGE: u8 = 0x01;

/// PCI subclass for SATA controllers.
const PCI_SUBCLASS_SATA: u8 = 0x06;

/// PCI programming interface for AHCI 1.0.
const PCI_PROGIF_AHCI: u8 = 0x01;

// ---------------------------------------------------------------------------
// HBA Generic Host Control registers (offsets from ABAR)
// ---------------------------------------------------------------------------

/// Host Capabilities (read-only).
const HBA_REG_CAP: usize = 0x00;

/// Global Host Control.
const HBA_REG_GHC: usize = 0x04;

/// Interrupt Status (read/write-1-to-clear).
const HBA_REG_IS: usize = 0x08;

/// Ports Implemented (read-only bitmask).
const HBA_REG_PI: usize = 0x0C;

/// AHCI Version.
const HBA_REG_VS: usize = 0x10;

/// Command Completion Coalescing Control.
const HBA_REG_CCC_CTL: usize = 0x14;

/// Command Completion Coalescing Ports.
const HBA_REG_CCC_PORTS: usize = 0x18;

/// Enclosure Management Location.
const HBA_REG_EM_LOC: usize = 0x1C;

/// Enclosure Management Control.
const HBA_REG_EM_CTL: usize = 0x20;

/// Host Capabilities Extended.
const HBA_REG_CAP2: usize = 0x24;

/// BIOS/OS Handoff Control and Status.
const HBA_REG_BOHC: usize = 0x28;

// ---------------------------------------------------------------------------
// GHC (Global Host Control) register bits
// ---------------------------------------------------------------------------

/// HBA Reset.
const GHC_HR: u32 = 1 << 0;

/// Interrupt Enable.
const GHC_IE: u32 = 1 << 1;

/// AHCI Enable.
const GHC_AE: u32 = 1 << 31;

// ---------------------------------------------------------------------------
// CAP (Host Capabilities) register bits
// ---------------------------------------------------------------------------

/// Number of ports (bits 4:0), zero-based.
const CAP_NP_MASK: u32 = 0x1F;

/// Supports NCQ (bit 30).
const CAP_SNCQ: u32 = 1 << 30;

/// Supports 64-bit addressing (bit 31).
const CAP_S64A: u32 = 1 << 31;

/// Number of command slots (bits 12:8), zero-based.
const CAP_NCS_SHIFT: u32 = 8;
const CAP_NCS_MASK: u32 = 0x1F;

// ---------------------------------------------------------------------------
// Port register block (per-port, base = 0x100 + port * 0x80)
// ---------------------------------------------------------------------------

const PORT_BASE: usize = 0x100;
const PORT_SIZE: usize = 0x80;

/// Port Command List Base Address (lower 32 bits).
const PORT_CLB: usize = 0x00;

/// Port Command List Base Address (upper 32 bits).
const PORT_CLBU: usize = 0x04;

/// Port FIS Base Address (lower 32 bits).
const PORT_FB: usize = 0x08;

/// Port FIS Base Address (upper 32 bits).
const PORT_FBU: usize = 0x0C;

/// Port Interrupt Status.
const PORT_IS: usize = 0x10;

/// Port Interrupt Enable.
const PORT_IE: usize = 0x14;

/// Port Command and Status.
const PORT_CMD: usize = 0x18;

/// Port Task File Data.
const PORT_TFD: usize = 0x20;

/// Port Signature.
const PORT_SIG: usize = 0x24;

/// Port SATA Status (SCR0: SStatus).
const PORT_SSTS: usize = 0x28;

/// Port SATA Control (SCR2: SControl).
const PORT_SCTL: usize = 0x2C;

/// Port SATA Error (SCR1: SError).
const PORT_SERR: usize = 0x30;

/// Port SATA Active (SCR3: SActive) -- for NCQ.
const PORT_SACT: usize = 0x34;

/// Port Command Issue.
const PORT_CI: usize = 0x38;

// ---------------------------------------------------------------------------
// PORT_CMD bits
// ---------------------------------------------------------------------------

/// Start (process command list).
const PORT_CMD_ST: u32 = 1 << 0;

/// Spin-Up Device.
const PORT_CMD_SUD: u32 = 1 << 1;

/// Power On Device.
const PORT_CMD_POD: u32 = 1 << 2;

/// FIS Receive Enable.
const PORT_CMD_FRE: u32 = 1 << 4;

/// FIS Receive Running.
const PORT_CMD_FR: u32 = 1 << 14;

/// Command List Running.
const PORT_CMD_CR: u32 = 1 << 15;

// ---------------------------------------------------------------------------
// PORT_TFD bits
// ---------------------------------------------------------------------------

/// Task file status: BSY (busy).
const TFD_BSY: u32 = 1 << 7;

/// Task file status: DRQ (data request).
const TFD_DRQ: u32 = 1 << 3;

/// Task file status: ERR (error).
const TFD_ERR: u32 = 1 << 0;

// ---------------------------------------------------------------------------
// PORT_SSTS (SStatus) detection and interface fields
// ---------------------------------------------------------------------------

/// Device detection (bits 3:0).
const SSTS_DET_MASK: u32 = 0x0F;

/// Device present and PHY communication established.
const SSTS_DET_PRESENT: u32 = 0x03;

/// Interface Power Management (bits 11:8).
const SSTS_IPM_MASK: u32 = 0x0F00;

/// Interface active.
const SSTS_IPM_ACTIVE: u32 = 0x0100;

// ---------------------------------------------------------------------------
// Device signatures (from PORT_SIG after device reset)
// ---------------------------------------------------------------------------

/// SATA drive (ATA).
const SIG_SATA: u32 = 0x00000101;

/// SATAPI (ATAPI) drive.
const SIG_SATAPI: u32 = 0xEB140101;

/// Enclosure Management Bridge (SEMB).
const SIG_SEMB: u32 = 0xC33C0101;

/// Port Multiplier.
const SIG_PM: u32 = 0x96690101;

// ---------------------------------------------------------------------------
// FIS Types
// ---------------------------------------------------------------------------

/// FIS type: Register -- Host to Device.
const FIS_TYPE_REG_H2D: u8 = 0x27;

/// FIS type: Register -- Device to Host.
const FIS_TYPE_REG_D2H: u8 = 0x34;

/// FIS type: DMA Activate -- Device to Host.
const FIS_TYPE_DMA_ACT: u8 = 0x39;

/// FIS type: DMA Setup -- Bidirectional.
const FIS_TYPE_DMA_SETUP: u8 = 0x41;

/// FIS type: Data -- Bidirectional.
const FIS_TYPE_DATA: u8 = 0x46;

/// FIS type: BIST Activate.
const FIS_TYPE_BIST: u8 = 0x58;

/// FIS type: PIO Setup -- Device to Host.
const FIS_TYPE_PIO_SETUP: u8 = 0x5F;

/// FIS type: Set Device Bits -- Device to Host.
const FIS_TYPE_DEV_BITS: u8 = 0xA1;

// ---------------------------------------------------------------------------
// ATA commands
// ---------------------------------------------------------------------------

/// ATA READ DMA EXT (48-bit LBA).
const ATA_CMD_READ_DMA_EXT: u8 = 0x25;

/// ATA WRITE DMA EXT (48-bit LBA).
const ATA_CMD_WRITE_DMA_EXT: u8 = 0x35;

/// ATA IDENTIFY DEVICE.
const ATA_CMD_IDENTIFY: u8 = 0xEC;

/// ATA READ FPDMA QUEUED (NCQ read).
const ATA_CMD_READ_FPDMA_QUEUED: u8 = 0x60;

/// ATA WRITE FPDMA QUEUED (NCQ write).
const ATA_CMD_WRITE_FPDMA_QUEUED: u8 = 0x61;

/// ATA FLUSH CACHE EXT.
const ATA_CMD_FLUSH_CACHE_EXT: u8 = 0xEA;

// ---------------------------------------------------------------------------
// Sector size
// ---------------------------------------------------------------------------

/// Standard sector size in bytes.
const SECTOR_SIZE: usize = 512;

/// Maximum PRDT entries per command table.
const MAX_PRDT_ENTRIES: usize = 65535;

/// Maximum sectors per single DMA transfer (limited by PRDT entry: 4MB / 512).
const MAX_SECTORS_PER_PRDT: u32 = 8192;

/// Polling timeout iterations.
const POLL_TIMEOUT: u32 = 1_000_000;

// ---------------------------------------------------------------------------
// FIS structures
// ---------------------------------------------------------------------------

/// FIS Register -- Host to Device (20 bytes).
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct FisRegH2D {
    /// FIS type (0x27).
    pub fis_type: u8,
    /// Port multiplier | Command/Control bit (bit 7).
    pub pm_and_c: u8,
    /// ATA command register.
    pub command: u8,
    /// Feature register (7:0).
    pub feature_lo: u8,

    /// LBA (7:0).
    pub lba0: u8,
    /// LBA (15:8).
    pub lba1: u8,
    /// LBA (23:16).
    pub lba2: u8,
    /// Device register.
    pub device: u8,

    /// LBA (31:24).
    pub lba3: u8,
    /// LBA (39:32).
    pub lba4: u8,
    /// LBA (47:40).
    pub lba5: u8,
    /// Feature register (15:8).
    pub feature_hi: u8,

    /// Sector count (7:0).
    pub count_lo: u8,
    /// Sector count (15:8).
    pub count_hi: u8,
    /// Isochronous command completion.
    pub icc: u8,
    /// Control register.
    pub control: u8,

    /// Reserved.
    pub _reserved: [u8; 4],
}

impl Default for FisRegH2D {
    fn default() -> Self {
        Self::new()
    }
}

impl FisRegH2D {
    /// Create a new H2D Register FIS with the command bit set.
    pub fn new() -> Self {
        Self {
            fis_type: FIS_TYPE_REG_H2D,
            pm_and_c: 0x80, // Command bit set
            command: 0,
            feature_lo: 0,
            lba0: 0,
            lba1: 0,
            lba2: 0,
            device: 0,
            lba3: 0,
            lba4: 0,
            lba5: 0,
            feature_hi: 0,
            count_lo: 0,
            count_hi: 0,
            icc: 0,
            control: 0,
            _reserved: [0; 4],
        }
    }

    /// Set 48-bit LBA address.
    pub fn set_lba(&mut self, lba: u64) {
        self.lba0 = (lba & 0xFF) as u8;
        self.lba1 = ((lba >> 8) & 0xFF) as u8;
        self.lba2 = ((lba >> 16) & 0xFF) as u8;
        self.lba3 = ((lba >> 24) & 0xFF) as u8;
        self.lba4 = ((lba >> 32) & 0xFF) as u8;
        self.lba5 = ((lba >> 40) & 0xFF) as u8;
        self.device = 1 << 6; // LBA mode
    }

    /// Set sector count (16-bit).
    pub fn set_count(&mut self, count: u16) {
        self.count_lo = (count & 0xFF) as u8;
        self.count_hi = ((count >> 8) & 0xFF) as u8;
    }
}

/// FIS Register -- Device to Host (20 bytes).
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct FisRegD2H {
    /// FIS type (0x34).
    pub fis_type: u8,
    /// Port multiplier | Interrupt bit (bit 6).
    pub pm_and_i: u8,
    /// Status register.
    pub status: u8,
    /// Error register.
    pub error: u8,

    /// LBA (7:0).
    pub lba0: u8,
    /// LBA (15:8).
    pub lba1: u8,
    /// LBA (23:16).
    pub lba2: u8,
    /// Device register.
    pub device: u8,

    /// LBA (31:24).
    pub lba3: u8,
    /// LBA (39:32).
    pub lba4: u8,
    /// LBA (47:40).
    pub lba5: u8,
    /// Reserved.
    pub _reserved0: u8,

    /// Sector count (7:0).
    pub count_lo: u8,
    /// Sector count (15:8).
    pub count_hi: u8,
    /// Reserved.
    pub _reserved1: [u8; 6],
}

/// FIS DMA Setup -- Bidirectional (28 bytes).
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct FisDmaSetup {
    /// FIS type (0x41).
    pub fis_type: u8,
    /// Port multiplier | Direction | Interrupt | Auto-activate.
    pub flags: u8,
    /// Reserved.
    pub _reserved0: [u8; 2],

    /// DMA Buffer Identifier (low).
    pub dma_buf_id_lo: u32,
    /// DMA Buffer Identifier (high).
    pub dma_buf_id_hi: u32,

    /// Reserved.
    pub _reserved1: u32,
    /// DMA buffer offset.
    pub dma_buf_offset: u32,
    /// Transfer count.
    pub transfer_count: u32,
    /// Reserved.
    pub _reserved2: u32,
}

/// FIS PIO Setup -- Device to Host (20 bytes).
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct FisPioSetup {
    /// FIS type (0x5F).
    pub fis_type: u8,
    /// Port multiplier | Direction | Interrupt.
    pub flags: u8,
    /// Status register.
    pub status: u8,
    /// Error register.
    pub error: u8,

    /// LBA (7:0).
    pub lba0: u8,
    /// LBA (15:8).
    pub lba1: u8,
    /// LBA (23:16).
    pub lba2: u8,
    /// Device register.
    pub device: u8,

    /// LBA (31:24).
    pub lba3: u8,
    /// LBA (39:32).
    pub lba4: u8,
    /// LBA (47:40).
    pub lba5: u8,
    /// Reserved.
    pub _reserved0: u8,

    /// Sector count (7:0).
    pub count_lo: u8,
    /// Sector count (15:8).
    pub count_hi: u8,
    /// Reserved.
    pub _reserved1: u8,
    /// New value of status register (E_Status).
    pub e_status: u8,

    /// Transfer count.
    pub transfer_count: u16,
    /// Reserved.
    pub _reserved2: [u8; 2],
}

/// FIS Data -- Bidirectional (variable length, header is 4 bytes).
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct FisData {
    /// FIS type (0x46).
    pub fis_type: u8,
    /// Port multiplier.
    pub pm: u8,
    /// Reserved.
    pub _reserved: [u8; 2],
    // Followed by DWORD-aligned payload data.
}

// ---------------------------------------------------------------------------
// Received FIS structure (256 bytes, per port)
// ---------------------------------------------------------------------------

/// Received FIS area for a single port (256 bytes, 256-byte aligned).
#[repr(C, align(256))]
#[derive(Debug, Clone, Copy)]
pub struct ReceivedFis {
    /// DMA Setup FIS (offset 0x00).
    pub dma_setup: FisDmaSetup,
    /// Padding.
    pub _pad0: [u8; 4],

    /// PIO Setup FIS (offset 0x20).
    pub pio_setup: FisPioSetup,
    /// Padding.
    pub _pad1: [u8; 12],

    /// D2H Register FIS (offset 0x40).
    pub d2h_reg: FisRegD2H,
    /// Padding.
    pub _pad2: [u8; 4],

    /// Set Device Bits FIS (offset 0x58).
    pub set_device_bits: [u8; 8],

    /// Unknown FIS (offset 0x60).
    pub unknown: [u8; 64],

    /// Reserved (offset 0xA0).
    pub _reserved: [u8; 96],
}

// ---------------------------------------------------------------------------
// Command structures
// ---------------------------------------------------------------------------

/// Physical Region Descriptor Table entry (16 bytes).
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct PrdtEntry {
    /// Data base address (lower 32 bits, must be word-aligned).
    pub data_base_lo: u32,
    /// Data base address (upper 32 bits).
    pub data_base_hi: u32,
    /// Reserved.
    pub _reserved: u32,
    /// Byte count (bit 0 must be 1 for odd byte count; bits 21:0 = count - 1).
    /// Bit 31: Interrupt on Completion.
    pub dbc_and_flags: u32,
}

impl PrdtEntry {
    /// Create a new PRDT entry.
    ///
    /// `phys_addr`: Physical address of the data buffer (must be word-aligned).
    /// `byte_count`: Number of bytes to transfer (max 4MB per entry, must be
    /// even). `interrupt`: If true, set the Interrupt on Completion bit.
    pub fn new(phys_addr: u64, byte_count: u32, interrupt: bool) -> Self {
        let mut dbc = (byte_count - 1) & 0x003F_FFFF; // bits 21:0
        if interrupt {
            dbc |= 1 << 31;
        }
        Self {
            data_base_lo: phys_addr as u32,
            data_base_hi: (phys_addr >> 32) as u32,
            _reserved: 0,
            dbc_and_flags: dbc,
        }
    }
}

/// Command Header (32 bytes, one per command slot).
///
/// The command list is an array of up to 32 command headers.
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct CommandHeader {
    /// DW0: Command FIS Length (bits 4:0, in DWORDs), ATAPI (bit 5),
    /// Write direction (bit 6), Prefetchable (bit 7), Reset (bit 8),
    /// BIST (bit 9), Clear Busy upon R_OK (bit 10), Reserved (bit 11),
    /// Port Multiplier Port (bits 15:12), PRDT Length (bits 31:16).
    pub flags_and_prdtl: u32,
    /// DW1: Physical Region Descriptor Byte Count (transferred so far).
    pub prdbc: u32,
    /// DW2: Command Table Descriptor Base Address (lower 32 bits, 128-byte
    /// aligned).
    pub ctba_lo: u32,
    /// DW3: Command Table Descriptor Base Address (upper 32 bits).
    pub ctba_hi: u32,
    /// DW4-7: Reserved.
    pub _reserved: [u32; 4],
}

impl CommandHeader {
    /// Create a new command header.
    ///
    /// `fis_len_dwords`: Length of the command FIS in DWORDs (typically 5 for
    /// H2D FIS). `prdt_count`: Number of PRDT entries.
    /// `write`: True if this is a write (host-to-device data) command.
    /// `ctba_phys`: Physical address of the command table (128-byte aligned).
    pub fn new(fis_len_dwords: u8, prdt_count: u16, write: bool, ctba_phys: u64) -> Self {
        let mut flags = (fis_len_dwords as u32) & 0x1F;
        if write {
            flags |= 1 << 6; // W bit
        }
        flags |= (prdt_count as u32) << 16;

        Self {
            flags_and_prdtl: flags,
            prdbc: 0,
            ctba_lo: ctba_phys as u32,
            ctba_hi: (ctba_phys >> 32) as u32,
            _reserved: [0; 4],
        }
    }
}

/// Command Table (variable size: 128-byte header + PRDT entries).
///
/// The header contains the Command FIS (64 bytes), ATAPI Command (16 bytes),
/// and reserved space (48 bytes). Followed by PRDT entries.
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct CommandTableHeader {
    /// Command FIS (up to 64 bytes).
    pub cfis: [u8; 64],
    /// ATAPI Command (12 or 16 bytes, zero-padded to 16).
    pub acmd: [u8; 16],
    /// Reserved.
    pub _reserved: [u8; 48],
    // Followed by PRDT entries (PrdtEntry array).
}

impl Default for CommandTableHeader {
    fn default() -> Self {
        Self::new()
    }
}

impl CommandTableHeader {
    /// Create a zeroed command table header.
    pub fn new() -> Self {
        Self {
            cfis: [0; 64],
            acmd: [0; 16],
            _reserved: [0; 48],
        }
    }

    /// Write a Register H2D FIS into the command FIS area.
    pub fn set_h2d_fis(&mut self, fis: &FisRegH2D) {
        let fis_bytes =
            unsafe { core::slice::from_raw_parts(fis as *const FisRegH2D as *const u8, 20) };
        self.cfis[..20].copy_from_slice(fis_bytes);
    }
}

// ---------------------------------------------------------------------------
// Device type identification
// ---------------------------------------------------------------------------

/// Type of device attached to an AHCI port.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AhciDeviceType {
    /// No device detected.
    None,
    /// SATA drive (hard disk or SSD).
    Sata,
    /// SATAPI device (optical drive, etc.).
    Satapi,
    /// Enclosure Management Bridge.
    Semb,
    /// Port Multiplier.
    PortMultiplier,
    /// Unknown signature.
    Unknown(u32),
}

impl AhciDeviceType {
    /// Determine device type from port signature register value.
    pub fn from_signature(sig: u32) -> Self {
        match sig {
            SIG_SATA => Self::Sata,
            SIG_SATAPI => Self::Satapi,
            SIG_SEMB => Self::Semb,
            SIG_PM => Self::PortMultiplier,
            0xFFFFFFFF | 0x00000000 => Self::None,
            other => Self::Unknown(other),
        }
    }
}

// ---------------------------------------------------------------------------
// Port state
// ---------------------------------------------------------------------------

/// State of an AHCI port.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PortState {
    /// Port is not implemented or no device.
    Inactive,
    /// Device detected, not yet initialized.
    Detected,
    /// Port fully initialized and ready for commands.
    Ready,
    /// Port encountered an error.
    Error,
}

// ---------------------------------------------------------------------------
// AHCI Port
// ---------------------------------------------------------------------------

/// Represents a single AHCI port with its state and device info.
#[derive(Debug, Clone, Copy)]
pub struct AhciPort {
    /// Port number (0-31).
    pub port_num: u8,
    /// Base address for this port's MMIO registers.
    pub port_mmio_base: usize,
    /// Device type detected on this port.
    pub device_type: AhciDeviceType,
    /// Current port state.
    pub state: PortState,
    /// Number of command slots supported by the HBA.
    pub num_cmd_slots: u8,
    /// Whether NCQ is supported by the HBA.
    pub ncq_supported: bool,
    /// Total sectors reported by IDENTIFY DEVICE (0 until identified).
    pub total_sectors: u64,
    /// Sector size in bytes (default 512).
    pub sector_size: usize,
}

impl AhciPort {
    /// Create a new port descriptor.
    pub fn new(port_num: u8, hba_mmio_base: usize, num_cmd_slots: u8, ncq_supported: bool) -> Self {
        Self {
            port_num,
            port_mmio_base: hba_mmio_base + PORT_BASE + (port_num as usize) * PORT_SIZE,
            device_type: AhciDeviceType::None,
            state: PortState::Inactive,
            num_cmd_slots,
            ncq_supported,
            total_sectors: 0,
            sector_size: SECTOR_SIZE,
        }
    }

    /// Read a 32-bit port register.
    fn read_port_reg(&self, offset: usize) -> u32 {
        // SAFETY: Reading an AHCI port MMIO register. The port_mmio_base was
        // computed from a validated PCI BAR5 address. read_volatile ensures
        // the compiler does not elide or reorder this hardware register access.
        unsafe { core::ptr::read_volatile((self.port_mmio_base + offset) as *const u32) }
    }

    /// Write a 32-bit port register.
    fn write_port_reg(&self, offset: usize, value: u32) {
        // SAFETY: Writing an AHCI port MMIO register. Same invariants as
        // read_port_reg.
        unsafe { core::ptr::write_volatile((self.port_mmio_base + offset) as *mut u32, value) }
    }

    /// Detect whether a device is present on this port.
    pub fn detect_device(&mut self) {
        let ssts = self.read_port_reg(PORT_SSTS);
        let det = ssts & SSTS_DET_MASK;
        let ipm = ssts & SSTS_IPM_MASK;

        if det != SSTS_DET_PRESENT || ipm != SSTS_IPM_ACTIVE {
            self.device_type = AhciDeviceType::None;
            self.state = PortState::Inactive;
            return;
        }

        let sig = self.read_port_reg(PORT_SIG);
        self.device_type = AhciDeviceType::from_signature(sig);
        self.state = PortState::Detected;
    }

    /// Stop the port command engine (clear ST and FRE, wait for CR and FR to
    /// clear).
    pub fn stop_cmd(&self) -> Result<(), KernelError> {
        let mut cmd = self.read_port_reg(PORT_CMD);

        // Clear ST (Start)
        cmd &= !PORT_CMD_ST;
        self.write_port_reg(PORT_CMD, cmd);

        // Wait for CR (Command List Running) to clear
        let mut timeout = POLL_TIMEOUT;
        while self.read_port_reg(PORT_CMD) & PORT_CMD_CR != 0 && timeout > 0 {
            timeout -= 1;
            core::hint::spin_loop();
        }
        if timeout == 0 {
            return Err(KernelError::Timeout {
                operation: "ahci_stop_cmd_cr",
                duration_ms: 0,
            });
        }

        // Clear FRE (FIS Receive Enable)
        cmd = self.read_port_reg(PORT_CMD);
        cmd &= !PORT_CMD_FRE;
        self.write_port_reg(PORT_CMD, cmd);

        // Wait for FR (FIS Receive Running) to clear
        timeout = POLL_TIMEOUT;
        while self.read_port_reg(PORT_CMD) & PORT_CMD_FR != 0 && timeout > 0 {
            timeout -= 1;
            core::hint::spin_loop();
        }
        if timeout == 0 {
            return Err(KernelError::Timeout {
                operation: "ahci_stop_cmd_fr",
                duration_ms: 0,
            });
        }

        Ok(())
    }

    /// Start the port command engine (set FRE then ST).
    pub fn start_cmd(&self) {
        // Wait until CR is clear before starting
        while self.read_port_reg(PORT_CMD) & PORT_CMD_CR != 0 {
            core::hint::spin_loop();
        }

        let mut cmd = self.read_port_reg(PORT_CMD);
        cmd |= PORT_CMD_FRE;
        self.write_port_reg(PORT_CMD, cmd);

        cmd = self.read_port_reg(PORT_CMD);
        cmd |= PORT_CMD_ST;
        self.write_port_reg(PORT_CMD, cmd);
    }

    /// Clear the port SERR register (write-1-to-clear all bits).
    pub fn clear_serr(&self) {
        self.write_port_reg(PORT_SERR, 0xFFFF_FFFF);
    }

    /// Clear the port interrupt status register.
    pub fn clear_interrupt_status(&self) {
        self.write_port_reg(PORT_IS, 0xFFFF_FFFF);
    }

    /// Wait for the port to become not busy (TFD BSY and DRQ clear).
    pub fn wait_ready(&self) -> Result<(), KernelError> {
        let mut timeout = POLL_TIMEOUT;
        while timeout > 0 {
            let tfd = self.read_port_reg(PORT_TFD);
            if tfd & (TFD_BSY | TFD_DRQ) == 0 {
                return Ok(());
            }
            timeout -= 1;
            core::hint::spin_loop();
        }
        Err(KernelError::Timeout {
            operation: "ahci_port_ready",
            duration_ms: 0,
        })
    }

    /// Find a free command slot.
    ///
    /// Returns the slot index (0-based), or an error if all slots are occupied.
    pub fn find_free_slot(&self) -> Result<u8, KernelError> {
        let ci = self.read_port_reg(PORT_CI);
        let sact = self.read_port_reg(PORT_SACT);
        let occupied = ci | sact;

        for slot in 0..self.num_cmd_slots {
            if occupied & (1 << slot) == 0 {
                return Ok(slot);
            }
        }

        Err(KernelError::ResourceExhausted {
            resource: "ahci_command_slots",
        })
    }

    /// Issue a command in the given slot and wait for completion.
    ///
    /// Sets the corresponding bit in PORT_CI and polls until it clears or
    /// an error is detected in PORT_TFD.
    pub fn issue_command_and_wait(&self, slot: u8) -> Result<(), KernelError> {
        // Clear any pending errors
        self.clear_serr();
        self.clear_interrupt_status();

        // Issue command
        self.write_port_reg(PORT_CI, 1 << slot);

        // Poll for completion
        let mut timeout = POLL_TIMEOUT;
        loop {
            let ci = self.read_port_reg(PORT_CI);
            if ci & (1 << slot) == 0 {
                // Command completed
                break;
            }

            let tfd = self.read_port_reg(PORT_TFD);
            if tfd & TFD_ERR != 0 {
                return Err(KernelError::HardwareError {
                    device: "ahci",
                    code: tfd,
                });
            }

            timeout -= 1;
            if timeout == 0 {
                return Err(KernelError::Timeout {
                    operation: "ahci_command_completion",
                    duration_ms: 0,
                });
            }
            core::hint::spin_loop();
        }

        // Check final status
        let tfd = self.read_port_reg(PORT_TFD);
        if tfd & TFD_ERR != 0 {
            return Err(KernelError::HardwareError {
                device: "ahci",
                code: tfd,
            });
        }

        Ok(())
    }

    /// Build and issue a READ DMA EXT command.
    ///
    /// `lba`: Starting logical block address.
    /// `sector_count`: Number of sectors to read (1-65535, 0 means 65536).
    /// `buffer_phys`: Physical address of the destination buffer.
    ///
    /// This prepares the FIS, command header, PRDT, and issues the command.
    /// In a production system, the command list and command tables would be
    /// allocated from DMA-capable memory; here we describe the operation
    /// structurally.
    pub fn build_read_dma_ext(
        &self,
        lba: u64,
        sector_count: u16,
        buffer_phys: u64,
    ) -> (CommandHeader, CommandTableHeader, PrdtEntry) {
        let mut fis = FisRegH2D::new();
        fis.command = ATA_CMD_READ_DMA_EXT;
        fis.set_lba(lba);
        fis.set_count(sector_count);

        let mut ct = CommandTableHeader::new();
        ct.set_h2d_fis(&fis);

        let byte_count = (sector_count as u32) * (self.sector_size as u32);
        let prdt = PrdtEntry::new(buffer_phys, byte_count, true);

        // Command header: FIS length = 5 DWORDs (20 bytes / 4), 1 PRDT entry,
        // not a write, command table physical address = 0 (placeholder).
        let cmd_hdr = CommandHeader::new(5, 1, false, 0);

        (cmd_hdr, ct, prdt)
    }

    /// Build and issue a WRITE DMA EXT command.
    ///
    /// `lba`: Starting logical block address.
    /// `sector_count`: Number of sectors to write.
    /// `buffer_phys`: Physical address of the source buffer.
    pub fn build_write_dma_ext(
        &self,
        lba: u64,
        sector_count: u16,
        buffer_phys: u64,
    ) -> (CommandHeader, CommandTableHeader, PrdtEntry) {
        let mut fis = FisRegH2D::new();
        fis.command = ATA_CMD_WRITE_DMA_EXT;
        fis.set_lba(lba);
        fis.set_count(sector_count);

        let mut ct = CommandTableHeader::new();
        ct.set_h2d_fis(&fis);

        let byte_count = (sector_count as u32) * (self.sector_size as u32);
        let prdt = PrdtEntry::new(buffer_phys, byte_count, true);

        // Write bit set in command header
        let cmd_hdr = CommandHeader::new(5, 1, true, 0);

        (cmd_hdr, ct, prdt)
    }
}

// ---------------------------------------------------------------------------
// NCQ (Native Command Queuing) stubs
// ---------------------------------------------------------------------------

/// NCQ command tag (0-31).
pub type NcqTag = u8;

/// Build a READ FPDMA QUEUED (NCQ) FIS.
///
/// NCQ uses the SATA Active (SActive) register rather than Command Issue.
/// The tag is encoded in the sector count register (bits 7:3).
pub fn build_ncq_read_fis(lba: u64, sector_count: u16, tag: NcqTag) -> FisRegH2D {
    let mut fis = FisRegH2D::new();
    fis.command = ATA_CMD_READ_FPDMA_QUEUED;
    fis.set_lba(lba);
    // NCQ: sector count goes in feature register
    fis.feature_lo = (sector_count & 0xFF) as u8;
    fis.feature_hi = ((sector_count >> 8) & 0xFF) as u8;
    // Tag in count register bits 7:3
    fis.count_lo = (tag & 0x1F) << 3;
    fis.count_hi = 0;
    fis.device = 1 << 6; // LBA mode
    fis
}

/// Build a WRITE FPDMA QUEUED (NCQ) FIS.
pub fn build_ncq_write_fis(lba: u64, sector_count: u16, tag: NcqTag) -> FisRegH2D {
    let mut fis = FisRegH2D::new();
    fis.command = ATA_CMD_WRITE_FPDMA_QUEUED;
    fis.set_lba(lba);
    fis.feature_lo = (sector_count & 0xFF) as u8;
    fis.feature_hi = ((sector_count >> 8) & 0xFF) as u8;
    fis.count_lo = (tag & 0x1F) << 3;
    fis.count_hi = 0;
    fis.device = 1 << 6; // LBA mode
    fis
}

/// Issue an NCQ command on a port (stub).
///
/// In a full implementation, this would:
/// 1. Find a free NCQ tag (0-31).
/// 2. Build the FPDMA FIS and set up the command table.
/// 3. Set the corresponding bit in PORT_SACT (SActive).
/// 4. Set the corresponding bit in PORT_CI (Command Issue).
/// 5. Return the tag for later completion tracking.
pub fn issue_ncq_command(_port: &AhciPort, _fis: &FisRegH2D) -> Result<NcqTag, KernelError> {
    Err(KernelError::NotImplemented {
        feature: "ahci_ncq_command_issue",
    })
}

/// Poll for NCQ command completion on a port (stub).
///
/// Checks PORT_SACT to determine which tags have completed.
pub fn poll_ncq_completion(_port: &AhciPort) -> u32 {
    // In a real implementation: return ~PORT_SACT & issued_tags
    0
}

// ---------------------------------------------------------------------------
// AHCI Controller
// ---------------------------------------------------------------------------

/// AHCI Host Bus Adapter controller.
#[cfg(feature = "alloc")]
pub struct AhciController {
    /// MMIO base address (BAR5 / ABAR mapped to virtual address).
    mmio_base: usize,
    /// HBA capabilities register value.
    capabilities: u32,
    /// Number of ports supported by this HBA.
    num_ports: u8,
    /// Number of command slots per port.
    num_cmd_slots: u8,
    /// Whether NCQ is supported.
    ncq_supported: bool,
    /// Whether 64-bit addressing is supported.
    supports_64bit: bool,
    /// AHCI version (major.minor packed as u32).
    version: u32,
    /// Detected ports.
    ports: Vec<AhciPort>,
}

#[cfg(feature = "alloc")]
impl AhciController {
    /// Create a new AHCI controller from a mapped MMIO base address.
    pub fn new(mmio_base: usize) -> Result<Self, KernelError> {
        let mut ctrl = Self {
            mmio_base,
            capabilities: 0,
            num_ports: 0,
            num_cmd_slots: 0,
            ncq_supported: false,
            supports_64bit: false,
            version: 0,
            ports: Vec::new(),
        };

        ctrl.initialize()?;
        Ok(ctrl)
    }

    /// Read a 32-bit HBA register.
    fn read_hba_reg(&self, offset: usize) -> u32 {
        // SAFETY: Reading an AHCI HBA MMIO register at mmio_base + offset.
        // The mmio_base was derived from a PCI BAR5 address. read_volatile
        // ensures the compiler does not elide or reorder this hardware access.
        unsafe { core::ptr::read_volatile((self.mmio_base + offset) as *const u32) }
    }

    /// Write a 32-bit HBA register.
    fn write_hba_reg(&self, offset: usize, value: u32) {
        // SAFETY: Writing an AHCI HBA MMIO register. Same invariants as
        // read_hba_reg.
        unsafe { core::ptr::write_volatile((self.mmio_base + offset) as *mut u32, value) }
    }

    /// Initialize the HBA: read capabilities, enable AHCI mode, detect ports.
    fn initialize(&mut self) -> Result<(), KernelError> {
        crate::println!("[AHCI] Initializing controller at 0x{:x}", self.mmio_base);

        // Read version
        self.version = self.read_hba_reg(HBA_REG_VS);
        let major = (self.version >> 16) & 0xFFFF;
        let minor = self.version & 0xFFFF;
        crate::println!("[AHCI] Version: {}.{}", major, minor);

        // Read capabilities
        self.capabilities = self.read_hba_reg(HBA_REG_CAP);
        self.num_ports = ((self.capabilities & CAP_NP_MASK) + 1) as u8;
        self.num_cmd_slots = (((self.capabilities >> CAP_NCS_SHIFT) & CAP_NCS_MASK) + 1) as u8;
        self.ncq_supported = self.capabilities & CAP_SNCQ != 0;
        self.supports_64bit = self.capabilities & CAP_S64A != 0;

        crate::println!(
            "[AHCI] Ports: {}, Command Slots: {}, NCQ: {}, 64-bit: {}",
            self.num_ports,
            self.num_cmd_slots,
            self.ncq_supported,
            self.supports_64bit,
        );

        // Enable AHCI mode (set GHC.AE)
        let ghc = self.read_hba_reg(HBA_REG_GHC);
        if ghc & GHC_AE == 0 {
            self.write_hba_reg(HBA_REG_GHC, ghc | GHC_AE);
            crate::println!("[AHCI] Enabled AHCI mode");
        }

        // Perform HBA reset
        self.reset()?;

        // Re-enable AHCI mode after reset
        let ghc = self.read_hba_reg(HBA_REG_GHC);
        self.write_hba_reg(HBA_REG_GHC, ghc | GHC_AE);

        // Enable global interrupts
        let ghc = self.read_hba_reg(HBA_REG_GHC);
        self.write_hba_reg(HBA_REG_GHC, ghc | GHC_IE);

        // Clear global interrupt status
        self.write_hba_reg(HBA_REG_IS, 0xFFFF_FFFF);

        // Detect ports
        self.detect_ports();

        crate::println!(
            "[AHCI] Controller initialized ({} active port(s))",
            self.ports
                .iter()
                .filter(|p| p.state != PortState::Inactive)
                .count()
        );

        Ok(())
    }

    /// Perform an HBA reset (GHC.HR = 1, wait for it to clear).
    fn reset(&self) -> Result<(), KernelError> {
        let ghc = self.read_hba_reg(HBA_REG_GHC);
        self.write_hba_reg(HBA_REG_GHC, ghc | GHC_HR);

        // Wait for HR bit to clear (HBA reset complete)
        let mut timeout = POLL_TIMEOUT;
        while self.read_hba_reg(HBA_REG_GHC) & GHC_HR != 0 && timeout > 0 {
            timeout -= 1;
            core::hint::spin_loop();
        }

        if timeout == 0 {
            return Err(KernelError::Timeout {
                operation: "ahci_hba_reset",
                duration_ms: 0,
            });
        }

        crate::println!("[AHCI] HBA reset complete");
        Ok(())
    }

    /// Detect devices on all implemented ports.
    fn detect_ports(&mut self) {
        let pi = self.read_hba_reg(HBA_REG_PI);
        crate::println!("[AHCI] Ports Implemented bitmask: 0x{:08x}", pi);

        for port_num in 0..32u8 {
            if pi & (1 << port_num) == 0 {
                continue;
            }

            let mut port = AhciPort::new(
                port_num,
                self.mmio_base,
                self.num_cmd_slots,
                self.ncq_supported,
            );
            port.detect_device();

            match port.device_type {
                AhciDeviceType::Sata => {
                    crate::println!("[AHCI] Port {}: SATA drive detected", port_num);
                }
                AhciDeviceType::Satapi => {
                    crate::println!("[AHCI] Port {}: SATAPI device detected", port_num);
                }
                AhciDeviceType::Semb => {
                    crate::println!("[AHCI] Port {}: SEMB detected", port_num);
                }
                AhciDeviceType::PortMultiplier => {
                    crate::println!("[AHCI] Port {}: Port Multiplier detected", port_num);
                }
                AhciDeviceType::Unknown(sig) => {
                    crate::println!(
                        "[AHCI] Port {}: Unknown device (sig=0x{:08x})",
                        port_num,
                        sig
                    );
                }
                AhciDeviceType::None => {}
            }

            self.ports.push(port);
        }
    }

    /// Get all detected SATA ports.
    pub fn sata_ports(&self) -> Vec<&AhciPort> {
        self.ports
            .iter()
            .filter(|p| p.device_type == AhciDeviceType::Sata)
            .collect()
    }

    /// Get a specific port by number.
    pub fn port(&self, port_num: u8) -> Option<&AhciPort> {
        self.ports.iter().find(|p| p.port_num == port_num)
    }

    /// Get a mutable reference to a specific port by number.
    pub fn port_mut(&mut self, port_num: u8) -> Option<&mut AhciPort> {
        self.ports.iter_mut().find(|p| p.port_num == port_num)
    }

    /// Get the number of detected ports.
    pub fn port_count(&self) -> usize {
        self.ports.len()
    }

    /// Get the AHCI version as (major, minor).
    pub fn version(&self) -> (u16, u16) {
        ((self.version >> 16) as u16, self.version as u16)
    }
}

// ---------------------------------------------------------------------------
// BlockDevice trait implementation for AHCI ports
// ---------------------------------------------------------------------------

/// Wrapper around an AhciPort for the BlockDevice trait.
///
/// In a production system, this would hold references to DMA-allocated
/// command list and received FIS areas. For now, it provides the structural
/// interface.
#[cfg(feature = "alloc")]
pub struct AhciBlockDevice {
    /// Port descriptor (contains MMIO base, device info).
    port: AhciPort,
    /// Device name (e.g., "sda").
    name: String,
}

#[cfg(feature = "alloc")]
impl AhciBlockDevice {
    /// Create a new AHCI block device from a port.
    pub fn new(port: AhciPort, name: String) -> Self {
        Self { port, name }
    }

    /// Get a reference to the underlying port.
    pub fn port(&self) -> &AhciPort {
        &self.port
    }
}

#[cfg(feature = "alloc")]
impl BlockDevice for AhciBlockDevice {
    fn name(&self) -> &str {
        &self.name
    }

    fn block_size(&self) -> usize {
        self.port.sector_size
    }

    fn block_count(&self) -> u64 {
        self.port.total_sectors
    }

    fn read_blocks(&self, start_block: u64, buffer: &mut [u8]) -> Result<(), KernelError> {
        if buffer.is_empty() {
            return Ok(());
        }

        if !buffer.len().is_multiple_of(self.port.sector_size) {
            return Err(KernelError::InvalidArgument {
                name: "buffer_length",
                value: "not_multiple_of_sector_size",
            });
        }

        if self.port.total_sectors > 0 {
            let sector_count = (buffer.len() / self.port.sector_size) as u64;
            if start_block + sector_count > self.port.total_sectors {
                return Err(KernelError::InvalidArgument {
                    name: "block_range",
                    value: "out_of_bounds",
                });
            }
        }

        // In a full implementation:
        // 1. Find a free command slot
        // 2. Build the READ DMA EXT command (FIS + PRDT pointing to buffer's phys addr)
        // 3. Write the command header into the command list
        // 4. Issue the command and wait for completion
        // 5. Data arrives in the buffer via DMA
        //
        // For now, report not implemented since DMA buffer allocation is required.
        Err(KernelError::NotImplemented {
            feature: "ahci_read_dma",
        })
    }

    fn write_blocks(&mut self, start_block: u64, buffer: &[u8]) -> Result<(), KernelError> {
        if buffer.is_empty() {
            return Ok(());
        }

        if !buffer.len().is_multiple_of(self.port.sector_size) {
            return Err(KernelError::InvalidArgument {
                name: "buffer_length",
                value: "not_multiple_of_sector_size",
            });
        }

        if self.port.total_sectors > 0 {
            let sector_count = (buffer.len() / self.port.sector_size) as u64;
            if start_block + sector_count > self.port.total_sectors {
                return Err(KernelError::InvalidArgument {
                    name: "block_range",
                    value: "out_of_bounds",
                });
            }
        }

        // Same as read_blocks -- requires DMA buffer allocation.
        Err(KernelError::NotImplemented {
            feature: "ahci_write_dma",
        })
    }

    fn flush(&mut self) -> Result<(), KernelError> {
        // Would issue ATA FLUSH CACHE EXT command.
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Global initialization
// ---------------------------------------------------------------------------

/// Whether the AHCI subsystem has been initialized.
static AHCI_INITIALIZED: AtomicBool = AtomicBool::new(false);

/// Discover and initialize AHCI controllers via PCI bus enumeration.
///
/// Scans the PCI bus for Mass Storage controllers with SATA subclass and
/// AHCI programming interface (class 0x01, subclass 0x06, prog-if 0x01).
pub fn init() -> Result<(), KernelError> {
    if AHCI_INITIALIZED.swap(true, Ordering::SeqCst) {
        return Ok(());
    }

    crate::println!("[AHCI] Scanning PCI bus for AHCI controllers...");

    #[cfg(target_arch = "x86_64")]
    {
        let pci_bus = crate::drivers::pci::get_pci_bus().lock();
        let storage_devices =
            pci_bus.find_devices_by_class(crate::drivers::pci::class_codes::MASS_STORAGE);

        let mut ahci_count = 0u32;
        for dev in &storage_devices {
            if dev.subclass == PCI_SUBCLASS_SATA && dev.prog_if == PCI_PROGIF_AHCI {
                ahci_count += 1;
                crate::println!(
                    "[AHCI] Found AHCI controller: {:04x}:{:04x} at {}:{}.{}",
                    dev.vendor_id,
                    dev.device_id,
                    dev.location.bus,
                    dev.location.device,
                    dev.location.function,
                );

                // AHCI uses BAR5 (ABAR) for MMIO registers
                if let Some(bar5) = dev.bars.get(5) {
                    match bar5 {
                        crate::drivers::pci::PciBar::Memory { address, size, .. } => {
                            crate::println!(
                                "[AHCI]   BAR5 (ABAR): MMIO at {:#x}, size {:#x}",
                                address,
                                size
                            );

                            #[cfg(target_os = "none")]
                            {
                                let virt_base = crate::mm::phys_to_virt_addr(*address) as usize;
                                match AhciController::new(virt_base) {
                                    Ok(ctrl) => {
                                        let (maj, min) = ctrl.version();
                                        crate::println!(
                                            "[AHCI] Controller v{}.{} with {} port(s)",
                                            maj,
                                            min,
                                            ctrl.port_count()
                                        );
                                    }
                                    Err(e) => {
                                        crate::println!("[AHCI] Controller init failed: {:?}", e);
                                    }
                                }
                            }
                        }
                        crate::drivers::pci::PciBar::Io { address, size } => {
                            crate::println!(
                                "[AHCI]   BAR5: I/O at {:#x}, size {:#x} (unexpected)",
                                address,
                                size
                            );
                        }
                        crate::drivers::pci::PciBar::None => {
                            crate::println!("[AHCI]   BAR5: not configured");
                        }
                    }
                } else {
                    crate::println!("[AHCI]   No BAR5 found (need at least 6 BARs)");
                }
            }
        }

        if ahci_count == 0 {
            crate::println!("[AHCI] No AHCI controllers found on PCI bus");
        } else {
            crate::println!("[AHCI] Found {} AHCI controller(s)", ahci_count);
        }
    }

    #[cfg(not(target_arch = "x86_64"))]
    {
        crate::println!("[AHCI] AHCI PCI scanning not available on this architecture");
    }

    Ok(())
}

/// Check whether the AHCI subsystem has been initialized.
pub fn is_initialized() -> bool {
    AHCI_INITIALIZED.load(Ordering::Relaxed)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_fis_reg_h2d_size() {
        assert_eq!(core::mem::size_of::<FisRegH2D>(), 20);
    }

    #[test]
    fn test_fis_reg_d2h_size() {
        assert_eq!(core::mem::size_of::<FisRegD2H>(), 20);
    }

    #[test]
    fn test_fis_dma_setup_size() {
        assert_eq!(core::mem::size_of::<FisDmaSetup>(), 28);
    }

    #[test]
    fn test_fis_pio_setup_size() {
        assert_eq!(core::mem::size_of::<FisPioSetup>(), 20);
    }

    #[test]
    fn test_prdt_entry_size() {
        assert_eq!(core::mem::size_of::<PrdtEntry>(), 16);
    }

    #[test]
    fn test_command_header_size() {
        assert_eq!(core::mem::size_of::<CommandHeader>(), 32);
    }

    #[test]
    fn test_command_table_header_size() {
        assert_eq!(core::mem::size_of::<CommandTableHeader>(), 128);
    }

    #[test]
    fn test_received_fis_size() {
        assert_eq!(core::mem::size_of::<ReceivedFis>(), 256);
    }

    #[test]
    fn test_device_type_from_signature() {
        assert_eq!(
            AhciDeviceType::from_signature(SIG_SATA),
            AhciDeviceType::Sata
        );
        assert_eq!(
            AhciDeviceType::from_signature(SIG_SATAPI),
            AhciDeviceType::Satapi
        );
        assert_eq!(
            AhciDeviceType::from_signature(SIG_SEMB),
            AhciDeviceType::Semb
        );
        assert_eq!(
            AhciDeviceType::from_signature(SIG_PM),
            AhciDeviceType::PortMultiplier
        );
        assert_eq!(
            AhciDeviceType::from_signature(0xFFFFFFFF),
            AhciDeviceType::None
        );
        assert_eq!(
            AhciDeviceType::from_signature(0x00000000),
            AhciDeviceType::None
        );
        assert_eq!(
            AhciDeviceType::from_signature(0xDEADBEEF),
            AhciDeviceType::Unknown(0xDEADBEEF)
        );
    }

    #[test]
    fn test_fis_h2d_lba_encoding() {
        let mut fis = FisRegH2D::new();
        fis.set_lba(0x0000_1234_5678_9ABC);

        assert_eq!(fis.lba0, 0xBC);
        assert_eq!(fis.lba1, 0x9A);
        assert_eq!(fis.lba2, 0x78);
        assert_eq!(fis.lba3, 0x56);
        assert_eq!(fis.lba4, 0x34);
        assert_eq!(fis.lba5, 0x12);
        assert_eq!(fis.device & (1 << 6), 1 << 6); // LBA mode bit
    }

    #[test]
    fn test_fis_h2d_count_encoding() {
        let mut fis = FisRegH2D::new();
        fis.set_count(0xABCD);

        assert_eq!(fis.count_lo, 0xCD);
        assert_eq!(fis.count_hi, 0xAB);
    }

    #[test]
    fn test_prdt_entry_construction() {
        let prdt = PrdtEntry::new(0x1000_0000_DEAD_0000, 4096, true);

        assert_eq!(prdt.data_base_lo, 0xDEAD_0000);
        assert_eq!(prdt.data_base_hi, 0x1000_0000);
        // byte_count - 1 = 4095 = 0xFFF, with IOC bit 31 set
        assert_eq!(prdt.dbc_and_flags, 0x8000_0FFF);
    }

    #[test]
    fn test_prdt_entry_no_interrupt() {
        let prdt = PrdtEntry::new(0x0000_0000_0001_0000, 512, false);

        assert_eq!(prdt.data_base_lo, 0x0001_0000);
        assert_eq!(prdt.data_base_hi, 0);
        // 512 - 1 = 511 = 0x1FF, no IOC
        assert_eq!(prdt.dbc_and_flags, 0x0000_01FF);
    }

    #[test]
    fn test_command_header_construction() {
        let hdr = CommandHeader::new(5, 3, true, 0x0000_0001_0000_0080);

        // flags: FIS len = 5, Write bit (6), PRDT count = 3 in bits 31:16
        let expected_flags = 5 | (1 << 6) | (3 << 16);
        assert_eq!(hdr.flags_and_prdtl, expected_flags);
        assert_eq!(hdr.ctba_lo, 0x0000_0080);
        assert_eq!(hdr.ctba_hi, 0x0000_0001);
        assert_eq!(hdr.prdbc, 0);
    }

    #[test]
    fn test_ncq_read_fis_tag_encoding() {
        let fis = build_ncq_read_fis(0, 8, 5);

        assert_eq!(fis.command, ATA_CMD_READ_FPDMA_QUEUED);
        // Tag 5 in bits 7:3 of count_lo
        assert_eq!(fis.count_lo, 5 << 3);
        // Sector count in feature register
        assert_eq!(fis.feature_lo, 8);
        assert_eq!(fis.feature_hi, 0);
    }

    #[test]
    fn test_ncq_write_fis_tag_encoding() {
        let fis = build_ncq_write_fis(100, 256, 31);

        assert_eq!(fis.command, ATA_CMD_WRITE_FPDMA_QUEUED);
        assert_eq!(fis.count_lo, 31 << 3);
        assert_eq!(fis.feature_lo, 0); // 256 & 0xFF = 0
        assert_eq!(fis.feature_hi, 1); // 256 >> 8 = 1
    }

    #[test]
    fn test_port_state_default() {
        let port = AhciPort::new(0, 0x1000_0000, 32, true);

        assert_eq!(port.port_num, 0);
        assert_eq!(port.device_type, AhciDeviceType::None);
        assert_eq!(port.state, PortState::Inactive);
        assert_eq!(port.num_cmd_slots, 32);
        assert!(port.ncq_supported);
        assert_eq!(port.total_sectors, 0);
        assert_eq!(port.sector_size, SECTOR_SIZE);
        // Port 0 MMIO base = hba_base + 0x100 + 0 * 0x80
        assert_eq!(port.port_mmio_base, 0x1000_0000 + PORT_BASE);
    }

    #[test]
    fn test_port_mmio_base_calculation() {
        let port3 = AhciPort::new(3, 0x2000_0000, 16, false);
        // Port 3: base + 0x100 + 3 * 0x80 = base + 0x100 + 0x180 = base + 0x280
        assert_eq!(port3.port_mmio_base, 0x2000_0000 + 0x100 + 3 * 0x80);
    }

    #[test]
    fn test_initialization_flag() {
        // Reset for test isolation
        AHCI_INITIALIZED.store(false, Ordering::SeqCst);
        assert!(!is_initialized());
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn test_ahci_block_device_name() {
        let port = AhciPort::new(0, 0, 32, false);
        let dev = AhciBlockDevice::new(port, String::from("sda"));
        assert_eq!(dev.name(), "sda");
        assert_eq!(dev.block_size(), SECTOR_SIZE);
        assert_eq!(dev.block_count(), 0);
    }
}