use crate::println;

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

#[derive(Debug, Default, Clone, Copy)]
pub struct VirtioPciRegion {
    pub bar_address: usize,
    pub offset: u32,
    pub length: u32,
}

// Helpers ECAM
fn pci_read<T>(bus: u8, dev: u8, func: u8, offset: u16) -> T {
    let addr = PCI_ECAM_BASE
        | ((bus as usize) << 20)
        | ((dev as usize) << 15)
        | ((func as usize) << 12)
        | (offset as usize);
    unsafe { core::ptr::read_volatile(addr as *const T) }
}

fn pci_write<T>(bus: u8, dev: u8, func: u8, offset: u16, val: T) {
    let addr = PCI_ECAM_BASE
        | ((bus as usize) << 20)
        | ((dev as usize) << 15)
        | ((func as usize) << 12)
        | (offset as usize);
    unsafe { core::ptr::write_volatile(addr as *mut T, val) }
}
#[derive(Copy, Clone, Debug)]
pub struct VirtioPciCaps {
    pub common_cfg: usize,
    pub notify_cfg: usize,
    pub isr_cfg: usize,
    pub notify_multiplier: u32,
}

#[derive(Debug, Clone, Copy)]
pub struct PciDevice {
    pub bus: u8,
    pub device: u8,
}

pub struct PciGeneralDevice(PciDevice);

impl PciDevice {
    pub fn new(bus: u8, device: u8) -> Self {
        Self { bus, device }
    }

    pub fn vendor_and_device_id(&self) -> (u16, u16) {
        let id = self.read::<u32>(0);
        ((id & 0xFFFF) as u16, (id >> 16) as u16)
    }
    pub fn device_id(&self) -> u16 {
        self.vendor_and_device_id().1
    }
    pub fn vendor_id(&self) -> u16 {
        self.vendor_and_device_id().0
    }

    pub fn interrupt_pin(&self) -> u8 {
        self.read(0x3D)
    }

    fn read<T>(&self, offset: u16) -> T {
        pci_read(self.bus, self.device, 0, offset)
    }

    fn write_u32(&self, offset: u16, value: u32) {
        pci_write(self.bus, self.device, 0, offset, value)
    }

    /// # Safety
    /// `self` must be a valid general PCI device.
    pub unsafe fn to_general_device(self) -> PciGeneralDevice {
        PciGeneralDevice(self)
    }
}

pub struct PciDeviceIterator {
    pub device: u8,
}

impl PciDeviceIterator {
    pub fn new() -> Self {
        let header_type: u32 = pci_read(0, 0, 0, 0);
        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 = self.device;
        self.device += 1;
        let vendor: u16 = pci_read(0, dev, 0, 0);
        if vendor == 0xFFFF {
            self.next()
        } else {
            Some(PciDevice::new(0, dev))
        }
    }
}
pub fn scan_pci_for_virtio_keyboard2() -> Option<VirtioPciCaps> {
    let device = PciDeviceIterator::new()
        .find(|device| device.vendor_and_device_id() == (0x1af4, 0x1052))
        .unwrap();
    let device = unsafe {
        // VirtIO keyboard is a general PCI device
        device.to_general_device()
    };
    None
}

