use core::{
    ops::{Deref, DerefMut},
    ptr::addr_of_mut,
};

use bytes_struct::VolatilePackedStruct;

use crate::{
    println,
    virtio::{
        VirtioCapability, VirtioCapabilityType, VirtioNotificationCapability, VirtioPciCommonCfg,
    },
};

// Configuration pour RISC-V Virt
const PCI_ECAM_BASE: usize = 0x3000_0000;

#[derive(Copy, Clone, Debug)]
pub struct VirtioPciCaps {
    pub common_cfg: *mut VirtioPciCommonCfg,
    pub notify_cfg: *mut u16,
    pub isr_cfg: *mut u8,
    pub notify_multiplier: u32,
    pub irq: u32,
}

#[derive(Debug, Clone, Copy)]
#[allow(unused)]
pub struct PciDevice {
    pub bus: u8,
    pub device: u8,
    pub inner: *mut PciHeader,
}

#[repr(C, packed)]
#[derive(VolatilePackedStruct)]
pub struct PciHeader {
    pub vendor_id: u16,
    pub device_id: u16,
    pub command: u16,
    pub status: u16,
    pub revision_id: u8,
    pub programming_interface_bytes: u8,
    pub subclass: u8,
    pub class_code: u8,
    pub cache_line_size: u8,
    pub latency_timer: u8,
    pub header_type: u8,
    pub bist: u8,
}

#[derive(Debug, Clone, Copy)]
pub enum PciBar {
    Mem32(MemoryBar32),
    IO32(IOBar32),
    Mem64(MemoryBar64),
}

#[derive(Debug, Clone, Copy)]
pub struct MemoryBar32(*mut u32);

#[derive(Debug, Clone, Copy)]
pub struct IOBar32(*mut u32);

#[derive(Debug, Clone, Copy)]
pub struct MemoryBar64(*mut u64);

#[repr(C, packed)]
#[derive(VolatilePackedStruct)]
pub struct PciGeneralHeader {
    pub common_header: PciHeader,
    pub bars: [u32; 6],
    pub cardbus_cis_pointer: u32,
    pub subsystem_vendor_id: u16,
    pub subsystem_id: u16,
    pub expansion_rom_base_address: u32,
    pub capabilities_pointer: u8,
    pub _reserved1: [u8; 7],
    pub interrupt_line: u8,
    pub interrupt_pin: u8,
    pub min_grant: u8,
    pub max_lantency: u8,
}
#[repr(C, packed)]
#[derive(VolatilePackedStruct)]
pub struct PciCapability {
    pub capability_id: u8,
    pub next_ptr: u8,
}

pub struct PciGeneralDevice {
    #[allow(unused)]
    pub bus: u8,
    pub device: u8,
    pub ptr: *mut PciGeneralHeader,
}

impl Deref for PciGeneralDevice {
    type Target = PciDevice;

    fn deref(&self) -> &Self::Target {
        unsafe { &*(self as *const Self as *const Self::Target) }
    }
}

impl DerefMut for PciGeneralDevice {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *(self as *mut Self as *mut Self::Target) }
    }
}

impl Deref for PciDevice {
    type Target = *mut PciHeader;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl PciDevice {
    pub fn new(bus: u8, device: u8) -> Self {
        let addr = PCI_ECAM_BASE | ((bus as usize) << 20) | ((device as usize) << 15);
        Self {
            bus,
            device,
            inner: addr as *mut PciHeader,
        }
    }

    pub fn vendor_and_device_id(&self) -> (u16, u16) {
        (self.get_vendor_id(), self.get_device_id())
    }

    /// # Safety
    /// `self` must be a valid general PCI device.
    pub unsafe fn to_general_device(self) -> PciGeneralDevice {
        unsafe { core::mem::transmute::<PciDevice, PciGeneralDevice>(self) }
    }
}

impl PciGeneralDevice {
    pub fn get_capabilities_pointer(&self) -> *mut PciCapability {
        (self.ptr as usize + self.ptr.get_capabilities_pointer() as usize) as *mut PciCapability
    }
    pub fn capabilities(&self) -> PciCapabilitiesIterator {
        PciCapabilitiesIterator {
            base: self.ptr as *mut u8,
            current: self.get_capabilities_pointer(),
        }
    }

