use crate::cap::cnode;
use crate::cap::object::ObjectTag;
use crate::cap::pool::POOL;
use crate::cap::retype::kernel_retype;
use crate::cap::table::Rights;
use crate::error::KernelError;
use crate::proc::PROCESSES;
use lancer_core::object_layout::{KernelObject, UntypedObject};

crate::kernel_test!(
    fn retype_single_endpoint() {
        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();
        crate::tests::helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = crate::tests::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::Endpoint,
                0,
                cnode_id,
                cnode_gen,
                dest_slot,
                depth,
                gv,
                gb,
                1,
            )
            .expect("retype endpoint");
        }

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

            assert!(pool.get_tag(cap.phys(), cap.generation()) == Ok(ObjectTag::Endpoint));
        }

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

crate::kernel_test!(
    fn retype_multiple_endpoints() {
        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();
        crate::tests::helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = crate::tests::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 = 3u32;

        {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::Endpoint,
                0,
                cnode_id,
                cnode_gen,
                dest_base,
                depth,
                gv,
                gb,
                count,
            )
            .expect("retype 3 endpoints");
        }

        {
            let pool = POOL.lock_after(&ptable);
            (0..count as u64).for_each(|i| {
                let cap = cnode::resolve_and_read(
                    &pool,
                    cnode_id,
                    cnode_gen,
                    dest_base + i,
                    depth,
                    gv,
                    gb,
                )
                .expect("read slot");
                assert!(cap.tag() == ObjectTag::Endpoint);
                assert!(cap.rights().contains(Rights::ALL));
            });
        }

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

crate::kernel_test!(
    fn retype_exceeds_capacity_fails() {
        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();
        crate::tests::helpers::bootstrap_test_cnode(pid, &mut ptable);

        let small_size_bits = 12u8;
        let small_frames = 1usize;
        let phys_base = crate::mem::phys::BitmapFrameAllocator
            .allocate_contiguous(small_frames)
            .expect("alloc small untyped");
        let base_virt = crate::mem::addr::phys_to_virt(phys_base);
        unsafe {
            core::ptr::write_bytes(base_virt.as_mut_ptr::<u8>(), 0, small_frames * 4096);
        }
        let header = lancer_core::header::KernelObjectHeader::new(ObjectTag::Untyped, 0, 64);
        let mut ut_obj = UntypedObject::init_default(header);
        ut_obj.phys_base = phys_base.as_u64();
        ut_obj.size_bits = small_size_bits;
        ut_obj.is_device = 0;
        let ut_phys = crate::cap::kernel_objects::alloc_slot().expect("alloc ut slot");
        crate::cap::kernel_objects::write_at(ut_phys, ut_obj);
        let (ut_id, ut_gen) = POOL
            .lock_after(&ptable)
            .register_object(ut_phys, ObjectTag::Untyped)
            .expect("register ut");

        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::Endpoint,
                0,
                cnode_id,
                cnode_gen,
                70,
                depth,
                gv,
                gb,
                100,
            )
        };
        assert!(result.is_err(), "retype exceeding capacity must fail");

        {
            let pool = POOL.lock_after(&ptable);
            let ut = pool
                .read_as::<UntypedObject>(ut_id, ut_gen)
                .expect("read ut");
            assert!(ut.watermark == 0, "watermark must be unchanged on failure");
        }

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

crate::kernel_test!(
    fn retype_occupied_slot_rolls_back() {
        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();
        crate::tests::helpers::bootstrap_test_cnode(pid, &mut ptable);

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

        let occupied_slot = 80u64;
        {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::Endpoint,
                0,
                cnode_id,
                cnode_gen,
                occupied_slot,
                depth,
                gv,
                gb,
                1,
            )
            .expect("first retype");
        }

        let watermark_before = {
            let pool = POOL.lock_after(&ptable);
            pool.read_as::<UntypedObject>(ut_id, ut_gen)
                .expect("read ut")
                .watermark
        };

        let result = {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                ut_gen,
                ObjectTag::Notification,
                0,
                cnode_id,
                cnode_gen,
                occupied_slot,
                depth,
                gv,
                gb,
                1,
            )
        };
        assert!(
            matches!(result, Err(KernelError::SlotOccupied)),
            "retype into occupied slot must fail"
        );

        {
            let pool = POOL.lock_after(&ptable);
            let cap =
                cnode::resolve_and_read(&pool, cnode_id, cnode_gen, occupied_slot, depth, gv, gb)
                    .expect("read slot");
            assert!(
                cap.tag() == ObjectTag::Endpoint,
                "original cap must survive"
            );

            let ut = pool
                .read_as::<UntypedObject>(ut_id, ut_gen)
                .expect("read ut");
            assert!(
                ut.watermark == watermark_before,
                "watermark must not advance on rollback"
            );
        }

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

crate::kernel_test!(
    fn retype_stale_generation_fails() {
        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();
        crate::tests::helpers::bootstrap_test_cnode(pid, &mut ptable);

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

        let stale_gen = crate::types::Generation::new(ut_gen.raw().wrapping_add(1));

        let result = {
            let mut pool = POOL.lock_after(&ptable);
            kernel_retype(
                &mut pool,
                None,
                ut_id,
                stale_gen,
                ObjectTag::Endpoint,
                0,
                cnode_id,
                cnode_gen,
                90,
                depth,
                gv,
                gb,
                1,
            )
        };
        assert!(result.is_err(), "retype with stale generation must fail");

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

crate::kernel_test!(
    fn device_untyped_rejects_non_frame() {
        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();
        crate::tests::helpers::bootstrap_test_cnode(pid, &mut ptable);

        let (ut_id, ut_gen, _phys) = crate::tests::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::Endpoint,
                0,
                cnode_id,
                cnode_gen,
                100,
                depth,
                gv,
                gb,
                1,
            )
        };
        assert!(
            matches!(result, Err(KernelError::InvalidType)),
            "device untyped must reject non-Frame retype"
        );

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