//! PCI/PCIe bus interfaces
#![no_std]
#![feature(let_chains, decl_macro)]
#![allow(clippy::missing_transmute_annotations, clippy::identity_op)]

extern crate alloc;

use core::fmt;

#[cfg(any(target_arch = "x86_64", rust_analyzer))]
use acpi::mcfg::McfgEntry;
use alloc::{format, sync::Arc, vec::Vec};

use bitflags::bitflags;
use device::{PciBusDevice, PciDeviceInfo, PciDeviceState};
use device_api::{device::DeviceInitContext, dma::DmaAllocator};
use interrupt::{PciInterruptMap, PciMsiMap};
use libk::{
    dma::DummyDmaAllocator,
    fs::sysfs::{self, object::KObject},
};
use libk_mm::address::PhysicalAddress;
use libk_util::{sync::IrqSafeSpinlock, OneTimeInit};
use space::legacy;
use yggdrasil_abi::{error::Error, primitive_enum};

#[cfg(target_arch = "x86_64")]
use device_api::interrupt::MessageInterruptController;

pub mod capability;
pub mod device;
pub mod driver;
pub mod interrupt;
pub mod macros;
mod nodes;
mod space;

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

bitflags! {
    /// Command register of the PCI configuration space
    #[derive(PartialEq, Clone, Copy)]
    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, PartialOrd, Ord, 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, Clone, Copy)]
pub enum PciBaseAddress {
    /// 32-bit memory address
    Memory32(u32),
    /// 64-bit memory address
    Memory64(u64),
    /// I/O space address
    Io(u16),
}

primitive_enum! {
    pub enum PciCapabilityId: u8 {
        PowerManagement = 0x01,
        Msi = 0x05,
        VendorSpecific = 0x09,
        PciExpress = 0x10,
        MsiX = 0x11,
    }
}

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

    fn check<S: PciConfigurationSpace + ?Sized>(space: &S, offset: usize, len: usize) -> bool {
        true
    }

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

struct BusAddressAllocator {
    pci_base_64: u64,
    pci_base_32: u32,
    // pci_base_io: u16,
    host_base_64: PhysicalAddress,
    host_base_32: PhysicalAddress,
    // host_base_io: PhysicalAddress,
    size_64: usize,
    size_32: usize,
    // size_io: usize,
    offset_64: u64,
    offset_32: u32,
}

#[derive(Debug)]
pub struct PciSegmentInfo {
    pub segment_number: u8,
    pub bus_number_start: u8,
    pub bus_number_end: u8,
    pub ecam_phys_base: Option<PhysicalAddress>,

    pub irq_translation_map: PciInterruptMap,
    pub msi_translation_map: PciMsiMap,
}

/// Represents a single PCIe bus segment
pub struct PciBusSegment {
    allocator: Option<BusAddressAllocator>,
    info: Arc<PciSegmentInfo>,
    devices: Vec<Arc<KObject<IrqSafeSpinlock<PciBusDevice>>>>,
}

#[derive(Debug)]
pub enum PciRangeType {
    Configuration,
    Io,
    Memory32,
    Memory64,
}

pub struct PciAddressRange {
    pub ty: PciRangeType,
    pub bus_number: u8,
    pub pci_base: u64,
    pub host_base: PhysicalAddress,
    pub size: usize,
}

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

#[cfg_attr(
    any(target_arch = "x86_64", target_arch = "x86", target_arch = "riscv64"),
    allow(dead_code)
)]
impl BusAddressAllocator {
    pub fn from_ranges(ranges: &[PciAddressRange]) -> Self {
        let mut range_32 = None;
        let mut range_64 = None;
        // let mut range_io = None;

        for range in ranges {
            let range_val = (range.pci_base, range.host_base, range.size);
            match range.ty {
                // PciRangeType::Io if range_io.is_none() => {
                //     range_io.replace(range_val);
                // }
                PciRangeType::Memory32 if range_32.is_none() => {
                    range_32.replace(range_val);
                }
                PciRangeType::Memory64 if range_64.is_none() => {
                    range_64.replace(range_val);
                }
                _ => (),
            }
        }

        let (pci_base_32, host_base_32, size_32) = range_32.unwrap();
        let (pci_base_64, host_base_64, size_64) = range_64.unwrap();
        // let (pci_base_io, host_base_io, size_io) = range_io.unwrap();

        Self {
            pci_base_64,
            pci_base_32: pci_base_32.try_into().unwrap(),
            // pci_base_io: pci_base_io.try_into().unwrap(),
            host_base_64,
            host_base_32,
            // host_base_io,
            size_64,
            size_32,
            // size_io,
            offset_64: 0,
            offset_32: 0,
        }
    }

