use crate::cap::cnode;
use crate::cap::pool::POOL;
use crate::cap::table::{CapRef, CapSlot, Rights};
use crate::error::KernelError;
use crate::mem::phys::BitmapFrameAllocator;
use crate::proc::PROCESSES;
use lancer_core::object_layout::KernelObject;
use lancer_core::object_tag::ObjectTag;

fn alloc_cnode_object(
    cnode_data: &crate::cap::object::CNodeData,
    pool: &mut crate::cap::pool::ObjectPool,
) -> (crate::types::ObjPhys, crate::types::Generation) {
    let obj_phys = crate::cap::kernel_objects::alloc_slot().expect("alloc slot");
    let header = lancer_core::header::KernelObjectHeader::new(ObjectTag::CNode, 0, 64);
    let mut obj = lancer_core::object_layout::CNodeObject::init_default(header);
    obj.slots_phys = cnode_data.slots_phys.as_u64();
    obj.size_bits = cnode_data.size_bits;
    obj.frame_count = cnode_data.frame_count;
    crate::cap::kernel_objects::write_at(obj_phys, obj);
    pool.register_object(obj_phys, ObjectTag::CNode)
        .expect("register")
}

fn bootstrap_test_cnode(pid: crate::types::Pid, ptable: &mut crate::proc::ProcessManager) {
    let size_bits = crate::proc::ROOT_CNODE_SIZE_BITS;
    let allocator = &BitmapFrameAllocator;
    let cnode_data = cnode::create_cnode(size_bits, allocator).expect("create cnode");
    let frame_count = cnode_data.frame_count;
    let (cnode_id, cnode_gen) = alloc_cnode_object(&cnode_data, &mut POOL.lock());
    let exec = ptable.exec_mut(pid).expect("get exec");
    exec.root_cnode = Some((cnode_id, cnode_gen));
    exec.cnode_depth = 64;
    exec.root_guard_bits = 64 - size_bits;
    exec.root_guard_value = 0;
    exec.charge_frames(frame_count as u16)
        .expect("charge frames");
}

crate::kernel_test!(
    fn cnode_create_and_destroy() {
        let allocator = &BitmapFrameAllocator;
        let cnode_data = cnode::create_cnode(4, allocator).expect("create cnode");
        assert!(cnode_data.size_bits == 4);
        assert!(cnode_data.frame_count > 0);
        cnode::destroy_cnode(&cnode_data, allocator);
    }
);

crate::kernel_test!(
    fn cnode_create_invalid_size_rejected() {
        let allocator = &BitmapFrameAllocator;
        assert!(matches!(
            cnode::create_cnode(0, allocator),
            Err(KernelError::InvalidParameter)
        ));
        assert!(matches!(
            cnode::create_cnode(16, allocator),
            Err(KernelError::InvalidParameter)
        ));
    }
);

crate::kernel_test!(
    fn single_level_resolve_insert_read() {
        let allocator = &BitmapFrameAllocator;
        let cnode_data = cnode::create_cnode(4, allocator).expect("create");
        let (cid, cgen) = alloc_cnode_object(&cnode_data, &mut POOL.lock());

        let (ep_id, ep_gen) =
            crate::tests::helpers::alloc_endpoint(&mut POOL.lock()).expect("alloc ep");
        let cap = CapRef::new(ObjectTag::Endpoint, ep_id, Rights::ALL, ep_gen);

        {
            let pool = POOL.lock();
            cnode::resolve_and_insert(&pool, cid, cgen, 5, 4, 0, 0, cap).expect("insert");
        }

        {
            let pool = POOL.lock();
            let read = cnode::resolve_and_read(&pool, cid, cgen, 5, 4, 0, 0).expect("read");
            assert!(read.tag() == ObjectTag::Endpoint);
            assert!(read.phys() == ep_id);
        }

        {
            let pool = POOL.lock();
            let slot = cnode::resolve(&pool, cid, cgen, 0, 4, 0, 0).expect("empty slot");
            assert!(matches!(slot, CapSlot::Empty));
        }

        {
            let mut pool = POOL.lock();
            pool.dec_ref_phys(ep_id, ep_gen);
            pool.dec_ref_phys(cid, cgen);
        }
    }
);