pub fn scan_pci_for_virtio_keyboard() -> Option<VirtioPciCaps> {
    // Sur RISC-V Virt, on scanne généralement le bus 0
    for dev in 0..32 {
        let vdev: u32 = pci_read(0, dev, 0, 0x00);
        let vendor = (vdev & 0xffff) as u16;
        let device = (vdev >> 16) as u16;

        // On cherche le Vendor ID VirtIO (0x1AF4)
        // et le Device ID Keyboard (0x1012 ou 0x1052 pour Modern)
        if vendor == 0x1af4 && (device >= 0x1000 && device <= 0x107f) {
            // Dans ta boucle de scan, après avoir trouvé le device
            let interrupt_pin = (pci_read::<u32>(0, dev, 0, 0x3C) >> 8) & 0xFF; // Offset 0x3D
            let irq = 32 + (dev as u32 + interrupt_pin - 1) % 4;

            let old_val = pci_read::<u32>(0, 2, 0, 0x3C);

            // On garde les 24 bits du haut (Interrupt Pin, etc.) et on change les 8 bits du bas
            let new_val = (old_val & 0xFFFFFF00) | irq;

            pci_write(0, 2, 0, 0x3C, new_val);

            println!(
                "VirtIO Keyboard sur Slot {}, PIN {}, mappé sur IRQ PLIC {}",
                dev, interrupt_pin, irq
            ); // ACTIVER l'accès mémoire et le bus mastering (PCI Command)
            let cmd = pci_read::<u32>(0, dev, 0, 0x04);
            pci_write(0, dev, 0, 0x04, cmd | 0x6);

            let mut common_cfg = 0;
            let mut notify_cfg = 0;
            let mut isr_cfg = 0;

            // On initialise le multiplicateur à 1 par défaut
            let mut notify_multiplier: u32 = 1;

            // 1. Activer le Bus Master et le Memory Space globalement pour ce périphérique avant de commencer
            let old_cmd = pci_read::<u32>(0, dev, 0, 0x04);
            pci_write(0, dev, 0, 0x04, old_cmd | 0x06);

            // 2. Boucle d'assignation des BARs (AVANT de lire les capabilities)
            let mut bar_idx = 0;
            while bar_idx < 6 {
                let bar_reg = 0x10 + (bar_idx * 4);
                let bar_val = pci_read::<u32>(0, dev, 0, bar_reg);

                // Si le BAR veut de la mémoire (bit 0 == 0) et n'est pas mappé
                if bar_val & 0x1 == 0 && (bar_val & 0xFFFF_FFF0) == 0 {
                    let new_addr =
                        0x5100_0000 + (dev as u32 * 0x100_000) + (bar_idx as u32 * 0x4000);
                    println!("ALLOCATE BAR {} at {:x}", bar_idx, new_addr);
                    pci_write(0, dev, 0, bar_reg, new_addr | (bar_val & 0xF));

                    // Gérer les BAR 64-bits (ils consomment deux slots)
                    if (bar_val >> 1) & 0x3 == 2 {
                        println!("bar {} is 64 bits", bar_idx);
                        pci_write(0, dev, 0, bar_reg + 4, 0);

                        bar_idx += 1
                    }
                    pci_write(0, dev, 0, bar_reg, new_addr);
                    let confirm = pci_read::<u32>(0, dev, 0, bar_reg);
                    println!("BAR confirmé : {:x}", confirm);
                }
                bar_idx += 1
            }
            let mut cap_ptr = (pci_read::<u32>(0, dev, 0, 0x34) & 0xFF) as u16;

            while cap_ptr != 0 {
                let header: u32 = pci_read(0, dev, 0, cap_ptr);
                let cap_id = (header & 0xFF) as u8;
                let next_ptr = ((header >> 8) & 0xFF) as u16;
                let v_type = ((header >> 24) & 0xFF) as u8;

                if cap_id == 0x09 {
                    // VirtIO Vendor Capability
                    let bar_idx = (pci_read::<u32>(0, dev, 0, cap_ptr + 4) & 0xFF) as u8;
                    let offset = pci_read::<u32>(0, dev, 0, cap_ptr + 8) as usize;

                    // VirtIO : bar_idx doit être entre 0 et 5. 0xFF signifie "ignoré".
                    if bar_idx <= 5 {
                        let bar_reg_offset = 0x10 + (bar_idx as u16 * 4);
                        let mut bar_val = pci_read::<u32>(0, dev, 0, bar_reg_offset);

                        let bar_addr = (bar_val & 0xFFFF_FFF0) as usize;
                        let final_addr = bar_addr + offset;

                        match v_type {
                            1 => {
                                common_cfg = final_addr;
                                println!("[VirtIO] CommonCfg trouvé à 0x{:x}", final_addr);
                            }
                            2 => {
                                notify_cfg = final_addr;
                                // TRÈS IMPORTANT : Lire le multiplicateur (offset 16 de la capability)
                                notify_multiplier = pci_read(0, dev, 0, cap_ptr + 16);
                                println!(
                                    "[VirtIO] NotifyCfg trouvé à 0x{:x} (mult: {})",
                                    final_addr, notify_multiplier
                                );
                            }
                            3 => {
                                isr_cfg = final_addr;
                                println!("[VirtIO] IsrCfg trouvé à 0x{:x}", final_addr);
                            }
                            _ => {}
                        }
                    }
                }
                cap_ptr = next_ptr;
            }
            if common_cfg != 0 && notify_cfg != 0 && isr_cfg != 0 {
                return Some(VirtioPciCaps {
                    common_cfg,
                    notify_cfg,
                    isr_cfg,
                    notify_multiplier,
                });
            }
        }
    }
    None
}