//! PCI capability structures and queries

use alloc::{sync::Arc, vec, vec::Vec};
use bitflags::bitflags;
use device_api::interrupt::{
    InterruptAffinity, InterruptHandler, MessageInterruptController, MsiInfo,
};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIoMut};
use tock_registers::{
    interfaces::{Readable, Writeable},
    registers::{ReadWrite, WriteOnly},
};
use yggdrasil_abi::error::Error;

use crate::PciBaseAddress;

use super::{PciCapability, PciCapabilityId, PciConfigurationSpace};

bitflags! {
    pub struct PcieLinkControl: u16 {
        const ASPM_DISABLE = 0 << 0;
        // Active state power management control
        const ASPM_MASK = 0x3 << 0;
        // Enable clock power management
        const ECPM = 1 << 8;
    }
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
use core::mem::offset_of;
#[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
use kernel_arch_x86::intrinsics;

pub trait VirtioCapabilityData<'s, S: PciConfigurationSpace + ?Sized + 's>: Sized {
    fn from_space_offset(space: &'s S, offset: usize) -> Self;

    fn space(&self) -> &'s S;
    fn offset(&self) -> usize;

    fn bar_index(&self) -> Option<usize> {
        let value = self.space().read_u8(self.offset() + 4);
        (value <= 0x5).then_some(value as _)
    }

    fn bar_offset(&self) -> usize {
        let value = self.space().read_u32(self.offset() + 8);
        value as _
    }

    fn length(&self) -> usize {
        let value = self.space().read_u32(self.offset() + 12);
        value as _
    }
}

pub trait VirtioCapability {
    const CFG_TYPE: u8;
    const MIN_LEN: usize = 0;
    type Output<'a, S: PciConfigurationSpace + ?Sized + 'a>: VirtioCapabilityData<'a, S>;
}

/// Power management capability entry
pub struct PowerManagementCapability;
/// MSI-X capability query
pub struct MsiXCapability;
/// MSI capability query
pub struct MsiCapability;
/// PCIe capability
pub struct PciExpressCapability;

// VirtIO-over-PCI capabilities
/// VirtIO PCI configuration access
pub struct VirtioDeviceConfigCapability;
/// VirtIO common configuration
pub struct VirtioCommonConfigCapability;
/// VirtIO notify configuration
pub struct VirtioNotifyConfigCapability;
/// VirtIO interrupt status
pub struct VirtioInterruptStatusCapability;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DevicePowerState {
    D0,
    D1,
    D2,
    D3Cold,
    D3Hot,
}

/// Represents an entry in MSI-X vector table
#[repr(C)]
pub struct MsiXEntry {
    /// Address to which the value is written on interrupt
    pub address: WriteOnly<u64>,
    /// Value which is written to trigger an interrupt
    pub data: WriteOnly<u32>,
    /// Vector control word
    pub control: ReadWrite<u32>,
}

enum MsiXVectorTableAccess<'a> {
    Memory(DeviceMemoryIoMut<'a, [MsiXEntry]>),
    #[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
    Io(u16),
}

pub struct MsiXVectorTable<'a> {
    access: MsiXVectorTableAccess<'a>,
    len: usize,
}

/// PCI Power Management capability data structure
pub struct PowerManagementData<'s, S: PciConfigurationSpace + ?Sized + 's> {
    space: &'s S,
    offset: usize,
}

/// MSI-X capability data structure
pub struct MsiXData<'s, S: PciConfigurationSpace + ?Sized + 's> {
    space: &'s S,
    offset: usize,
}

/// MSI capability data structure
pub struct MsiData<'s, S: PciConfigurationSpace + ?Sized + 's> {
    space: &'s S,
    offset: usize,
}

/// PCI Express capability data structure
pub struct PcieData<'s, S: PciConfigurationSpace + ?Sized + 's> {
    space: &'s S,
    offset: usize,
}

pub struct VirtioDeviceConfigData<'s, S: PciConfigurationSpace + ?Sized + 's> {
    space: &'s S,
    offset: usize,
}

pub struct VirtioCommonConfigData<'s, S: PciConfigurationSpace + ?Sized + 's> {
    space: &'s S,
    offset: usize,
}

pub struct VirtioNotifyConfigData<'s, S: PciConfigurationSpace + ?Sized + 's> {
    space: &'s S,
    offset: usize,
}

pub struct VirtioInterruptStatusData<'s, S: PciConfigurationSpace + ?Sized + 's> {
    space: &'s S,
    offset: usize,
}

impl<T: VirtioCapability> PciCapability for T {
    const ID: PciCapabilityId = PciCapabilityId::VendorSpecific;
    type CapabilityData<'a, S: PciConfigurationSpace + ?Sized + 'a> = T::Output<'a, S>;

    fn check<S: PciConfigurationSpace + ?Sized>(space: &S, offset: usize, len: usize) -> bool {
        let cfg_type = space.read_u8(offset + 3);
        cfg_type == T::CFG_TYPE && len >= T::MIN_LEN
    }

    fn data<'s, S: PciConfigurationSpace + ?Sized + 's>(
        space: &'s S,
        offset: usize,
        _len: usize,
    ) -> Self::CapabilityData<'s, S> {
        T::Output::from_space_offset(space, offset)
    }
}

