use crate::cap::cnode;
use crate::cap::frame_table::FRAME_TABLE;
use crate::cap::pool::POOL;
use crate::cap::retype::kernel_retype;
use crate::cap::table::Rights;
use crate::proc::PROCESSES;
use crate::tests::helpers;
use lancer_core::header::NONE_SENTINEL;
use lancer_core::object_layout::FrameObject;
use lancer_core::object_tag::ObjectTag;

crate::kernel_test!(
    fn frame_retype_from_untyped() {
        let mut allocator = crate::mem::phys::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 mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::Frame,
                12,
                cnode_id,
                cnode_gen,
                dest_slot,
                depth,
                gv,
                gb,
                1,
            )
            .expect("retype frame");
        }

        {
            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::Frame);
            assert!(cap.rights().contains(Rights::ALL));

            let frame = pool
                .read_as::<FrameObject>(cap.phys(), cap.generation())
                .expect("read frame");
            assert!(frame.phys_addr >= phys_base.as_u64());
            assert!(frame.header.parent_untyped != NONE_SENTINEL);
            let ft = FRAME_TABLE.lock();
            assert!(ft.get(frame.frame_table_idx).mapping_count() == 0);
            assert!(ft.get(frame.frame_table_idx).iommu_mapping().is_none());
        }

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

crate::kernel_test!(
    fn frame_retype_marks_bitmap() {
        let mut allocator = crate::mem::phys::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 mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::Frame,
                12,
                cnode_id,
                cnode_gen,
                55,
                depth,
                gv,
                gb,
                1,
            )
            .expect("retype frame");
        }

        {
            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 frame = pool
                .read_as::<FrameObject>(cap.phys(), cap.generation())
                .expect("read frame");
            let frame_idx = (frame.phys_addr / 4096) as usize;
            assert!(
                crate::mem::phys::BitmapFrameAllocator::is_used(frame_idx),
                "frame must be marked used in bitmap"
            );
        }

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

crate::kernel_test!(
    fn frame_retype_multiple() {
        let mut allocator = crate::mem::phys::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_base = 60u64;
        let count = 4u32;

        {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::Frame,
                12,
                cnode_id,
                cnode_gen,
                dest_base,
                depth,
                gv,
                gb,
                count,
            )
            .expect("retype 4 frames");
        }

        {
            let pool = POOL.lock_after(&ptable);
            let addrs: [u64; 4] = core::array::from_fn(|i| {
                let cap = cnode::resolve_and_read(
                    &pool,
                    cnode_id,
                    cnode_gen,
                    dest_base + i as u64,
                    depth,
                    gv,
                    gb,
                )
                .expect("read slot");
                assert!(cap.tag() == ObjectTag::Frame);
                pool.read_as::<FrameObject>(cap.phys(), cap.generation())
                    .expect("read frame")
                    .phys_addr
            });

            (0..4).for_each(|i| {
                (i + 1..4).for_each(|j| {
                    assert!(
                        addrs[i] != addrs[j],
                        "frames {} and {} have same phys addr {:#x}",
                        i,
                        j,
                        addrs[i]
                    );
                });
            });

            (0..3).for_each(|i| {
                assert!(
                    addrs[i + 1] == addrs[i] + 4096,
                    "expected 4KB stride between frames {} and {}",
                    i,
                    i + 1
                );
            });
        }

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

crate::kernel_test!(
    fn device_untyped_allows_frame_retype() {
        let mut allocator = crate::mem::phys::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 mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::Frame,
                12,
                cnode_id,
                cnode_gen,
                100,
                depth,
                gv,
                gb,
                1,
            )
            .expect("device untyped must allow Frame retype");
        }

        {
            let pool = POOL.lock_after(&ptable);
            let cap = cnode::resolve_and_read(&pool, cnode_id, cnode_gen, 100, depth, gv, gb)
                .expect("read slot");
            assert!(cap.tag() == ObjectTag::Frame);
        }

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

crate::kernel_test!(
    fn frame_cleanup_unmaps_processes() {
        let mut allocator = crate::mem::phys::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::Frame,
                12,
                cnode_id,
                cnode_gen,
                110,
                depth,
                gv,
                gb,
                1,
            )
            .expect("retype frame");
        }

        let (frame_obj_id, frame_gen, ft_idx) = {
            let pool = POOL.lock_after(&ptable);
            let cap = cnode::resolve_and_read(&pool, cnode_id, cnode_gen, 110, depth, gv, gb)
                .expect("read slot");
            let frame = pool
                .read_as::<FrameObject>(cap.phys(), cap.generation())
                .expect("read frame for idx");
            (cap.phys(), cap.generation(), frame.frame_table_idx)
        };

        {
            let vaddr = x86_64::VirtAddr::new(0x40_0000);
            FRAME_TABLE
                .lock()
                .get_mut(ft_idx)
                .add_mapping(pid, vaddr)
                .expect("add mapping");
            assert!(FRAME_TABLE.lock().get(ft_idx).mapping_count() > 0);
        }

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

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

crate::kernel_test!(
    fn frame_from_untyped_not_freed_to_bitmap() {
        let mut allocator = crate::mem::phys::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::Frame,
                12,
                cnode_id,
                cnode_gen,
                120,
                depth,
                gv,
                gb,
                1,
            )
            .expect("retype frame");
        }

        let frame_phys = {
            let pool = POOL.lock_after(&ptable);
            let cap = cnode::resolve_and_read(&pool, cnode_id, cnode_gen, 120, depth, gv, gb)
                .expect("read slot");
            pool.read_as::<FrameObject>(cap.phys(), cap.generation())
                .expect("read frame")
                .phys_addr
        };

        let frame_idx = (frame_phys / 4096) as usize;
        assert!(
            crate::mem::phys::BitmapFrameAllocator::is_used(frame_idx),
            "frame must be used before revoke"
        );

        let (frame_obj_id, frame_gen) = {
            let pool = POOL.lock_after(&ptable);
            let cap = cnode::resolve_and_read(&pool, cnode_id, cnode_gen, 120, depth, gv, gb)
                .expect("read slot");
            (cap.phys(), cap.generation())
        };

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

        assert!(
            crate::mem::phys::BitmapFrameAllocator::is_used(frame_idx),
            "frame from untyped must still be marked used after revoke"
        );

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