    pub fn get_bars(&mut self) -> PciBarIterator {
        PciBarIterator {
            current: 0,
            base_addr: unsafe { addr_of_mut!((*self.ptr).bars) as *mut u32 },
        }
    }
    pub fn get_bar(&mut self, i: u8) -> PciBar {
        unsafe { PciBar::new((&raw mut (*self.ptr).bars as *mut u32).add(i as usize)) }
    }
}

pub struct PciBarIterator {
    pub current: u8,
    pub base_addr: *mut u32,
}

impl Iterator for PciBarIterator {
    type Item = (u8, PciBar);

    fn next(&mut self) -> Option<Self::Item> {
        if self.current >= 6 {
            None
        } else {
            let addr = self.base_addr.wrapping_add(self.current as usize);
            let res = (self.current, unsafe { PciBar::new(addr) });
            self.current += 1;
            if matches!(res.1, PciBar::Mem64(_)) {
                self.current += 1;
            }
            Some(res)
        }
    }
}

impl MemoryBar32 {
    pub fn is_prefetchable(self) -> bool {
        unsafe { self.0.read_volatile() & 0b1000 != 0 }
    }
    pub fn base_addr(self) -> *const u8 {
        unsafe { (self.0.read_volatile() & !0b1111) as *const u8 }
    }
    pub fn allocate(self, addr: *const u8) {
        debug_assert!(addr as usize <= 0xFFFF_FFFF);
        unsafe { self.0.write_volatile((*self.0 & 0b11) | addr as u32) }
    }
}

impl MemoryBar64 {
    pub fn is_prefetchable(self) -> bool {
        unsafe { self.0.read_volatile() & 0b1000 != 0 }
    }
    pub fn base_addr(self) -> *const u8 {
        unsafe { (self.0.read_volatile() & !0b1111) as *const u8 }
    }
    pub fn allocate(self, addr: *const u8) {
        unsafe { self.0.write_volatile((*self.0 & 0b1111) | addr as u64) }
    }
}

impl IOBar32 {
    pub fn base_addr(self) -> *const u8 {
        unsafe { (self.0.read_volatile() & !0b11) as *const u8 }
    }
    pub fn allocate(self, addr: *const u8) {
        debug_assert!(addr as usize <= 0xFFFF_FFFF);
        unsafe { self.0.write_volatile((*self.0 & 0b11) | addr as u32) }
    }
}

impl PciBar {
    pub unsafe fn new(addr: *mut u32) -> Self {
        if unsafe { *addr & 0b1 != 0 } {
            PciBar::IO32(IOBar32(addr))
        } else if unsafe { *addr & 0b100 != 0 } {
            PciBar::Mem64(MemoryBar64(addr as *mut u64))
        } else {
            PciBar::Mem32(MemoryBar32(addr))
        }
    }
    pub fn is_memory_space(self) -> bool {
        matches!(self, PciBar::Mem32(_) | PciBar::Mem64(_))
    }
    #[allow(unused)]
    pub fn is_io_space(self) -> bool {
        matches!(self, PciBar::IO32(_))
    }
    #[allow(unused)]
    pub fn is_prefetchable(self) -> bool {
        match self {
            PciBar::Mem32(pci_bar32) => pci_bar32.is_prefetchable(),
            PciBar::Mem64(pci_bar64) => pci_bar64.is_prefetchable(),
            PciBar::IO32(_) => panic!("IO Space Bars are not prefetchable"),
        }
    }
    pub fn base_address(self) -> *const u8 {
        match self {
            PciBar::Mem32(pci_bar32) => pci_bar32.base_addr(),
            PciBar::Mem64(pci_bar64) => pci_bar64.base_addr(),
            PciBar::IO32(iobar32) => iobar32.base_addr(),
        }
    }
    pub fn allocate(self, addr: *const u8) {
        match self {
            PciBar::Mem32(pci_bar32) => pci_bar32.allocate(addr),
            PciBar::Mem64(pci_bar64) => pci_bar64.allocate(addr),
            PciBar::IO32(iobar32) => iobar32.allocate(addr),
        }
    }
}

pub struct PciDeviceIterator {
    pub device: u8,
}

impl PciDeviceIterator {
    pub fn new() -> Self {
        let header_type: u32 = unsafe { (PCI_ECAM_BASE as *mut u32).read_volatile() };
        if header_type & 0x80 != 0 {
            unimplemented!("Multiple PCI host controllers are not supported")
        }
        Self { device: 0 }
    }
}

impl Iterator for PciDeviceIterator {
    type Item = PciDevice;

