use crate::cap::cnode;
use crate::cap::object::ObjectTag;
use crate::cap::pool::POOL;
use crate::cap::table::{CapRef, Rights};
use crate::error::KernelError;
use crate::mem::phys::BitmapFrameAllocator;
use crate::pci::{BarInfo, DEVICE_TABLE};
use crate::proc::PROCESSES;
use crate::types::Pid;
use lancer_core::header::KernelObjectHeader;
use lancer_core::object_layout::{KernelObject, PciDeviceObject};
use x86_64::structures::paging::{PhysFrame, Size4KiB};
use x86_64::{PhysAddr, VirtAddr};

use crate::proc::address_space::{map_fb_page_inner, unmap_user_page};

fn find_virtio_net_idx() -> Option<usize> {
    let table = DEVICE_TABLE.lock();
    table
        .iter()
        .enumerate()
        .find(|(_, dev)| dev.vendor_id == 0x1AF4 && dev.class_code == 0x02)
        .map(|(i, _)| i)
}

fn find_first_memory_bar(idx: usize) -> Option<(usize, u64, u64)> {
    let table = DEVICE_TABLE.lock();
    let dev = table.get(idx)?;
    dev.bars
        .iter()
        .enumerate()
        .find_map(|(bar_idx, bar)| match bar {
            BarInfo::Memory {
                phys_base, size, ..
            } => Some((bar_idx, *phys_base, *size)),
            _ => None,
        })
}

fn bootstrap_test_cnode(pid: Pid, ptable: &mut crate::proc::ProcessManager) {
    crate::tests::helpers::bootstrap_test_cnode(pid, ptable);
}

fn setup_process_with_pci_cap(device_table_idx: u8) -> (Pid, u64) {
    let mut allocator = BitmapFrameAllocator;
    let mut ptable = PROCESSES.lock();
    let created = ptable.allocate(&mut allocator).expect("alloc process");
    ptable.start(created).expect("start");
    let pid = created.pid();
    bootstrap_test_cnode(pid, &mut ptable);

    let (cnode_id, cnode_gen, depth, gv, gb) =
        cnode::cnode_coords(pid, &ptable).expect("cnode coords");
    let address = 0u64;
    {
        let mut pool = POOL.lock_after(&ptable);
        let header = KernelObjectHeader::new(ObjectTag::PciDevice, 0, 64);
        let mut pci_obj = PciDeviceObject::init_default(header);
        pci_obj.device_table_idx = device_table_idx;
        let (obj_id, obj_gen) =
            crate::tests::helpers::alloc_typed(&mut pool, ObjectTag::PciDevice, pci_obj)
                .expect("alloc pci");
        let cap = CapRef::new(ObjectTag::PciDevice, obj_id, Rights::ALL, obj_gen);
        cnode::resolve_and_insert(&pool, cnode_id, cnode_gen, address, depth, gv, gb, cap)
            .expect("insert pci cap");
    }

    (pid, address)
}

crate::kernel_test!(
    fn pci_discovers_devices() {
        let table = DEVICE_TABLE.lock();
        assert!(
            !table.is_empty(),
            "PCI enumeration found 0 devices (Q35 has built-in devices)"
        );
    }
);

crate::kernel_test!(
    fn pci_finds_virtio_net() {
        let table = DEVICE_TABLE.lock();
        let found = table
            .iter()
            .any(|dev| dev.vendor_id == 0x1AF4 && dev.class_code == 0x02);
        assert!(
            found,
            "virtio-net device (vendor=1AF4 class=02) not found in PCI device table"
        );
    }
);

crate::kernel_test!(
    fn pci_virtio_net_has_io_bar() {
        let table = DEVICE_TABLE.lock();
        let dev = table
            .iter()
            .find(|d| d.vendor_id == 0x1AF4 && d.class_code == 0x02)
            .expect("virtio-net device must exist");

        let has_io = dev.bars.iter().any(|bar| matches!(bar, BarInfo::Io { .. }));
        assert!(
            has_io,
            "transitional virtio-net (device_id=0x1000) should have an IO BAR for legacy interface"
        );
    }
);

