//! PCI/PCIe bus interfaces
#![no_std]

extern crate alloc;

use core::fmt;

use acpi::mcfg::McfgEntry;
use alloc::{rc::Rc, vec::Vec};
use bitflags::bitflags;
use device_api::Device;
use kernel_util::{
    mem::address::{FromRaw, PhysicalAddress},
    sync::IrqSafeSpinlock,
};
use yggdrasil_abi::error::Error;

pub mod capability;
mod space;

pub use space::{
    ecam::PciEcam, PciConfigSpace, PciConfigurationSpace, PciLegacyConfigurationSpace,
};

bitflags! {
    /// Command register of the PCI configuration space
    pub struct PciCommandRegister: u16 {
        /// If set, I/O access to the device is enabled
        const ENABLE_IO = 1 << 0;
        /// If set, memory-mapped access to the device is enabled
        const ENABLE_MEMORY = 1 << 1;
        /// If set, the device can generate PCI bus accesses on its own
        const BUS_MASTER = 1 << 2;
        /// If set, interrupts are masked from being raised
        const DISABLE_INTERRUPTS = 1 << 10;
    }
}

bitflags! {
    /// Status register of the PCI configuration space
    pub struct PciStatusRegister: u16 {
        /// Read-only. If set, the configuration space has a pointer to the capabilities list.
        const CAPABILITIES_LIST = 1 << 4;
    }
}

/// Represents the address of a single object on a bus (or the bus itself)
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct PciAddress {
    /// PCIe segment group, ignored (?) with PCI
    pub segment: u8,
    /// Bus number
    pub bus: u8,
    /// Slot/device number
    pub device: u8,
    /// Function number
    pub function: u8,
}

/// Address provided by PCI configuration space Base Address Register
#[derive(Debug)]
pub enum PciBaseAddress {
    /// 32/64-bit memory address
    Memory(usize),
    /// I/O space address
    Io(u16),
}

/// Unique ID assigned to PCI capability structures
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[non_exhaustive]
#[repr(u8)]
pub enum PciCapabilityId {
    /// MSI (32-bit or 64-bit)
    Msi = 0x05,
    /// MSI-X
    MsiX = 0x11,
    /// Unknown capability missing from this list
    Unknown,
}

/// Interface used for querying PCI capabilities
pub trait PciCapability {
    /// Capability ID
    const ID: PciCapabilityId;
    /// Wrapper for accessing the capability data structure
    type CapabilityData<'a, S: PciConfigurationSpace + ?Sized + 'a>;

    /// Constructs an access wrapper for this capability with given offset
    fn data<'s, S: PciConfigurationSpace + ?Sized + 's>(
        space: &'s S,
        offset: usize,
    ) -> Self::CapabilityData<'s, S>;
}

/// Describes a PCI device
#[derive(Debug)]
pub struct PciDeviceInfo {
    /// Address of the device
    pub address: PciAddress,
    /// Configuration space access method
    pub config_space: PciConfigSpace,
}

pub enum PciMatch {
    Generic(fn(&PciDeviceInfo) -> bool),
    Class(u8, Option<u8>, Option<u8>),
}

pub struct PciDriver {
    name: &'static str,
    check: PciMatch,
    probe: fn(&PciDeviceInfo) -> Result<&'static dyn Device, Error>,
}

/// Used to store PCI bus devices which were enumerated by the kernel
pub struct PciBusDevice {
    info: PciDeviceInfo,
    driver: Option<Rc<dyn Device>>,
}

/// Represents a single PCIe bus segment
pub struct PciBusSegment {
    segment_number: u8,
    bus_number_start: u8,
    bus_number_end: u8,
    ecam_phys_base: Option<PhysicalAddress>,

    devices: Vec<PciBusDevice>,
}

/// Manager struct to store and control all PCI devices in the system
pub struct PciBusManager {
    segments: Vec<PciBusSegment>,
}

impl PciBusSegment {
    fn probe_config_space(&self, address: PciAddress) -> Result<Option<PciConfigSpace>, Error> {
        match self.ecam_phys_base {
            Some(ecam_phys_base) => Ok(unsafe {
                PciEcam::probe_raw_parts(ecam_phys_base, self.bus_number_start, address)?
            }
            .map(PciConfigSpace::Ecam)),
            None => todo!(),
        }
    }

    fn enumerate_function(&mut self, address: PciAddress) -> Result<(), Error> {
        let Some(config) = self.probe_config_space(address)? else {
            return Ok(());
        };

        let header_type = config.header_type();

        // Enumerate multi-function devices
        if address.function == 0 && header_type & 0x80 != 0 {
            for function in 1..8 {
                self.enumerate_function(address.with_function(function))?;
            }
        }

        // PCI-to-PCI bridge
        // if config.class_code() == 0x06 && config.subclass() == 0x04 {
        //     let secondary_bus = config.secondary_bus();
        //     // TODO
        // }

        let info = PciDeviceInfo {
            address,
            config_space: config,
        };
        self.devices.push(PciBusDevice { info, driver: None });

        Ok(())
    }

    fn enumerate_bus(&mut self, bus: u8) -> Result<(), Error> {
        let address = PciAddress::for_bus(self.segment_number, bus);

        for i in 0..32 {
            let device_address = address.with_device(i);

            self.enumerate_function(device_address)?;
        }

        Ok(())
    }

