//! PCI bus enumeration (the legacy "configuration mechanism #1").
//!
//! Every PCI function has a 256-byte *configuration space*: vendor/device
//! ids, class code, BARs (where its registers live), the command register...
//! On x86 we reach it through two I/O ports: write the address
//! (bus, device, function, register) to 0xCF8, then read/write the data at
//! 0xCFC. Enumerating = asking every possible (bus, device, function) for its
//! vendor id; 0xFFFF means "nobody here".
//!
//! Modern machines also offer memory-mapped config space (PCIe ECAM, found
//! via ACPI's MCFG table); the port method is simpler and QEMU supports both.
//! Reference: https://wiki.osdev.org/PCI

use alloc::vec::Vec;
use x86_64::instructions::port::Port;

const CONFIG_ADDRESS: u16 = 0xCF8;
const CONFIG_DATA: u16 = 0xCFC;

#[derive(Debug, Clone, Copy)]
pub struct PciDevice {
    pub bus: u8,
    pub device: u8,
    pub function: u8,
    pub vendor_id: u16,
    pub device_id: u16,
    pub class: u8,
    pub subclass: u8,
}

/// Read a 32-bit register (`offset` must be 4-byte aligned).
pub fn read_config(bus: u8, device: u8, function: u8, offset: u8) -> u32 {
    let address = 1u32 << 31 // enable bit
        | (bus as u32) << 16
        | (device as u32) << 11
        | (function as u32) << 8
        | (offset as u32 & 0xfc);
    unsafe {
        Port::<u32>::new(CONFIG_ADDRESS).write(address);
        Port::<u32>::new(CONFIG_DATA).read()
    }
}

pub fn write_config(bus: u8, device: u8, function: u8, offset: u8, value: u32) {
    let address = 1u32 << 31 | (bus as u32) << 16 | (device as u32) << 11 | (function as u32) << 8 | (offset as u32 & 0xfc);
    unsafe {
        Port::<u32>::new(CONFIG_ADDRESS).write(address);
        Port::<u32>::new(CONFIG_DATA).write(value);
    }
}

impl PciDevice {
    fn read(&self, offset: u8) -> u32 {
        read_config(self.bus, self.device, self.function, offset)
    }

    /// Base address register `n` (0-5), raw. Bit 0 set = I/O port space.
    pub fn bar(&self, n: u8) -> u32 {
        self.read(0x10 + 4 * n)
    }

    /// Let the device respond to I/O port accesses and do DMA (bus mastering).
    /// The legacy interrupt line (8259 IRQ number) the firmware routed this
    /// device's INTx pin to (config register 0x3C, low byte).
    pub fn interrupt_line(&self) -> u8 {
        read_config(self.bus, self.device, self.function, 0x3c) as u8
    }

    pub fn enable_io_and_bus_master(&self) {
        let command = self.read(0x04);
        write_config(self.bus, self.device, self.function, 0x04, command | 0b101);
    }

    /// Same for devices whose registers are memory-mapped (MMIO): command
    /// bit 1 = memory space, bit 2 = bus master.
    pub fn enable_memory_and_bus_master(&self) {
        let command = self.read(0x04);
        write_config(self.bus, self.device, self.function, 0x04, command | 0b110);
    }

    /// Physical address of memory BAR `n`. Bits 1-2 of a memory BAR give its
    /// type: 0b10 = 64-bit, and then BAR n+1 holds the upper half. The low 4
    /// bits are flags, not address.
    pub fn bar_address(&self, n: u8) -> u64 {
        let low = self.bar(n);
        if low & 1 != 0 {
            return (low & !0x3) as u64; // an I/O port BAR
        }
        let mut addr = (low & !0xf) as u64;
        if (low >> 1) & 0b11 == 0b10 {
            addr |= (self.bar(n + 1) as u64) << 32;
        }
        addr
    }

    /// One byte of configuration space (any offset).
    pub fn read_u8(&self, offset: u8) -> u8 {
        (self.read(offset & 0xfc) >> ((offset & 3) * 8)) as u8
    }

    /// Read 32 bits at any 4-byte aligned config offset.
    pub fn read_u32(&self, offset: u8) -> u32 {
        self.read(offset)
    }

    /// The *capability list*: a linked list in config space that newer
    /// devices use to describe extra features (MSI, MSI-X, power
    /// management, and -- what we need -- where modern virtio devices put
    /// their register blocks). Status bit 4 says the list exists; register
    /// 0x34 points at the first entry; each entry is [id, next, data...].
    /// Returns the config-space offsets of the entries with id `cap_id`.
    pub fn capabilities(&self, cap_id: u8) -> Vec<u8> {
        let mut found = Vec::new();
        let status = (self.read(0x04) >> 16) as u16;
        if status & (1 << 4) == 0 {
            return found;
        }
        let mut offset = self.read_u8(0x34) & 0xfc;
        let mut guard = 0; // a broken list must not loop forever
        while offset != 0 && guard < 48 {
            if self.read_u8(offset) == cap_id {
                found.push(offset);
            }
            offset = self.read_u8(offset + 1) & 0xfc;
            guard += 1;
        }
        found
    }

    /// Human-readable class (just the common ones).
    pub fn class_name(&self) -> &'static str {
        match (self.class, self.subclass) {
            (0x01, 0x01) => "IDE controller",
            (0x01, 0x00) => "SCSI storage (virtio-blk)",
            (0x01, 0x06) => "SATA controller",
            (0x01, _) => "storage controller",
            (0x02, _) => "network controller",
            (0x03, _) => "display controller",
            (0x00, _) => "unclassified",
            (0x07, _) => "communication controller",
            (0x06, 0x00) => "host bridge",
            (0x06, 0x01) => "ISA bridge",
            (0x06, 0x80) => "bridge (other)",
            (0x06, _) => "bridge",
            (0x0c, 0x03) => "USB controller",
            _ => "other",
        }
    }
}

fn probe(bus: u8, device: u8, function: u8) -> Option<PciDevice> {
    let id = read_config(bus, device, function, 0x00);
    let vendor_id = id as u16;
    if vendor_id == 0xffff {
        return None; // no such device
    }
    let class = read_config(bus, device, function, 0x08);
    Some(PciDevice {
        bus,
        device,
        function,
        vendor_id,
        device_id: (id >> 16) as u16,
        class: (class >> 24) as u8,
        subclass: (class >> 16) as u8,
    })
}

/// Every PCI function on every bus (brute force: 256 buses x 32 devices --
/// a few thousand port reads, instant in QEMU).
pub fn scan() -> Vec<PciDevice> {
    let mut found = Vec::new();
    for bus in 0..=255u8 {
        for device in 0..32u8 {
            let Some(first) = probe(bus, device, 0) else { continue };
            found.push(first);
            // Header type bit 7: multi-function device (functions 1-7 exist).
            let header_type = (read_config(bus, device, 0, 0x0c) >> 16) as u8;
            if header_type & 0x80 != 0 {
                found.extend((1..8).filter_map(|f| probe(bus, device, f)));
            }
        }
    }
    found
}