use crate::acpi::mcfg::McfgEntry;
use crate::mem::phys::BitmapFrameAllocator;
use crate::static_vec::StaticVec;
use x86_64::structures::paging::OffsetPageTable;

use super::config::{self, EcamAddress};

pub const MAX_PCI_DEVICES: usize = 64;

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct BlockedRange {
    pub start: u16,
    pub end: u16,
}

#[allow(dead_code)]
#[derive(Debug, Clone, Copy)]
pub struct MsixCapInfo {
    pub cap_offset: u16,
    pub table_size: u16,
    pub table_bir: u8,
    pub table_offset: u32,
    pub pba_bir: u8,
    pub pba_offset: u32,
}

#[allow(dead_code)]
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct PciDeviceInfo {
    pub bus: u8,
    pub device: u8,
    pub function: u8,
    pub vendor_id: u16,
    pub device_id: u16,
    pub class_code: u8,
    pub subclass: u8,
    pub prog_if: u8,
    pub header_type: u8,
    pub interrupt_line: u8,
    pub interrupt_pin: u8,
    pub bars: [BarInfo; 6],
    pub blocked_config_ranges: [Option<BlockedRange>; 4],
    pub msix_cap: Option<MsixCapInfo>,
}

#[allow(dead_code)]
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub enum BarInfo {
    None,
    Memory {
        phys_base: u64,
        size: u64,
        is_64bit: bool,
        prefetchable: bool,
    },
    Io {
        port_base: u16,
        size: u16,
    },
}

impl BarInfo {
    pub fn to_wire(self) -> lancer_core::pci::BarInfoWire {
        match self {
            BarInfo::None => lancer_core::pci::BarInfoWire::zeroed(),
            BarInfo::Memory {
                phys_base,
                size,
                is_64bit,
                prefetchable,
            } => {
                let flags = match is_64bit {
                    true => lancer_core::pci::BAR_FLAG_64BIT,
                    false => 0,
                } | match prefetchable {
                    true => lancer_core::pci::BAR_FLAG_PREFETCHABLE,
                    false => 0,
                };
                lancer_core::pci::BarInfoWire {
                    tag: lancer_core::pci::BAR_TAG_MEMORY,
                    flags,
                    io_port_base: 0,
                    io_size: 0,
                    _reserved: [0; 2],
                    mem_phys_base: phys_base,
                    mem_size: size,
                }
            }
            BarInfo::Io { port_base, size } => lancer_core::pci::BarInfoWire {
                tag: lancer_core::pci::BAR_TAG_IO,
                flags: 0,
                io_port_base: port_base,
                io_size: size,
                _reserved: [0; 2],
                mem_phys_base: 0,
                mem_size: 0,
            },
        }
    }
}

impl PciDeviceInfo {
    pub fn to_wire(self) -> lancer_core::pci::PciDeviceInfoWire {
        let bars = [
            self.bars[0].to_wire(),
            self.bars[1].to_wire(),
            self.bars[2].to_wire(),
            self.bars[3].to_wire(),
            self.bars[4].to_wire(),
            self.bars[5].to_wire(),
        ];
        let blocked_config_ranges = [
            self.blocked_config_ranges[0].map_or(
                lancer_core::pci::BlockedRangeWire::zeroed(),
                |r| lancer_core::pci::BlockedRangeWire {
                    start: r.start,
                    end: r.end,
                },
            ),
            self.blocked_config_ranges[1].map_or(
                lancer_core::pci::BlockedRangeWire::zeroed(),
                |r| lancer_core::pci::BlockedRangeWire {
                    start: r.start,
                    end: r.end,
                },
            ),
            self.blocked_config_ranges[2].map_or(
                lancer_core::pci::BlockedRangeWire::zeroed(),
                |r| lancer_core::pci::BlockedRangeWire {
                    start: r.start,
                    end: r.end,
                },
            ),
            self.blocked_config_ranges[3].map_or(
                lancer_core::pci::BlockedRangeWire::zeroed(),
                |r| lancer_core::pci::BlockedRangeWire {
                    start: r.start,
                    end: r.end,
                },
            ),
        ];
        lancer_core::pci::PciDeviceInfoWire {
            bus: self.bus,
            device: self.device,
            function: self.function,
            class_code: self.class_code,
            subclass: self.subclass,
            prog_if: self.prog_if,
            header_type: self.header_type,
            interrupt_line: self.interrupt_line,
            interrupt_pin: self.interrupt_pin,
            _pad0: 0,
            vendor_id: self.vendor_id,
            device_id: self.device_id,
            _pad1: [0; 2],
            bars,
            blocked_config_ranges,
        }
    }
}

pub type DeviceTable = StaticVec<PciDeviceInfo, MAX_PCI_DEVICES>;