impl PciCapability for PowerManagementCapability {
    const ID: PciCapabilityId = PciCapabilityId::PowerManagement;
    type CapabilityData<'a, S: PciConfigurationSpace + ?Sized + 'a> = PowerManagementData<'a, S>;

    fn data<'s, S: PciConfigurationSpace + ?Sized + 's>(
        space: &'s S,
        offset: usize,
        _len: usize,
    ) -> Self::CapabilityData<'s, S> {
        PowerManagementData { space, offset }
    }
}

impl PciCapability for MsiXCapability {
    const ID: PciCapabilityId = PciCapabilityId::MsiX;
    type CapabilityData<'a, S: PciConfigurationSpace + ?Sized + 'a> = MsiXData<'a, S>;

    fn data<'s, S: PciConfigurationSpace + ?Sized + 's>(
        space: &'s S,
        offset: usize,
        _len: usize,
    ) -> Self::CapabilityData<'s, S> {
        MsiXData { space, offset }
    }
}

impl PciCapability for MsiCapability {
    const ID: PciCapabilityId = PciCapabilityId::Msi;
    type CapabilityData<'a, S: PciConfigurationSpace + ?Sized + 'a> = MsiData<'a, S>;

    fn data<'s, S: PciConfigurationSpace + ?Sized + 's>(
        space: &'s S,
        offset: usize,
        _len: usize,
    ) -> Self::CapabilityData<'s, S> {
        MsiData { space, offset }
    }
}

impl PciCapability for PciExpressCapability {
    const ID: PciCapabilityId = PciCapabilityId::PciExpress;
    type CapabilityData<'a, S: PciConfigurationSpace + ?Sized + 'a> = PcieData<'a, S>;

    fn data<'s, S: PciConfigurationSpace + ?Sized + 's>(
        space: &'s S,
        offset: usize,
        _len: usize,
    ) -> Self::CapabilityData<'s, S> {
        PcieData { space, offset }
    }
}

impl VirtioCapability for VirtioDeviceConfigCapability {
    const CFG_TYPE: u8 = 0x04;
    type Output<'a, S: PciConfigurationSpace + ?Sized + 'a> = VirtioDeviceConfigData<'a, S>;
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> VirtioCapabilityData<'s, S>
    for VirtioDeviceConfigData<'s, S>
{
    fn from_space_offset(space: &'s S, offset: usize) -> Self {
        Self { space, offset }
    }

    fn space(&self) -> &'s S {
        self.space
    }

    fn offset(&self) -> usize {
        self.offset
    }
}

impl VirtioCapability for VirtioCommonConfigCapability {
    const CFG_TYPE: u8 = 0x01;
    type Output<'a, S: PciConfigurationSpace + ?Sized + 'a> = VirtioCommonConfigData<'a, S>;
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> VirtioCapabilityData<'s, S>
    for VirtioCommonConfigData<'s, S>
{
    fn from_space_offset(space: &'s S, offset: usize) -> Self {
        Self { space, offset }
    }

    fn space(&self) -> &'s S {
        self.space
    }

    fn offset(&self) -> usize {
        self.offset
    }
}

impl VirtioCapability for VirtioNotifyConfigCapability {
    const CFG_TYPE: u8 = 0x02;
    const MIN_LEN: usize = 0x14;
    type Output<'a, S: PciConfigurationSpace + ?Sized + 'a> = VirtioNotifyConfigData<'a, S>;
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> VirtioNotifyConfigData<'s, S> {
    pub fn offset_multiplier(&self) -> usize {
        self.space.read_u32(self.offset + 16) as usize
    }
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> VirtioCapabilityData<'s, S>
    for VirtioNotifyConfigData<'s, S>
{
    fn from_space_offset(space: &'s S, offset: usize) -> Self {
        Self { space, offset }
    }

    fn space(&self) -> &'s S {
        self.space
    }

    fn offset(&self) -> usize {
        self.offset
    }
}