crate::kernel_test!(
    fn pci_virtio_net_has_memory_bar() {
        let idx = find_virtio_net_idx().expect("virtio-net must exist");
        let (bar_idx, phys_base, size) =
            find_first_memory_bar(idx).expect("virtio-net should have at least one Memory BAR");

        assert!(
            phys_base != 0,
            "BAR{} phys_base should be non-zero",
            bar_idx
        );
        assert!(size > 0, "BAR{} size should be non-zero", bar_idx);
        assert!(
            phys_base & 0xF == 0,
            "BAR{} phys_base should be 16-byte aligned",
            bar_idx
        );
    }
);

crate::kernel_test!(
    fn pci_device_cap_stores_table_idx() {
        let idx = find_virtio_net_idx().expect("virtio-net must exist") as u8;
        let (pid, address) = setup_process_with_pci_cap(idx);

        let ptable = PROCESSES.lock();
        let pool = POOL.lock_after(&ptable);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("cnode coords");
        let cap = cnode::resolve_and_validate(
            &pool,
            cnode_id,
            cnode_gen,
            address,
            depth,
            gv,
            gb,
            ObjectTag::PciDevice,
            Rights::READ,
        )
        .expect("validate pci cap");

        let pci_data = pool
            .read_as::<PciDeviceObject>(cap.phys(), cap.generation())
            .expect("retrieve pci data");

        assert!(
            pci_data.device_table_idx == idx,
            "stored idx {} != expected {}",
            pci_data.device_table_idx,
            idx
        );
        drop(pool);
        drop(ptable);

        let mut ptable = PROCESSES.lock();
        ptable.destroy(pid, &mut BitmapFrameAllocator);
    }
);

crate::kernel_test!(
    fn pci_bar_map_mmio_pages() {
        let idx = find_virtio_net_idx().expect("virtio-net must exist");
        let (phys_base, size) = find_first_memory_bar(idx)
            .map(|(_, base, sz)| (base, sz))
            .expect("need a Memory BAR");
        let page_count = size.div_ceil(4096) as usize;
        let (pid, _address) = setup_process_with_pci_cap(idx as u8);

        let pml4_phys = {
            let ptable = PROCESSES.lock();
            ptable.exec(pid).unwrap().pml4_phys
        };
        let mut allocator = BitmapFrameAllocator;
        let base_vaddr = 0x0000_2000_0000_0000u64;

        let mapped = (0..page_count).try_fold(0usize, |count, i| {
            let phys = PhysAddr::new(phys_base + (i as u64) * 4096);
            let virt = VirtAddr::new(base_vaddr + (i as u64) * 4096);
            let frame = PhysFrame::<Size4KiB>::containing_address(phys);
            match map_fb_page_inner(pml4_phys, virt, frame, &mut allocator) {
                Ok(()) => Ok(count + 1),
                Err(e) => Err((e, count)),
            }
        });

        assert!(
            mapped.is_ok(),
            "mapping Memory BAR pages should succeed, got {:?}",
            mapped
        );
        assert!(
            mapped.unwrap() == page_count,
            "should have mapped {} pages",
            page_count
        );

        (0..page_count).for_each(|i| {
            let virt = VirtAddr::new(base_vaddr + (i as u64) * 4096);
            let _ = unmap_user_page(pml4_phys, virt);
        });

        let mut ptable = PROCESSES.lock();
        ptable.destroy(pid, &mut BitmapFrameAllocator);
    }
);

crate::kernel_test!(
    fn pci_bar_unmap_returns_correct_frames() {
        let idx = find_virtio_net_idx().expect("virtio-net must exist");
        let (phys_base, size) = find_first_memory_bar(idx)
            .map(|(_, base, sz)| (base, sz))
            .expect("need a Memory BAR");
        let page_count = size.div_ceil(4096) as usize;
        let (pid, _address) = setup_process_with_pci_cap(idx as u8);

        let pml4_phys = {
            let ptable = PROCESSES.lock();
            ptable.exec(pid).unwrap().pml4_phys
        };
        let mut allocator = BitmapFrameAllocator;
        let base_vaddr = 0x0000_3000_0000_0000u64;

        (0..page_count).for_each(|i| {
            let phys = PhysAddr::new(phys_base + (i as u64) * 4096);
            let virt = VirtAddr::new(base_vaddr + (i as u64) * 4096);
            let frame = PhysFrame::<Size4KiB>::containing_address(phys);
            map_fb_page_inner(pml4_phys, virt, frame, &mut allocator).expect("map should succeed");
        });

        (0..page_count).for_each(|i| {
            let expected_phys = PhysAddr::new(phys_base + (i as u64) * 4096);
            let virt = VirtAddr::new(base_vaddr + (i as u64) * 4096);
            let returned_frame = unmap_user_page(pml4_phys, virt).expect("unmap should succeed");
            assert!(
                returned_frame.start_address() == expected_phys,
                "unmap page {} returned phys {:#x}, expected {:#x}",
                i,
                returned_frame.start_address().as_u64(),
                expected_phys.as_u64(),
            );
        });

        let mut ptable = PROCESSES.lock();
        ptable.destroy(pid, &mut BitmapFrameAllocator);
    }
);