pub static DEVICE_TABLE: spin::Mutex<DeviceTable> = spin::Mutex::new(StaticVec::new());

pub fn enumerate_segment(
    segment: &McfgEntry,
    mapper: &mut OffsetPageTable,
    allocator: &mut BitmapFrameAllocator,
    hhdm_offset: u64,
    table: &mut DeviceTable,
) {
    (segment.start_bus..=segment.end_bus).for_each(|bus| {
        (0u8..32).for_each(|dev| {
            let addr = match EcamAddress::new(segment.base_address, bus, dev, 0) {
                Some(a) => a,
                None => return,
            };
            if config::ensure_ecam_mapped(addr, mapper, allocator, hhdm_offset).is_err() {
                return;
            }

            let vendor = config::read_config_u16(addr, 0x00);
            match vendor {
                0xFFFF => {}
                _ => {
                    let header_type = (config::read_config_u32(addr, 0x0C) >> 16) as u8;
                    let multifunction = header_type & 0x80 != 0;
                    let func_limit = match multifunction {
                        true => 8u8,
                        false => 1u8,
                    };

                    (0..func_limit).for_each(|func| {
                        let faddr = match EcamAddress::new(segment.base_address, bus, dev, func) {
                            Some(a) => a,
                            None => return,
                        };
                        match func {
                            0 => {}
                            _ => {
                                if config::ensure_ecam_mapped(faddr, mapper, allocator, hhdm_offset)
                                    .is_err()
                                {
                                    return;
                                }
                            }
                        }
                        probe_function(faddr, table);
                    });
                }
            }
        });
    });
}

const PCI_CAP_MSI: u8 = 0x05;
const PCI_CAP_MSIX: u8 = 0x11;

struct CapScanResult {
    blocked: [Option<BlockedRange>; 4],
    msix: Option<MsixCapInfo>,
}

fn scan_capabilities(addr: EcamAddress) -> CapScanResult {
    let mut result = CapScanResult {
        blocked: [None; 4],
        msix: None,
    };
    let range_idx = 0usize;

    let status_cmd = config::read_config_u32(addr, 0x04);
    let status = (status_cmd >> 16) as u16;
    if status & 0x10 == 0 {
        return result;
    }

    let cap_ptr_word = config::read_config_u32(addr, 0x34);
    let first_ptr = (cap_ptr_word & 0xFC) as u16;

    (0..48u16).fold((first_ptr, range_idx), |(ptr, ri), _| match ptr {
        0 => (0, ri),
        _ if ri >= 4 => (0, ri),
        _ => {
            let header = config::read_config_u32(addr, ptr);
            let cap_id = (header & 0xFF) as u8;
            let next = ((header >> 8) & 0xFC) as u16;
            let new_ri = match cap_id {
                PCI_CAP_MSI => {
                    let msg_ctrl = config::read_config_u16(addr, ptr + 2);
                    let is_64bit = (msg_ctrl >> 7) & 1 != 0;
                    let has_pvmask = (msg_ctrl >> 8) & 1 != 0;
                    let cap_size: u16 = match (is_64bit, has_pvmask) {
                        (false, false) => 10,
                        (false, true) => 14,
                        (true, false) => 14,
                        (true, true) => 20,
                    };
                    result.blocked[ri] = Some(BlockedRange {
                        start: ptr,
                        end: ptr + cap_size,
                    });
                    ri + 1
                }
                PCI_CAP_MSIX => {
                    result.blocked[ri] = Some(BlockedRange {
                        start: ptr,
                        end: ptr + 12,
                    });

                    let msg_ctrl = config::read_config_u16(addr, ptr + 2);
                    let table_size = (msg_ctrl & 0x7FF) + 1;

                    let table_word = config::read_config_u32(addr, ptr + 4);
                    let table_bir = (table_word & 0x7) as u8;
                    let table_offset = table_word & !0x7;

                    let pba_word = config::read_config_u32(addr, ptr + 8);
                    let pba_bir = (pba_word & 0x7) as u8;
                    let pba_offset = pba_word & !0x7;

                    result.msix = Some(MsixCapInfo {
                        cap_offset: ptr,
                        table_size,
                        table_bir,
                        table_offset,
                        pba_bir,
                        pba_offset,
                    });

                    ri + 1
                }
                _ => ri,
            };

            (next, new_ri)
        }
    });

    result
}

