use crate::cap::cnode;
use crate::cap::pool::POOL;
use crate::cap::retype::kernel_retype;
use crate::cap::table::Rights;
use crate::mem::phys::BitmapFrameAllocator;
use crate::proc::PROCESSES;
use crate::proc::address_space::{self, PageAccess};
use crate::tests::helpers;
use lancer_core::header::NONE_SENTINEL;
use lancer_core::object_layout::VRegionObject;
use lancer_core::object_tag::ObjectTag;
use x86_64::PhysAddr;
use x86_64::structures::paging::PhysFrame;

crate::kernel_test!(
    fn vregion_retype_from_untyped() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc");
        ptable.start(created).expect("start");
        let pid = created.pid();
        helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, phys_base) = helpers::allocate_untyped(&ptable, false);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");
        let dest_slot = 50u64;
        let page_count = 4u16;

        {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::VRegion,
                0,
                cnode_id,
                cnode_gen,
                dest_slot,
                depth,
                gv,
                gb,
                page_count as u32,
            )
            .expect("retype vregion");
        }

        {
            let pool = POOL.lock_after(&ptable);
            let cap =
                cnode::resolve_and_read(&pool, cnode_id, cnode_gen, dest_slot, depth, gv, gb)
                    .expect("read slot");
            assert!(cap.tag() == ObjectTag::VRegion);
            assert!(cap.rights().contains(Rights::ALL));

            let vr = pool
                .read_as::<VRegionObject>(cap.phys(), cap.generation())
                .expect("read vregion");
            assert!(vr.phys_base >= phys_base.as_u64());
            assert!(vr.page_count == page_count);
            assert!(vr.owner_pid == NONE_SENTINEL);
            assert!(vr.child_pid == NONE_SENTINEL);
            assert!(vr.header.parent_untyped != NONE_SENTINEL);
        }

        ptable.destroy(pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn vregion_retype_marks_bitmap() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc");
        ptable.start(created).expect("start");
        let pid = created.pid();
        helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = helpers::allocate_untyped(&ptable, false);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");
        let page_count = 8u16;

        {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::VRegion,
                0,
                cnode_id,
                cnode_gen,
                55,
                depth,
                gv,
                gb,
                page_count as u32,
            )
            .expect("retype vregion");
        }

        {
            let pool = POOL.lock_after(&ptable);
            let cap = cnode::resolve_and_read(&pool, cnode_id, cnode_gen, 55, depth, gv, gb)
                .expect("read slot");
            let vr = pool
                .read_as::<VRegionObject>(cap.phys(), cap.generation())
                .expect("read vregion");

            (0..page_count as u64).for_each(|i| {
                let frame_idx = ((vr.phys_base + i * 4096) / 4096) as usize;
                assert!(
                    BitmapFrameAllocator::is_used(frame_idx),
                    "backing frame {} must be marked used",
                    i
                );
            });
        }

        ptable.destroy(pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn vregion_map_and_access() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc");
        ptable.start(created).expect("start");
        let pid = created.pid();
        helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = helpers::allocate_untyped(&ptable, false);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");
        let dest_slot = 60u64;
        let page_count = 4u16;
        let vaddr_base = 0x20_0000u64;

        {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::VRegion,
                0,
                cnode_id,
                cnode_gen,
                dest_slot,
                depth,
                gv,
                gb,
                page_count as u32,
            )
            .expect("retype vregion");
        }

        let (vr_phys, vr_gen, backing_phys) = {
            let pool = POOL.lock_after(&ptable);
            let cap =
                cnode::resolve_and_read(&pool, cnode_id, cnode_gen, dest_slot, depth, gv, gb)
                    .expect("read slot");
            let vr = pool
                .read_as::<VRegionObject>(cap.phys(), cap.generation())
                .expect("read vregion");
            (cap.phys(), cap.generation(), vr.phys_base)
        };

        let pml4_phys = ptable.exec(pid).map(|e| e.pml4_phys).expect("pml4");
        let mut frame_alloc = BitmapFrameAllocator;

        (0..page_count as u64).for_each(|i| {
            let page_phys = PhysAddr::new(backing_phys + i * 4096);
            crate::mem::addr::zero_frame(page_phys);
            let phys_frame = PhysFrame::containing_address(page_phys);
            let page_vaddr = x86_64::VirtAddr::new(vaddr_base + i * 4096);
            address_space::map_user_page_inner(
                pml4_phys,
                page_vaddr,
                phys_frame,
                PageAccess::ReadWrite,
                &mut frame_alloc,
            )
            .unwrap_or_else(|e| panic!("map page {} failed: {:?}", i, e));
        });

        {
            let mut pool = POOL.lock_after(&ptable);
            let vr = pool
                .write_as::<VRegionObject>(vr_phys, vr_gen)
                .expect("write vregion");
            vr.owner_pid = pid.raw();
            vr.owner_vaddr = vaddr_base;
            vr.flags = lancer_core::types::VRegionFlags::try_new(lancer_core::types::VRegionFlags::WRITABLE).unwrap();
        }

        (0..page_count as u64).for_each(|i| {
            let resolved = crate::syscall::resolve_user_page(
                pml4_phys.raw(),
                vaddr_base + i * 4096,
                true,
                false,
            );
            assert!(
                resolved.is_some(),
                "page {} must be resolvable in PML4",
                i
            );
            let resolved_phys = resolved.unwrap();
            assert!(
                resolved_phys == backing_phys + i * 4096,
                "page {} resolved phys {:#x} != expected {:#x}",
                i,
                resolved_phys,
                backing_phys + i * 4096,
            );
        });

        (0..page_count as u64).for_each(|i| {
            let page_phys = PhysAddr::new(backing_phys + i * 4096);
            let page_virt = crate::mem::addr::phys_to_virt(page_phys);
            let ptr = page_virt.as_mut_ptr::<u64>();
            unsafe {
                core::ptr::write_volatile(ptr, 0xDEAD_BEEF_0000 + i);
                let val = core::ptr::read_volatile(ptr);
                assert!(
                    val == 0xDEAD_BEEF_0000 + i,
                    "page {} readback mismatch: got {:#x}",
                    i,
                    val
                );
            }
        });

        ptable.destroy(pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn vregion_backing_pages_zeroed() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc");
        ptable.start(created).expect("start");
        let pid = created.pid();
        helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = helpers::allocate_untyped(&ptable, false);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");

        {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::VRegion,
                0,
                cnode_id,
                cnode_gen,
                70,
                depth,
                gv,
                gb,
                4,
            )
            .expect("retype vregion");
        }

        let backing_phys = {
            let pool = POOL.lock_after(&ptable);
            let cap = cnode::resolve_and_read(&pool, cnode_id, cnode_gen, 70, depth, gv, gb)
                .expect("read slot");
            pool.read_as::<VRegionObject>(cap.phys(), cap.generation())
                .expect("read vregion")
                .phys_base
        };

        (0..4u64).for_each(|i| {
            let page_phys = PhysAddr::new(backing_phys + i * 4096);
            crate::mem::addr::zero_frame(page_phys);

            let page_virt = crate::mem::addr::phys_to_virt(page_phys);
            let slice =
                unsafe { core::slice::from_raw_parts(page_virt.as_ptr::<u8>(), 4096) };
            assert!(
                slice.iter().all(|&b| b == 0),
                "backing page {} not zeroed after zero_frame",
                i
            );
        });

        ptable.destroy(pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn vregion_destroy_unmaps_owner() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc");
        ptable.start(created).expect("start");
        let pid = created.pid();
        helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = helpers::allocate_untyped(&ptable, false);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");
        let dest_slot = 80u64;
        let page_count = 2u16;
        let vaddr_base = 0x30_0000u64;

        {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::VRegion,
                0,
                cnode_id,
                cnode_gen,
                dest_slot,
                depth,
                gv,
                gb,
                page_count as u32,
            )
            .expect("retype vregion");
        }

        let (vr_obj_phys, vr_gen, backing_phys) = {
            let pool = POOL.lock_after(&ptable);
            let cap =
                cnode::resolve_and_read(&pool, cnode_id, cnode_gen, dest_slot, depth, gv, gb)
                    .expect("read slot");
            let vr = pool
                .read_as::<VRegionObject>(cap.phys(), cap.generation())
                .expect("read vregion");
            (cap.phys(), cap.generation(), vr.phys_base)
        };

        let pml4_phys = ptable.exec(pid).map(|e| e.pml4_phys).expect("pml4");
        let mut frame_alloc = BitmapFrameAllocator;

        (0..page_count as u64).for_each(|i| {
            let page_phys = PhysAddr::new(backing_phys + i * 4096);
            crate::mem::addr::zero_frame(page_phys);
            let phys_frame = PhysFrame::containing_address(page_phys);
            let page_vaddr = x86_64::VirtAddr::new(vaddr_base + i * 4096);
            address_space::map_user_page_inner(
                pml4_phys,
                page_vaddr,
                phys_frame,
                PageAccess::ReadWrite,
                &mut frame_alloc,
            )
            .expect("map page");
        });

        {
            let mut pool = POOL.lock_after(&ptable);
            let vr = pool
                .write_as::<VRegionObject>(vr_obj_phys, vr_gen)
                .expect("write vregion");
            vr.owner_pid = pid.raw();
            vr.owner_vaddr = vaddr_base;
            vr.flags = lancer_core::types::VRegionFlags::try_new(lancer_core::types::VRegionFlags::WRITABLE).unwrap();
        }

        (0..page_count as u64).for_each(|i| {
            assert!(
                crate::syscall::resolve_user_page(
                    pml4_phys.raw(),
                    vaddr_base + i * 4096,
                    false,
                    false,
                )
                .is_some(),
                "page {} must be mapped before destroy",
                i
            );
        });

        {
            let (_new_gen, freed) = POOL
                .lock_after(&ptable)
                .revoke_phys(vr_obj_phys, vr_gen)
                .expect("revoke");
            freed.inspect(|(phys, tag)| {
                crate::cap::ops::cleanup_by_tag_with_ptable(*tag, *phys, &mut ptable);
            });
        }

        (0..page_count as u64).for_each(|i| {
            assert!(
                crate::syscall::resolve_user_page(
                    pml4_phys.raw(),
                    vaddr_base + i * 4096,
                    false,
                    false,
                )
                .is_none(),
                "page {} must be unmapped after destroy",
                i
            );
        });

        ptable.destroy(pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn vregion_device_untyped_rejected() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc");
        ptable.start(created).expect("start");
        let pid = created.pid();
        helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = helpers::allocate_untyped(&ptable, true);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");

        let result = {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::VRegion,
                0,
                cnode_id,
                cnode_gen,
                90,
                depth,
                gv,
                gb,
                4,
            )
        };

        assert!(
            result.is_err(),
            "device untyped must reject VRegion retype"
        );

        ptable.destroy(pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn vregion_zero_page_count_noop() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc");
        ptable.start(created).expect("start");
        let pid = created.pid();
        helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = helpers::allocate_untyped(&ptable, false);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");

        let result = {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::VRegion,
                0,
                cnode_id,
                cnode_gen,
                91,
                depth,
                gv,
                gb,
                0,
            )
        };

        assert!(
            result.is_ok(),
            "page_count=0 must return Ok(()) (no-op)"
        );

        ptable.destroy(pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn vregion_contiguous_backing() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc");
        ptable.start(created).expect("start");
        let pid = created.pid();
        helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = helpers::allocate_untyped(&ptable, false);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");
        let page_count = 8u16;

        {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::VRegion,
                0,
                cnode_id,
                cnode_gen,
                95,
                depth,
                gv,
                gb,
                page_count as u32,
            )
            .expect("retype vregion");
        }

        let backing_phys = {
            let pool = POOL.lock_after(&ptable);
            let cap = cnode::resolve_and_read(&pool, cnode_id, cnode_gen, 95, depth, gv, gb)
                .expect("read slot");
            pool.read_as::<VRegionObject>(cap.phys(), cap.generation())
                .expect("read vregion")
                .phys_base
        };

        assert!(
            backing_phys % 4096 == 0,
            "backing phys {:#x} must be page-aligned",
            backing_phys
        );

        let pml4_phys = ptable.exec(pid).map(|e| e.pml4_phys).expect("pml4");
        let vaddr_base = 0x40_0000u64;
        let mut frame_alloc = BitmapFrameAllocator;

        (0..page_count as u64).for_each(|i| {
            let page_phys = PhysAddr::new(backing_phys + i * 4096);
            crate::mem::addr::zero_frame(page_phys);
            let phys_frame = PhysFrame::containing_address(page_phys);
            address_space::map_user_page_inner(
                pml4_phys,
                x86_64::VirtAddr::new(vaddr_base + i * 4096),
                phys_frame,
                PageAccess::ReadWrite,
                &mut frame_alloc,
            )
            .expect("map page");
        });

        (0..page_count as u64).for_each(|i| {
            let resolved = crate::syscall::resolve_user_page(
                pml4_phys.raw(),
                vaddr_base + i * 4096,
                false,
                false,
            )
            .unwrap_or_else(|| panic!("page {} not mapped", i));
            assert!(
                resolved == backing_phys + i * 4096,
                "page {} should map to contiguous phys: {:#x} vs expected {:#x}",
                i,
                resolved,
                backing_phys + i * 4096,
            );
        });

        ptable.destroy(pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn vregion_exceeds_max_rejected() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();
        let created = ptable.allocate(&mut allocator).expect("alloc");
        ptable.start(created).expect("start");
        let pid = created.pid();
        helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = helpers::allocate_small_untyped(&ptable, 22);
        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");

        let result = {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::VRegion,
                0,
                cnode_id,
                cnode_gen,
                100,
                depth,
                gv,
                gb,
                513,
            )
        };
        assert!(result.is_err(), "page_count=513 exceeds MAX_VREGION_PAGES");

        ptable.destroy(pid, &mut allocator);
    }
);