crate::kernel_test!(
    fn two_level_resolve() {
        let allocator = &BitmapFrameAllocator;

        let inner_data = cnode::create_cnode(4, allocator).expect("inner cnode");
        let (inner_id, inner_gen) = alloc_cnode_object(&inner_data, &mut POOL.lock());

        let outer_data = cnode::create_cnode(4, allocator).expect("outer cnode");
        let (outer_id, outer_gen) = alloc_cnode_object(&outer_data, &mut POOL.lock());

        let inner_cap = CapRef::new(ObjectTag::CNode, inner_id, Rights::ALL, inner_gen);
        {
            let pool = POOL.lock();
            cnode::resolve_and_insert(&pool, outer_id, outer_gen, 3, 4, 0, 0, inner_cap)
                .expect("insert inner cnode at slot 3");
        }

        let (ep_id, ep_gen) =
            crate::tests::helpers::alloc_endpoint(&mut POOL.lock()).expect("alloc ep");
        let ep_cap = CapRef::new(ObjectTag::Endpoint, ep_id, Rights::ALL, ep_gen);

        let address: u64 = (3 << 4) | 7;
        {
            let pool = POOL.lock();
            cnode::resolve_and_insert(&pool, outer_id, outer_gen, address, 8, 0, 0, ep_cap)
                .expect("insert ep at two-level address");
        }

        {
            let pool = POOL.lock();
            let read = cnode::resolve_and_read(&pool, outer_id, outer_gen, address, 8, 0, 0)
                .expect("two-level read");
            assert!(read.tag() == ObjectTag::Endpoint);
            assert!(read.phys() == ep_id);
        }

        {
            let mut pool = POOL.lock();
            pool.dec_ref_phys(ep_id, ep_gen);
            let _ = cnode::resolve_and_clear(&pool, outer_id, outer_gen, 3, 4, 0, 0);
            pool.dec_ref_phys(inner_id, inner_gen);
            pool.dec_ref_phys(outer_id, outer_gen);
        }
    }
);

crate::kernel_test!(
    fn resolve_empty_slot_fails() {
        let allocator = &BitmapFrameAllocator;
        let cnode_data = cnode::create_cnode(4, allocator).expect("create");
        let (cid, cgen) = alloc_cnode_object(&cnode_data, &mut POOL.lock());

        {
            let pool = POOL.lock();
            let result = cnode::resolve_and_read(&pool, cid, cgen, 0, 4, 0, 0);
            assert!(matches!(result, Err(KernelError::SlotEmpty)));
        }

        POOL.lock().dec_ref_phys(cid, cgen);
    }
);

crate::kernel_test!(
    fn resolve_non_cnode_in_path_fails() {
        let allocator = &BitmapFrameAllocator;
        let cnode_data = cnode::create_cnode(4, allocator).expect("create");
        let (cid, cgen) = alloc_cnode_object(&cnode_data, &mut POOL.lock());

        let (ep_id, ep_gen) =
            crate::tests::helpers::alloc_endpoint(&mut POOL.lock()).expect("alloc ep");
        let ep_cap = CapRef::new(ObjectTag::Endpoint, ep_id, Rights::ALL, ep_gen);

        {
            let pool = POOL.lock();
            cnode::resolve_and_insert(&pool, cid, cgen, 2, 4, 0, 0, ep_cap).expect("insert ep");
        }

        let address: u64 = (2 << 4) | 5;
        {
            let pool = POOL.lock();
            let result = cnode::resolve_and_read(&pool, cid, cgen, address, 8, 0, 0);
            assert!(matches!(result, Err(KernelError::InvalidSlot)));
        }

        {
            let mut pool = POOL.lock();
            pool.dec_ref_phys(ep_id, ep_gen);
            pool.dec_ref_phys(cid, cgen);
        }
    }
);

crate::kernel_test!(
    fn stale_generation_rejected() {
        let allocator = &BitmapFrameAllocator;
        let cnode_data = cnode::create_cnode(4, allocator).expect("create");
        let (cid, cgen) = alloc_cnode_object(&cnode_data, &mut POOL.lock());

        let stale_gen = cgen;
        {
            let mut pool = POOL.lock();
            let _ = pool.revoke_phys(cid, cgen);
        }

        {
            let pool = POOL.lock();
            let result = cnode::resolve(&pool, cid, stale_gen, 0, 4, 0, 0);
            assert!(matches!(result, Err(KernelError::StaleGeneration)));
        }

        POOL.lock().dec_ref_phys(cid, stale_gen.wrapping_inc());
    }
);