    fn next(&mut self) -> Option<Self::Item> {
        if self.device == 32 {
            return None;
        }
        let dev = PciDevice::new(0, self.device);
        self.device += 1;
        if dev.get_vendor_id() == 0xFFFF {
            self.next()
        } else {
            Some(dev)
        }
    }
}
pub struct PciCapabilitiesIterator {
    base: *mut u8,
    current: *mut PciCapability,
}

impl Iterator for PciCapabilitiesIterator {
    type Item = *mut PciCapability;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current.is_null() {
            None
        } else {
            let res = self.current;
            self.current = unsafe {
                let offset = self.current.get_next_ptr();
                if offset != 0 {
                    self.base.add(offset as usize) as *mut PciCapability
                } else {
                    core::ptr::null_mut()
                }
            };
            Some(res)
        }
    }
}

pub fn scan_virtio_keyboard_and_mouse(device: PciDevice) -> Option<VirtioPciCaps> {
    let mut device = unsafe {
        // VirtIO keyboard is a general PCI device
        device.to_general_device()
    };
    let interrupt_pin = device.ptr.get_interrupt_pin();
    // Interrupt swizzling
    let irq = 32 + (device.device + interrupt_pin - 1) % 4;
    device.ptr.set_interrupt_line(irq);

    let command = device.get_command() | 0b110;
    device.set_command(command);

    for (i, bar) in device.get_bars() {
        if bar.is_memory_space() && bar.base_address().is_null() {
            bar.allocate(
                (0x5100_0000 + (device.device as u32 * 0x100_000) + (i as u32 * 0x4000))
                    as *const u8,
            )
        }
    }

    let mut common_cfg = core::ptr::null_mut();
    let mut notify_cfg = core::ptr::null_mut();
    let mut isr_cfg = core::ptr::null_mut();
    let mut notify_multiplier = 1;

    for capability_addr in device.capabilities() {
        if unsafe { (*capability_addr).capability_id == 0x9 } {
            let capability = unsafe { *(capability_addr as *mut VirtioCapability) };

            if capability.bar <= 5 {
                let bar_addr = device.get_bar(capability.bar).base_address();
                let addr = bar_addr.wrapping_add(capability.offset as usize);
                match capability.capability_type {
                    VirtioCapabilityType::Common => {
                        common_cfg = addr as *mut VirtioPciCommonCfg;
                        println!("[VirtIO] CommonCfg trouvé à 0x{:x?}", addr);
                    }
                    VirtioCapabilityType::Notify => {
                        notify_cfg = addr as *mut u16;
                        notify_multiplier = unsafe {
                            (*(capability_addr as *mut VirtioNotificationCapability))
                                .notify_offset_multiplier
                        };
                        println!("[VirtIO] NotifyCfg trouvé à 0x{:x?}", addr);
                    }
                    VirtioCapabilityType::Isr => {
                        isr_cfg = addr as *mut u8;
                        println!("[VirtIO] IsrCfg trouvé à 0x{:x?}", addr);
                    }
                    _ => {}
                }
            }
        }
    }

    if !common_cfg.is_null() && !notify_cfg.is_null() && !isr_cfg.is_null() {
        Some(VirtioPciCaps {
            common_cfg,
            notify_cfg,
            isr_cfg,
            notify_multiplier,
            irq: irq as u32,
        })
    } else {
        None
    }
}

pub fn scan_virtio_devices() -> (Option<VirtioPciCaps>, Option<VirtioPciCaps>) {
    let mut devices =
        PciDeviceIterator::new().filter(|device| device.vendor_and_device_id() == (0x1af4, 0x1052));
    (
        devices.next().and_then(scan_virtio_keyboard_and_mouse),
        devices.next().and_then(scan_virtio_keyboard_and_mouse),
    )
}