    pub fn allocate(&mut self, ty: PciRangeType, size: usize) -> (PciBaseAddress, PhysicalAddress) {
        match ty {
            PciRangeType::Io => todo!(),
            PciRangeType::Memory32 => {
                if self.offset_32 as usize + size >= self.size_32 {
                    todo!();
                }
                let bar = PciBaseAddress::Memory32(self.pci_base_32 + self.offset_32);
                let host = self.host_base_32.add(self.offset_32 as usize);
                self.offset_32 += size as u32;
                (bar, host)
            }
            PciRangeType::Memory64 => {
                if self.offset_64 as usize + size >= self.size_64 {
                    todo!();
                }
                let bar = PciBaseAddress::Memory64(self.pci_base_64 + self.offset_64);
                let host = self.host_base_64.add(self.offset_64 as usize);
                self.offset_64 += size as u64;
                (bar, host)
            }
            PciRangeType::Configuration => unimplemented!(),
        }
    }
}

impl PciBaseAddress {
    pub fn as_memory(self) -> Option<PhysicalAddress> {
        match self {
            Self::Memory32(address) => Some(PhysicalAddress::from_u64(address as u64)),
            Self::Memory64(address) => Some(PhysicalAddress::from_u64(address)),
            _ => None,
        }
    }

    pub fn is_zero(&self) -> bool {
        match *self {
            Self::Memory32(base) => base == 0,
            Self::Memory64(base) => base == 0,
            Self::Io(base) => base == 0,
        }
    }
}

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

    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
        // }

        if let Some(allocator) = self.allocator.as_mut() {
            log::debug!("Remapping BARs for {}", address);

            // Find valid BARs
            let mut i = 0;
            let mut bar_mask = 0;

            while i < 6 {
                let w0 = config.read_u32(0x10 + i * 4);

                let bar_width = match w0 & 1 == 0 {
                    // Memory BAR
                    true => match (w0 >> 1) & 3 {
                        // 32-bit BAR
                        0 => 1,
                        // Reserved
                        1 => unimplemented!(),
                        // 64-bit BAR
                        2 => 2,
                        // Unknown
                        _ => unreachable!(),
                    },
                    false => 1,
                };

                bar_mask |= 1 << i;
                i += bar_width;
            }

            for i in 0..6 {
                if (1 << i) & bar_mask != 0 {
                    let Some(orig_value) = config.bar(i) else {
                        continue;
                    };
                    let size = unsafe { config.bar_size(i) };

                    if size != 0 {
                        log::debug!("BAR{}: size={:#x}", i, size);

                        match orig_value {
                            PciBaseAddress::Io(_) => (),
                            PciBaseAddress::Memory64(_) => {
                                let (bar, host) = allocator.allocate(PciRangeType::Memory64, size);
                                let bar_address = bar.as_memory().unwrap();
                                unsafe {
                                    config.set_bar(i, bar);
                                }
                                log::debug!(
                                    "Mapped BAR{} -> pci {:#x} host {:#x}",
                                    i,
                                    bar_address,
                                    host
                                );
                                // TODO Don't yet differentiate between Host/PCI addresses, lol
                                assert_eq!(bar_address, host);
                            }
                            PciBaseAddress::Memory32(_) => {
                                let (bar, host) = allocator.allocate(PciRangeType::Memory32, size);
                                let bar_address = bar.as_memory().unwrap();
                                unsafe {
                                    config.set_bar(i, bar);
                                }
                                log::debug!(
                                    "Mapped BAR{} -> pci {:#x} host {:#x}",
                                    i,
                                    bar_address,
                                    host
                                );
                                // TODO Don't yet differentiate between Host/PCI addresses, lol
                                assert_eq!(bar_address, host);
                            }
                        }
                    }
                }
            }
        }

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

        let info = PciDeviceInfo {
            address,
            vendor_id,
            device_id,
            class,
            subclass,
            prog_if,
            segment: self.info.clone(),
            config_space: config,
            interrupt_config: Arc::new(OneTimeInit::new()),
        };

        let object = nodes::make_sysfs_object(PciBusDevice {
            info,
            driver: None,
            device: None,
            state: PciDeviceState::None,
        });
        let pci_object = PCI_SYSFS_NODE.or_init_with(|| {
            let bus_object = sysfs::bus().unwrap();
            let pci_object = KObject::new(());
            bus_object.add_object("pci", pci_object.clone()).ok();
            pci_object
        });

        let name = format!("{address}");
        pci_object.add_object(name, object.clone()).ok();
        self.devices.push(object);

        Ok(())
    }

    fn enumerate_bus(&mut self, bus: u8) -> Result<(), Error> {
        let address = PciAddress::for_bus(self.info.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.info.bus_number_start..self.info.bus_number_end {
            self.enumerate_bus(bus)?;
        }
        Ok(())
    }
}