crate::kernel_test!(
    fn pci_bar_unmap_nonexistent_fails() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc process");
        ptable.start(created).expect("start");
        let pid = created.pid();
        let pml4_phys = ptable.exec(pid).unwrap().pml4_phys;
        drop(ptable);

        let virt = VirtAddr::new(0x0000_5000_0000_0000);
        let result = unmap_user_page(pml4_phys, virt);
        assert!(
            result == Err(KernelError::InvalidAddress),
            "unmapping a never-mapped page should fail"
        );

        let mut ptable = PROCESSES.lock();
        ptable.destroy(pid, &mut BitmapFrameAllocator);
    }
);

crate::kernel_test!(
    fn pci_bar_map_double_map_fails() {
        let idx = find_virtio_net_idx().expect("virtio-net must exist");
        let (phys_base, _) = find_first_memory_bar(idx)
            .map(|(_, base, sz)| (base, sz))
            .expect("need a Memory BAR");
        let (pid, _address) = setup_process_with_pci_cap(idx as u8);

        let pml4_phys = {
            let ptable = PROCESSES.lock();
            ptable.exec(pid).unwrap().pml4_phys
        };
        let mut allocator = BitmapFrameAllocator;
        let vaddr = VirtAddr::new(0x0000_4000_0000_0000);
        let frame = PhysFrame::<Size4KiB>::containing_address(PhysAddr::new(phys_base));

        map_fb_page_inner(pml4_phys, vaddr, frame, &mut allocator)
            .expect("first map should succeed");

        let second = map_fb_page_inner(pml4_phys, vaddr, frame, &mut allocator);
        assert!(second.is_err(), "mapping the same vaddr twice should fail");

        let _ = unmap_user_page(pml4_phys, vaddr);

        let mut ptable = PROCESSES.lock();
        ptable.destroy(pid, &mut BitmapFrameAllocator);
    }
);

fn make_test_pci_device(
    ranges: [Option<crate::pci::device::BlockedRange>; 4],
) -> crate::pci::PciDeviceInfo {
    crate::pci::PciDeviceInfo {
        bus: 0,
        device: 0,
        function: 0,
        vendor_id: 0,
        device_id: 0,
        class_code: 0,
        subclass: 0,
        prog_if: 0,
        header_type: 0,
        interrupt_line: 0,
        interrupt_pin: 0,
        bars: [BarInfo::None; 6],
        blocked_config_ranges: ranges,
        msix_cap: None,
    }
}

crate::kernel_test!(
    fn pci_config_write_rejects_command_register() {
        let dev = make_test_pci_device([None; 4]);
        assert!(
            !crate::syscall::pci::is_config_write_safe(0x04, &dev),
            "command register (0x04) must be blocked"
        );
    }
);

crate::kernel_test!(
    fn pci_config_write_rejects_bar_offsets() {
        let dev = make_test_pci_device([None; 4]);
        [0x10u16, 0x14, 0x18, 0x1C, 0x20, 0x24]
            .iter()
            .for_each(|&offset| {
                assert!(
                    !crate::syscall::pci::is_config_write_safe(offset, &dev),
                    "BAR offset {:#x} must be blocked",
                    offset
                );
            });
    }
);

crate::kernel_test!(
    fn pci_config_write_rejects_expansion_rom() {
        let dev = make_test_pci_device([None; 4]);
        assert!(
            !crate::syscall::pci::is_config_write_safe(0x30, &dev),
            "expansion ROM (0x30) must be blocked"
        );
    }
);