crate::kernel_test!(
    fn invalidate_stale_across_cnode() {
        let allocator = &BitmapFrameAllocator;
        let cnode_data = cnode::create_cnode(4, allocator).expect("create");
        let (cid, cgen) = alloc_cnode_object(&cnode_data, &mut POOL.lock());

        let (ep_id, ep_gen) =
            crate::tests::helpers::alloc_endpoint(&mut POOL.lock()).expect("alloc ep");
        let cap = CapRef::new(ObjectTag::Endpoint, ep_id, Rights::ALL, ep_gen);

        {
            let pool = POOL.lock();
            cnode::resolve_and_insert(&pool, cid, cgen, 1, 4, 0, 0, cap).expect("insert");
        }
        POOL.lock().inc_ref(ep_id, ep_gen).expect("inc ref");
        {
            let pool = POOL.lock();
            cnode::resolve_and_insert(&pool, cid, cgen, 7, 4, 0, 0, cap).expect("insert 2");
        }

        {
            let pool = POOL.lock();
            cnode::invalidate_stale_in_cnode(&pool, cid, cgen, ep_id, ep_gen).expect("invalidate");
        }

        {
            let pool = POOL.lock();
            let slot1 = cnode::resolve(&pool, cid, cgen, 1, 4, 0, 0).expect("check 1");
            let slot7 = cnode::resolve(&pool, cid, cgen, 7, 4, 0, 0).expect("check 7");
            assert!(matches!(slot1, CapSlot::Empty));
            assert!(matches!(slot7, CapSlot::Empty));
        }

        {
            let mut pool = POOL.lock();
            pool.dec_ref_phys(ep_id, ep_gen);
            pool.dec_ref_phys(cid, cgen);
        }
    }
);

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

        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(pid, &ptable).expect("coords");

        (0..3u64).for_each(|i| {
            let (eid, egen) = crate::tests::helpers::alloc_endpoint(&mut POOL.lock_after(&ptable))
                .expect("alloc ep");
            let cap = CapRef::new(ObjectTag::Endpoint, eid, Rights::ALL, egen);
            let pool = POOL.lock_after(&ptable);
            cnode::resolve_and_insert(&pool, cnode_id, cnode_gen, i, depth, gv, gb, cap)
                .expect("insert");
        });

        let mut drained_count = 0u32;
        {
            let mut pool = POOL.lock_after(&ptable);
            cnode::drain_cnode_tree(&mut pool, cnode_id, cnode_gen, &mut |cap, pool| {
                pool.dec_ref_phys(cap.phys(), cap.generation());
                drained_count += 1;
            })
            .expect("drain");
        }
        assert!(
            drained_count == 3,
            "expected 3 drained caps, got {}",
            drained_count
        );

        {
            let pool = POOL.lock_after(&ptable);
            let slot0 =
                cnode::resolve(&pool, cnode_id, cnode_gen, 0, depth, gv, gb).expect("slot0");
            assert!(matches!(slot0, CapSlot::Empty));
        }

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

crate::kernel_test!(
    fn thread_shares_parent_cnode() {
        let mut allocator = BitmapFrameAllocator;
        let mut ptable = PROCESSES.lock();

        let parent_created = ptable.allocate(&mut allocator).expect("alloc parent");
        ptable.start(parent_created).expect("start parent");
        let parent_pid = parent_created.pid();
        bootstrap_test_cnode(parent_pid, &mut ptable);

        let parent_cnode = ptable
            .exec(parent_pid)
            .expect("get parent exec")
            .root_cnode();

        let (ut_id, ut_gen, _) = crate::tests::helpers::allocate_small_untyped(&ptable, 16);

        let (cnode_id_p, cnode_gen_p, depth_p, gv_p, gb_p) =
            cnode::cnode_coords(parent_pid, &ptable).expect("coords");
        let child_pid = {
            let mut pool = POOL.lock_after(&ptable);
            crate::cap::retype::retype_thread_from_untyped(
                &mut pool,
                &mut ptable,
                ut_id,
                ut_gen,
                parent_pid,
                cnode_id_p,
                cnode_gen_p,
                210,
                depth_p,
                gv_p,
                gb_p,
            )
            .expect("retype thread")
        };

        let child_cnode = ptable.exec(child_pid).expect("get child exec").root_cnode();

        assert!(
            parent_cnode == child_cnode,
            "thread must share parent's root CNode"
        );
        assert!(child_cnode.is_some(), "thread must have a root CNode");

        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(parent_pid, &ptable).expect("coords");
        let (ep_id, ep_gen) =
            crate::tests::helpers::alloc_endpoint(&mut POOL.lock_after(&ptable)).expect("alloc ep");
        let cap = CapRef::new(ObjectTag::Endpoint, ep_id, Rights::ALL, ep_gen);
        {
            let pool = POOL.lock_after(&ptable);
            cnode::resolve_and_insert(&pool, cnode_id, cnode_gen, 5, depth, gv, gb, cap)
                .expect("insert via parent");
        }

        {
            let (child_cid, child_cgen, child_depth, child_gv, child_gb) =
                cnode::cnode_coords(child_pid, &ptable).expect("child coords");
            let pool = POOL.lock_after(&ptable);
            let read = cnode::resolve_and_read(
                &pool,
                child_cid,
                child_cgen,
                5,
                child_depth,
                child_gv,
                child_gb,
            )
            .expect("read via child");
            assert!(read.tag() == ObjectTag::Endpoint);
            assert!(read.phys() == ep_id);
        }

        ptable.destroy(child_pid, &mut allocator);
        ptable.destroy(parent_pid, &mut allocator);
    }
);