impl PciSegmentInfo {
    pub fn has_msi(&self) -> bool {
        !matches!(self.msi_translation_map, PciMsiMap::Legacy)
    }
}

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 probe_bus_devices() -> Result<(), Error> {
        Self::walk_bus_devices(|device| {
            probe_bus_device(device, false)?;
            Ok(true)
        })
    }

    pub fn setup_bus_devices(rescan: bool) -> Result<(), Error> {
        Self::walk_bus_devices(|device| {
            setup_bus_device(device, rescan)?;
            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() {
                let mut device = device.lock();
                if !f(&mut device)? {
                    return Ok(());
                }
            }
        }

        Ok(())
    }

    pub fn add_legacy_segment(access: &'static dyn LegacyPciAccess) -> Result<(), Error> {
        legacy::PCI.init(access);

        let mut bus_segment = PciBusSegment {
            info: Arc::new(PciSegmentInfo {
                segment_number: 0,
                bus_number_start: 0,
                bus_number_end: 255,
                ecam_phys_base: None,
                irq_translation_map: PciInterruptMap::Legacy,
                msi_translation_map: PciMsiMap::Legacy,
            }),
            allocator: None,
            devices: Vec::new(),
        };

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

        Ok(())
    }

    /// Enumerates a bus segment provided by ACPI MCFG table entry
    #[cfg(target_arch = "x86_64")]
    pub fn add_segment_from_mcfg(
        entry: &McfgEntry,
        msi_controller: Arc<dyn MessageInterruptController>,
    ) -> Result<(), Error> {
        let msi_translation_map = PciMsiMap::Identity(msi_controller);

        let mut bus_segment = PciBusSegment {
            info: Arc::new(PciSegmentInfo {
                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_u64(entry.base_address)),

                // TODO get the segment's PCI root bridge AML name
                irq_translation_map: PciInterruptMap::Acpi("\\_SB.PCI0._PRT".into()),
                msi_translation_map,
            }),
            // Firmware done this for us
            allocator: None,

            devices: Vec::new(),
        };

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

        Ok(())
    }

    #[cfg(any(target_arch = "aarch64", target_arch = "riscv64", rust_analyzer))]
    pub fn add_segment_from_device_tree(
        cfg_base: PhysicalAddress,
        bus_range: core::ops::Range<u8>,
        ranges: Vec<PciAddressRange>,
        irq_translation_map: PciInterruptMap,
        msi_translation_map: PciMsiMap,
    ) -> Result<(), Error> {
        let mut bus_segment = PciBusSegment {
            info: Arc::new(PciSegmentInfo {
                segment_number: 0,
                bus_number_start: bus_range.start,
                bus_number_end: bus_range.end,
                ecam_phys_base: Some(cfg_base),

                irq_translation_map,
                msi_translation_map,
            }),
            allocator: Some(BusAddressAllocator::from_ranges(&ranges)),

            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 specific function
    pub const fn for_function(segment: u8, bus: u8, device: u8, function: u8) -> Self {
        Self {
            segment,
            bus,
            device,
            function,
        }
    }

    /// 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),
            Self::Legacy(legacy) => legacy.read_u32(offset),
        }
    }

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

fn probe_bus_device(device: &mut PciBusDevice, _rescan: bool) -> Result<(), Error> {
    // Already has a driver/device set up
    if device.device.is_some() || device.state != PciDeviceState::None {
        return Ok(());
    }

    if let Some(driver) = driver::lookup_driver(&device.info) {
        let dma: Arc<dyn DmaAllocator> = Arc::new(DummyDmaAllocator);

        match driver.probe(&device.info, &dma) {
            Ok(instance) => {
                log::info!("{} -> {}", device.info.address, driver.driver_name());
                device.device.replace(instance);
                device.driver.replace(driver);
                device.state = PciDeviceState::Probed;
            }
            Err(error) => {
                log::error!(
                    "{} ({}) probe error: {error:?}",
                    device.info.address,
                    driver.driver_name()
                );
            }
        }
    }

    Ok(())
}

fn setup_bus_device(device: &mut PciBusDevice, _rescan: bool) -> Result<(), Error> {
    // No driver yet (TODO probe if rescan is asked)
    let (Some(dev), Some(driver)) = (device.device.as_ref(), device.driver) else {
        return Ok(());
    };
    // Already initialized/failed
    if device.state != PciDeviceState::Probed {
        return Ok(());
    }

    let dma: Arc<dyn DmaAllocator> = Arc::new(DummyDmaAllocator);
    let cx = DeviceInitContext {
        dma_allocator: dma.clone(),
    };

    match unsafe { dev.clone().init(cx) } {
        Ok(()) => {
            device.state = PciDeviceState::Initialized;
        }
        Err(error) => {
            log::error!(
                "{} ({}) setup error: {error:?}",
                device.info.address,
                driver.driver_name()
            );
            device.state = PciDeviceState::Failed;
        }
    }

    Ok(())
}

static PCI_MANAGER: IrqSafeSpinlock<PciBusManager> = IrqSafeSpinlock::new(PciBusManager::new());
static PCI_SYSFS_NODE: OneTimeInit<Arc<KObject<()>>> = OneTimeInit::new();