fn probe_function(addr: EcamAddress, table: &mut DeviceTable) {
    let vendor_id = config::read_config_u16(addr, 0x00);
    match vendor_id {
        0xFFFF => {}
        _ => {
            let device_id = config::read_config_u16(addr, 0x02);
            let class_rev = config::read_config_u32(addr, 0x08);
            let class_code = (class_rev >> 24) as u8;
            let subclass = (class_rev >> 16) as u8;
            let prog_if = (class_rev >> 8) as u8;

            let header_word = config::read_config_u32(addr, 0x0C);
            let header_type = ((header_word >> 16) as u8) & 0x7F;

            let (interrupt_line, interrupt_pin) = match header_type {
                0 => {
                    let irq_word = config::read_config_u32(addr, 0x3C);
                    (irq_word as u8, (irq_word >> 8) as u8)
                }
                _ => (0, 0),
            };

            let bars = match header_type {
                0 => decode_bars(addr, 6),
                1 => {
                    let mut b = [BarInfo::None; 6];
                    let partial = decode_bars(addr, 2);
                    b[0] = partial[0];
                    b[1] = partial[1];
                    b
                }
                _ => [BarInfo::None; 6],
            };

            let cap_scan = scan_capabilities(addr);

            crate::kprintln!(
                "  PCI: {:02x}:{:02x}.{} vendor={:04x} device={:04x} class={:02x}/{:02x}",
                addr.bus,
                addr.device(),
                addr.function(),
                vendor_id,
                device_id,
                class_code,
                subclass,
            );

            let info = PciDeviceInfo {
                bus: addr.bus,
                device: addr.device(),
                function: addr.function(),
                vendor_id,
                device_id,
                class_code,
                subclass,
                prog_if,
                header_type,
                interrupt_line,
                interrupt_pin,
                bars,
                blocked_config_ranges: cap_scan.blocked,
                msix_cap: cap_scan.msix,
            };

            match table.push(info) {
                Ok(()) => {}
                Err(_) => {
                    crate::kprintln!(
                        "  PCI: device table full, dropping {:02x}:{:02x}.{}",
                        addr.bus,
                        addr.device(),
                        addr.function(),
                    );
                }
            }
        }
    }
}

fn decode_bars(addr: EcamAddress, count: usize) -> [BarInfo; 6] {
    debug_assert!(count <= 6);
    let count = count.min(6);
    let mut bars = [BarInfo::None; 6];
    let mut idx = 0usize;

    let saved_cmd = config::read_config_u16(addr, 0x04);
    config::write_config_u16(addr, 0x04, saved_cmd & !0x03u16);

    core::iter::from_fn(|| match idx < count {
        true => {
            let current = idx;
            let offset = 0x10 + (current as u16) * 4;

            let original = config::read_config_u32(addr, offset);
            match original {
                0 => {
                    idx += 1;
                    Some((current, BarInfo::None))
                }
                _ => {
                    let is_io = original & 1 != 0;
                    match is_io {
                        true => {
                            config::write_config_u32(addr, offset, 0xFFFF_FFFF);
                            let sizing = config::read_config_u32(addr, offset);
                            config::write_config_u32(addr, offset, original);

                            let io_mask = sizing & !0x3u32;
                            let size = (!io_mask).wrapping_add(1) as u16;
                            let port_base = (original & !0x3) as u16;

                            idx += 1;
                            Some((current, BarInfo::Io { port_base, size }))
                        }
                        false => {
                            let is_64bit = (original >> 1) & 0x3 == 0x2;
                            let prefetchable = (original >> 3) & 1 != 0;

                            config::write_config_u32(addr, offset, 0xFFFF_FFFF);
                            let sizing_lo = config::read_config_u32(addr, offset);
                            config::write_config_u32(addr, offset, original);

                            let (phys_base, size) = match is_64bit {
                                true if current + 1 < count => {
                                    let hi_offset = offset + 4;
                                    let original_hi = config::read_config_u32(addr, hi_offset);
                                    config::write_config_u32(addr, hi_offset, 0xFFFF_FFFF);
                                    let sizing_hi = config::read_config_u32(addr, hi_offset);
                                    config::write_config_u32(addr, hi_offset, original_hi);

                                    let base =
                                        ((original_hi as u64) << 32) | ((original & !0xF) as u64);
                                    let sizing_full =
                                        ((sizing_hi as u64) << 32) | ((sizing_lo & !0xF) as u64);
                                    let sz = (!sizing_full).wrapping_add(1);

                                    idx += 2;
                                    (base, sz)
                                }
                                _ => {
                                    let base = (original & !0xF) as u64;
                                    let mask = sizing_lo & !0xF;
                                    let sz = (!(mask as u64)).wrapping_add(1) & 0xFFFF_FFFF;

                                    idx += match is_64bit {
                                        true => 2,
                                        false => 1,
                                    };
                                    (base, sz)
                                }
                            };

                            Some((
                                current,
                                BarInfo::Memory {
                                    phys_base,
                                    size,
                                    is_64bit,
                                    prefetchable,
                                },
                            ))
                        }
                    }
                }
            }
        }
        false => None,
    })
    .for_each(|(i, bar)| {
        bars[i] = bar;
    });

    config::write_config_u16(addr, 0x04, saved_cmd);

    bars
}