impl VirtioCapability for VirtioInterruptStatusCapability {
    const CFG_TYPE: u8 = 0x03;
    const MIN_LEN: usize = 1;
    type Output<'a, S: PciConfigurationSpace + ?Sized + 'a> = VirtioInterruptStatusData<'a, S>;
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> VirtioInterruptStatusData<'s, S> {
    pub fn read_status(&self) -> (bool, bool) {
        todo!()
    }
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> VirtioCapabilityData<'s, S>
    for VirtioInterruptStatusData<'s, S>
{
    fn from_space_offset(space: &'s S, offset: usize) -> Self {
        Self { space, offset }
    }

    fn space(&self) -> &'s S {
        self.space
    }

    fn offset(&self) -> usize {
        self.offset
    }
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> PowerManagementData<'s, S> {
    pub fn set_device_power_state(&self, state: DevicePowerState) {
        let pmcsr = self.space.read_u16(self.offset + 4) & !0x3;
        let current = self.get_device_power_state();

        if state == current {
            return;
        }

        log::info!("Set device power state: {state:?}");

        match state {
            DevicePowerState::D0 => {
                // power = 0b00 | PME_EN
                self.space.write_u16(self.offset + 4, pmcsr);
            }
            _ => {
                log::warn!("TODO: {state:?} power state");
            }
        }
    }

    pub fn set_pme_en(&self, state: bool) {
        let pmcsr = self.space.read_u16(self.offset + 4);
        let new = if state {
            pmcsr | (1 << 8)
        } else {
            pmcsr & !(1 << 8)
        };
        if pmcsr == new {
            return;
        }

        log::info!("Set PMCSR.PME_En = {state}");

        self.space.write_u16(self.offset + 4, new);
    }

    pub fn get_device_power_state(&self) -> DevicePowerState {
        let pmcsr = self.space.read_u16(self.offset + 4);
        match pmcsr & 0x3 {
            0b00 => DevicePowerState::D0,
            0b01 => DevicePowerState::D1,
            0b10 => DevicePowerState::D2,
            0b11 => DevicePowerState::D3Hot,
            _ => unreachable!(),
        }
    }
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> MsiXData<'s, S> {
    // TODO use pending bits as well
    /// Maps and returns the vector table associated with the device's MSI-X capability
    pub fn vector_table<'a>(&self) -> Result<MsiXVectorTable<'a>, Error> {
        let w0 = self.space.read_u16(self.offset + 2);
        let dw1 = self.space.read_u32(self.offset + 4);

        let table_size = (w0 as usize & 0x3FF) + 1;
        let bir = dw1 as usize & 0x3;
        let table_offset = dw1 as usize & !0x3;

        let Some(base) = self.space.bar(bir) else {
            return Err(Error::DoesNotExist);
        };

        match base {
            PciBaseAddress::Memory32(mem32) => unsafe {
                log::info!("MSI-X table address: {:#x}", mem32 + table_offset as u32);
                MsiXVectorTable::memory_from_raw_parts(
                    PhysicalAddress::from_u32(mem32).add(table_offset),
                    table_size,
                )
            },
            PciBaseAddress::Memory64(mem64) => unsafe {
                log::info!("MSI-X table address: {:#x}", mem64 + table_offset as u64);
                MsiXVectorTable::memory_from_raw_parts(
                    PhysicalAddress::from_u64(mem64).add(table_offset),
                    table_size,
                )
            },
            #[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
            PciBaseAddress::Io(io) => unsafe {
                log::info!("MSI-X table I/O: {:#x}", io + table_offset as u16);
                MsiXVectorTable::io_from_raw_parts(io + table_offset as u16, table_size)
            },
            #[cfg(any(not(any(target_arch = "x86", target_arch = "x86_64")), rust_analyzer))]
            PciBaseAddress::Io(_) => Err(Error::DoesNotExist),
        }
    }

    /// Changes the global enable status for the device's MSI-X capability. If set, regular IRQs
    /// are not generated.
    pub fn set_enabled(&mut self, enabled: bool) {
        let mut w0 = self.space.read_u32(self.offset);
        if enabled {
            w0 |= 1 << 31;
        } else {
            w0 &= !(1 << 31);
        }
        self.space.write_u32(self.offset, w0);
    }

    pub fn set_function_mask(&mut self, masked: bool) {
        let mut w0 = self.space.read_u32(self.offset);
        if masked {
            w0 |= 1 << 30;
        } else {
            w0 &= !(1 << 30);
        }
        self.space.write_u32(self.offset, w0);
    }
}