    /// Enumerates the bus segment, placing found devices into the manager
    pub fn enumerate(&mut self) -> Result<(), Error> {
        for bus in self.bus_number_start..self.bus_number_end {
            self.enumerate_bus(bus)?;
        }
        Ok(())
    }
}

impl PciBusManager {
    const fn new() -> Self {
        Self {
            segments: Vec::new(),
        }
    }

    /// Walks the bus device list and calls init/init_irq functions on any devices with associated
    /// drivers
    pub fn setup_bus_devices() -> Result<(), Error> {
        Self::walk_bus_devices(|device| {
            setup_bus_device(device)?;
            Ok(true)
        })
    }

    /// Iterates over the bus devices, calling the function on each of them until either an error
    /// or `Ok(false)` is returned
    pub fn walk_bus_devices<F: FnMut(&mut PciBusDevice) -> Result<bool, Error>>(
        mut f: F,
    ) -> Result<(), Error> {
        let mut this = PCI_MANAGER.lock();

        for segment in this.segments.iter_mut() {
            for device in segment.devices.iter_mut() {
                if !f(device)? {
                    return Ok(());
                }
            }
        }

        Ok(())
    }

    /// Enumerates a bus segment provided by ACPI MCFG table entry
    pub fn add_segment_from_mcfg(entry: &McfgEntry) -> Result<(), Error> {
        let mut bus_segment = PciBusSegment {
            segment_number: entry.pci_segment_group as u8,
            bus_number_start: entry.bus_number_start,
            bus_number_end: entry.bus_number_end,
            ecam_phys_base: Some(PhysicalAddress::from_raw(entry.base_address)),

            devices: Vec::new(),
        };

        let mut this = PCI_MANAGER.lock();
        bus_segment.enumerate()?;
        this.segments.push(bus_segment);

        Ok(())
    }
}

impl fmt::Display for PciAddress {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}:{}", self.bus, self.device, self.function)
    }
}

impl PciAddress {
    /// Constructs a [PciAddress] representing a bus
    pub const fn for_bus(segment: u8, bus: u8) -> Self {
        Self {
            segment,
            bus,
            device: 0,
            function: 0,
        }
    }

    /// Constructs a [PciAddress] representing a device on a given bus
    pub const fn with_device(self, device: u8) -> Self {
        Self {
            device,
            function: 0,
            ..self
        }
    }

    /// Constructs a [PciAddress] representing a function of a given bus device
    pub const fn with_function(self, function: u8) -> Self {
        Self { function, ..self }
    }
}

impl PciConfigurationSpace for PciConfigSpace {
    fn read_u32(&self, offset: usize) -> u32 {
        match self {
            Self::Ecam(ecam) => ecam.read_u32(offset),
            _ => todo!(),
        }
    }

    fn write_u32(&self, offset: usize, value: u32) {
        match self {
            Self::Ecam(ecam) => ecam.write_u32(offset, value),
            _ => todo!(),
        }
    }
}

fn setup_bus_device(device: &mut PciBusDevice) -> Result<(), Error> {
    if device.driver.is_some() {
        return Ok(());
    }

    let config = &device.info.config_space;

    log::debug!(
        "{}: {:04x}:{:04x}",
        device.info.address,
        config.vendor_id(),
        config.device_id()
    );

    let class = config.class_code();
    let subclass = config.subclass();
    let prog_if = config.prog_if();

    let drivers = PCI_DRIVERS.lock();
    for driver in drivers.iter() {
        if driver
            .check
            .check_device(&device.info, class, subclass, prog_if)
        {
            // TODO add the device to the bus
            log::debug!(" -> {:?}", driver.name);
            let device = (driver.probe)(&device.info)?;

            unsafe {
                device.init()?;
            }
        } else {
            log::debug!(" -> No driver");
        }
    }

    Ok(())
}

impl PciMatch {
    pub fn check_device(&self, info: &PciDeviceInfo, class: u8, subclass: u8, prog_if: u8) -> bool {
        match self {
            Self::Generic(f) => f(info),
            &Self::Class(class_, Some(subclass_), Some(prog_if_)) => {
                class_ == class && subclass_ == subclass && prog_if_ == prog_if
            }
            &Self::Class(class_, Some(subclass_), _) => class_ == class && subclass_ == subclass,
            &Self::Class(class_, _, _) => class_ == class,
        }
    }
}

pub fn register_class_driver(
    name: &'static str,
    class: u8,
    subclass: Option<u8>,
    prog_if: Option<u8>,
    probe: fn(&PciDeviceInfo) -> Result<&'static dyn Device, Error>,
) {
    PCI_DRIVERS.lock().push(PciDriver {
        name,
        check: PciMatch::Class(class, subclass, prog_if),
        probe,
    });
}

pub fn register_generic_driver(
    name: &'static str,
    check: fn(&PciDeviceInfo) -> bool,
    probe: fn(&PciDeviceInfo) -> Result<&'static dyn Device, Error>,
) {
    PCI_DRIVERS.lock().push(PciDriver {
        name,
        check: PciMatch::Generic(check),
        probe,
    });
}

static PCI_DRIVERS: IrqSafeSpinlock<Vec<PciDriver>> = IrqSafeSpinlock::new(Vec::new());
static PCI_MANAGER: IrqSafeSpinlock<PciBusManager> = IrqSafeSpinlock::new(PciBusManager::new());