crate::kernel_test!(
    fn pci_config_write_rejects_msi_range() {
        use crate::pci::device::BlockedRange;
        let dev = make_test_pci_device([
            Some(BlockedRange {
                start: 0x40,
                end: 0x4E,
            }),
            None,
            None,
            None,
        ]);
        assert!(
            !crate::syscall::pci::is_config_write_safe(0x40, &dev),
            "MSI start (0x40) must be blocked"
        );
        assert!(
            !crate::syscall::pci::is_config_write_safe(0x48, &dev),
            "MSI mid (0x48) must be blocked"
        );
        assert!(
            crate::syscall::pci::is_config_write_safe(0x38, &dev),
            "offset 0x38 before MSI range should be allowed"
        );
        assert!(
            crate::syscall::pci::is_config_write_safe(0x50, &dev),
            "offset 0x50 after MSI range should be allowed"
        );
    }
);

crate::kernel_test!(
    fn pci_config_write_allows_safe_offsets() {
        let dev = make_test_pci_device([None; 4]);
        assert!(
            crate::syscall::pci::is_config_write_safe(0x3C, &dev),
            "interrupt line (0x3C) should be allowed"
        );
    }
);

crate::kernel_test!(
    fn pci_bar_mapping_tracks_duplicates() {
        use crate::pci::{BAR_MAPPINGS, BarMappingEntry};

        let pid = Pid::new(1);
        let entry = BarMappingEntry {
            pid,
            device_idx: 5,
            bar_idx: 2,
            base_vaddr: 0x2000_0000,
        };

        let mut mappings = BAR_MAPPINGS.lock();
        let _ = mappings.push(entry);

        assert!(
            mappings.iter().any(|e| e.matches(pid, 5, 2)),
            "should find matching entry"
        );
        assert!(
            !mappings.iter().any(|e| e.matches(pid, 5, 3)),
            "different bar_idx should not match"
        );
        assert!(
            !mappings.iter().any(|e| e.matches(pid, 6, 2)),
            "different device_idx should not match"
        );
        assert!(
            !mappings.iter().any(|e| e.matches(Pid::new(2), 5, 2)),
            "different pid should not match"
        );

        let pos = mappings
            .iter()
            .enumerate()
            .find(|(_, e)| e.matches(pid, 5, 2))
            .map(|(i, _)| i);
        if let Some(i) = pos {
            mappings.swap_remove(i);
        }
    }
);

crate::kernel_test!(
    fn pci_device_info_wire_roundtrip() {
        use lancer_core::pci::*;
        use zerocopy::IntoBytes;

        let table = DEVICE_TABLE.lock();
        table.iter().for_each(|dev| {
            let wire = dev.to_wire();
            assert!(wire.bus == dev.bus, "bus mismatch");
            assert!(wire.device == dev.device, "device mismatch");
            assert!(wire.function == dev.function, "function mismatch");
            assert!(wire.vendor_id == dev.vendor_id, "vendor_id mismatch");
            assert!(wire.device_id == dev.device_id, "device_id mismatch");
            assert!(wire.class_code == dev.class_code, "class_code mismatch");
            assert!(wire.subclass == dev.subclass, "subclass mismatch");
            assert!(wire.prog_if == dev.prog_if, "prog_if mismatch");
            assert!(wire.header_type == dev.header_type, "header_type mismatch");
            assert!(
                wire.interrupt_line == dev.interrupt_line,
                "interrupt_line mismatch"
            );
            assert!(
                wire.interrupt_pin == dev.interrupt_pin,
                "interrupt_pin mismatch"
            );

            dev.bars.iter().zip(wire.bars.iter()).enumerate().for_each(
                |(i, (orig, w))| match orig {
                    BarInfo::None => {
                        assert!(w.tag == BAR_TAG_NONE, "BAR{} tag should be None", i);
                    }
                    BarInfo::Memory {
                        phys_base,
                        size,
                        is_64bit,
                        prefetchable,
                    } => {
                        assert!(w.tag == BAR_TAG_MEMORY, "BAR{} tag should be Memory", i);
                        assert!(w.mem_phys_base == *phys_base, "BAR{} phys_base mismatch", i);
                        assert!(w.mem_size == *size, "BAR{} size mismatch", i);
                        assert!(
                            (w.flags & BAR_FLAG_64BIT != 0) == *is_64bit,
                            "BAR{} is_64bit mismatch",
                            i
                        );
                        assert!(
                            (w.flags & BAR_FLAG_PREFETCHABLE != 0) == *prefetchable,
                            "BAR{} prefetchable mismatch",
                            i
                        );
                    }
                    BarInfo::Io { port_base, size } => {
                        assert!(w.tag == BAR_TAG_IO, "BAR{} tag should be Io", i);
                        assert!(w.io_port_base == *port_base, "BAR{} port_base mismatch", i);
                        assert!(w.io_size == *size, "BAR{} io_size mismatch", i);
                    }
                },
            );

            dev.blocked_config_ranges
                .iter()
                .zip(wire.blocked_config_ranges.iter())
                .for_each(|(orig, w)| match orig {
                    Some(r) => {
                        assert!(w.start == r.start, "blocked range start mismatch");
                        assert!(w.end == r.end, "blocked range end mismatch");
                    }
                    None => {
                        assert!(w.start == 0 && w.end == 0, "None range should be zeroed");
                    }
                });

            let bytes = wire.as_bytes();
            assert!(bytes.len() == 176, "wire bytes should be 176");
        });
    }
);