impl MsiXVectorTableAccess<'_> {
    fn set_vector_masked(&mut self, vector: usize, masked: bool) {
        let old = self.read_control(vector);
        let new = if masked { old | 1 } else { old & !1 };
        if old != new {
            self.write_control(vector, new);
        }
    }

    fn read_control(&mut self, vector: usize) -> u32 {
        match self {
            #[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
            &mut Self::Io(base) => unsafe {
                let a = base
                    + (vector * size_of::<MsiXEntry>() + offset_of!(MsiXEntry, control)) as u16;
                intrinsics::inl(a)
            },

            Self::Memory(vectors) => vectors[vector].control.get(),
        }
    }

    fn write_address(&mut self, vector: usize, value: u64) {
        match self {
            #[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
            &mut Self::Io(base) => unsafe {
                let a = base + (vector * size_of::<MsiXEntry>()) as u16;
                intrinsics::outl(a, value as u32);
                intrinsics::outl(a + 4, (value >> 32) as u32);
            },

            Self::Memory(vectors) => vectors[vector].address.set(value),
        }
    }

    fn write_data(&mut self, vector: usize, value: u32) {
        match self {
            #[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
            &mut Self::Io(base) => unsafe {
                let a =
                    base + (vector * size_of::<MsiXEntry>() + offset_of!(MsiXEntry, data)) as u16;
                intrinsics::outl(a, value)
            },

            Self::Memory(vectors) => vectors[vector].data.set(value),
        }
    }

    fn write_control(&mut self, vector: usize, value: u32) {
        match self {
            #[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
            &mut Self::Io(base) => unsafe {
                let a = base
                    + (vector * size_of::<MsiXEntry>() + offset_of!(MsiXEntry, control)) as u16;
                intrinsics::outl(a, value)
            },

            Self::Memory(vectors) => vectors[vector].control.set(value),
        }
    }
}

impl MsiXVectorTable<'_> {
    unsafe fn memory_from_raw_parts(base: PhysicalAddress, len: usize) -> Result<Self, Error> {
        let vectors = DeviceMemoryIoMut::map_slice(base, len, Default::default())?;
        Ok(Self {
            access: MsiXVectorTableAccess::Memory(vectors),
            len,
        })
    }

    #[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
    unsafe fn io_from_raw_parts(base: u16, len: usize) -> Result<Self, Error> {
        Ok(Self {
            access: MsiXVectorTableAccess::Io(base),
            len,
        })
    }

    pub fn mask_all(&mut self) {
        for i in 0..self.len {
            self.access.set_vector_masked(i, true);
        }
    }

    pub fn register_range(
        &mut self,
        start: usize,
        end: usize,
        ic: &Arc<dyn MessageInterruptController>,
        affinity: InterruptAffinity,
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<Vec<MsiInfo>, Error> {
        assert!(end > start);
        let mut range = vec![
            MsiInfo {
                affinity,
                ..Default::default()
            };
            end - start
        ];
        ic.clone().register_msi_range(&mut range, handler)?;

        for (i, info) in range.iter().enumerate() {
            let index = i + start;
            self.access.write_address(index, info.address as _);
            self.access.write_data(index, info.value);
            self.access.set_vector_masked(index, false);
        }

        Ok(range)
    }
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> MsiData<'s, S> {
    pub fn set_enabled(&mut self, enabled: bool) {
        let mut w0 = self.space.read_u16(self.offset + 2);
        if enabled {
            w0 |= 1 << 0;
        } else {
            w0 &= !(1 << 0);
        }
        self.space.write_u16(self.offset + 2, w0);
    }

    pub fn register(
        &mut self,
        ic: &Arc<dyn MessageInterruptController>,
        affinity: InterruptAffinity,
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<MsiInfo, Error> {
        let info = ic.clone().register_msi(affinity, handler)?;

        let mut w0 = self.space.read_u16(self.offset + 2);
        // Enable the vector first
        w0 |= 1 << 0;

        // Reset to one vector
        w0 &= !(0x7 << 4);

        self.space.write_u16(self.offset + 2, w0);

        if info.value > u16::MAX as u32 {
            log::warn!("Could not setup a MSI: value={:#x} > u16", info.value);
            return Err(Error::InvalidOperation);
        }

        if info.address > u32::MAX as usize {
            if w0 & (1 << 7) == 0 {
                log::warn!(
                    "Could not setup a MSI: address={:#x} and MSI is not 64 bit capable",
                    info.address
                );
                return Err(Error::InvalidOperation);
            }

            todo!("FIXME: PCI 64-bit addresses");
            // self.space
            //     .write_u32(self.offset + 8, (info.address >> 32) as u32);
        }
        self.space.write_u32(self.offset + 4, info.address as u32);

        self.space.write_u16(self.offset + 12, info.value as u16);

        Ok(info)
    }
}

impl<'s, S: PciConfigurationSpace + ?Sized + 's> PcieData<'s, S> {
    pub fn link_control(&self) -> PcieLinkControl {
        PcieLinkControl::from_bits_retain(self.space.read_u16(self.offset + 0x10))
    }

    pub fn set_link_control(&mut self, value: PcieLinkControl) {
        self.space.write_u16(self.offset + 0x10, value.bits());
    }
}