crate::kernel_test!(
    fn pci_all_bar_types_valid() {
        let table = DEVICE_TABLE.lock();
        table.iter().for_each(|dev| {
            dev.bars.iter().enumerate().for_each(|(i, bar)| match bar {
                BarInfo::Memory {
                    phys_base, size, ..
                } => {
                    assert!(
                        *size > 0,
                        "device {:04x}:{:04x} BAR{} Memory with size=0",
                        dev.vendor_id,
                        dev.device_id,
                        i
                    );
                    assert!(
                        *phys_base & 0xF == 0,
                        "device {:04x}:{:04x} BAR{} Memory phys_base not aligned",
                        dev.vendor_id,
                        dev.device_id,
                        i
                    );
                }
                BarInfo::Io { size, .. } => {
                    assert!(
                        *size > 0,
                        "device {:04x}:{:04x} BAR{} IO with size=0",
                        dev.vendor_id,
                        dev.device_id,
                        i
                    );
                }
                BarInfo::None => {}
            });
        });
    }
);

crate::kernel_test!(
    fn pci_virtio_net_has_msix_cap() {
        let table = DEVICE_TABLE.lock();
        let dev = table
            .iter()
            .find(|d| d.vendor_id == 0x1AF4 && d.class_code == 0x02)
            .expect("virtio-net device must exist");

        let msix = dev
            .msix_cap
            .expect("QEMU virtio-net should advertise MSI-X capability");

        assert!(
            msix.table_size > 0,
            "MSI-X table_size should be > 0, got {}",
            msix.table_size
        );
        assert!(
            msix.table_bir < 6,
            "MSI-X table BIR should be 0-5, got {}",
            msix.table_bir
        );
        assert!(
            msix.cap_offset >= 0x40,
            "MSI-X cap_offset should be >= 0x40 (past standard header), got {:#x}",
            msix.cap_offset
        );
    }
);

crate::kernel_test!(
    fn pci_msix_table_map_and_configure() {
        let idx = find_virtio_net_idx().expect("virtio-net must exist") as u8;

        let has_msix = {
            let table = DEVICE_TABLE.lock();
            table.get(idx as usize).and_then(|d| d.msix_cap).is_some()
        };
        assert!(has_msix, "virtio-net should have MSI-X");

        let hhdm_offset = crate::mem::addr::hhdm_offset();
        let mut mapper = unsafe { crate::arch::paging::init(hhdm_offset) };
        let mut alloc = BitmapFrameAllocator;

        let map_result =
            crate::pci::msix::ensure_table_mapped(idx, &mut mapper, &mut alloc, hhdm_offset);
        assert!(
            map_result.is_ok(),
            "ensure_table_mapped should succeed, got {:?}",
            map_result,
        );

        let vector = crate::arch::idt::IrqVector::new(38);
        let cfg_result = crate::pci::msix::configure_entry(idx, 0, vector, 0);
        assert!(
            cfg_result.is_ok(),
            "configure_entry should succeed, got {:?}",
            cfg_result,
        );

        let unmask_result = crate::pci::msix::unmask_entry(idx, 0);
        assert!(
            unmask_result.is_ok(),
            "unmask_entry should succeed, got {:?}",
            unmask_result,
        );
    }
);

crate::kernel_test!(
    fn pci_msix_invalid_device_idx_fails() {
        let result = crate::pci::msix::unmask_entry(255, 0);
        assert!(
            result.is_err(),
            "unmask_entry on non-existent device index should fail"
        );
    }
);