use crate::cap::cnode;
use crate::cap::object::{EndpointData, ObjectData, ObjectTag};
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;

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) = POOL
        .lock()
        .allocate(ObjectData::CNode(cnode_data))
        .expect("alloc cnode");
    let proc = ptable.get_mut(pid).expect("get proc");
    proc.root_cnode = Some((cnode_id, cnode_gen));
    proc.cnode_depth = size_bits;
    proc.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(13, 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) = POOL
            .lock()
            .allocate(ObjectData::CNode(cnode_data))
            .expect("alloc");

        let (ep_id, ep_gen) = POOL
            .lock()
            .allocate(ObjectData::Endpoint(EndpointData::new()))
            .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, cap).expect("insert");
        }

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

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

        {
            let mut pool = POOL.lock();
            pool.dec_ref(ep_id, ep_gen);
            pool.dec_ref(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) = POOL
            .lock()
            .allocate(ObjectData::CNode(inner_data))
            .expect("alloc inner");

        let outer_data = cnode::create_cnode(4, allocator).expect("outer cnode");
        let (outer_id, outer_gen) = POOL
            .lock()
            .allocate(ObjectData::CNode(outer_data))
            .expect("alloc outer");

        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, inner_cap)
                .expect("insert inner cnode at slot 3");
        }

        let (ep_id, ep_gen) = POOL
            .lock()
            .allocate(ObjectData::Endpoint(EndpointData::new()))
            .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, 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)
                .expect("two-level read");
            assert!(read.tag() == ObjectTag::Endpoint);
            assert!(read.object_id() == ep_id);
        }

        {
            let mut pool = POOL.lock();
            pool.dec_ref(ep_id, ep_gen);
            let _ = cnode::resolve_and_clear(&pool, outer_id, outer_gen, 3, 4);
            pool.dec_ref(inner_id, inner_gen);
            pool.dec_ref(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) = POOL
            .lock()
            .allocate(ObjectData::CNode(cnode_data))
            .expect("alloc");

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

        POOL.lock().dec_ref(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) = POOL
            .lock()
            .allocate(ObjectData::CNode(cnode_data))
            .expect("alloc");

        let (ep_id, ep_gen) = POOL
            .lock()
            .allocate(ObjectData::Endpoint(EndpointData::new()))
            .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, 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);
            assert!(matches!(result, Err(KernelError::InvalidSlot)));
        }

        {
            let mut pool = POOL.lock();
            pool.dec_ref(ep_id, ep_gen);
            pool.dec_ref(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) = POOL
            .lock()
            .allocate(ObjectData::CNode(cnode_data))
            .expect("alloc");

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

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

        POOL.lock().dec_ref(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) = POOL
            .lock()
            .allocate(ObjectData::CNode(cnode_data))
            .expect("alloc");

        let (ep_id, ep_gen) = POOL
            .lock()
            .allocate(ObjectData::Endpoint(EndpointData::new()))
            .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, 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, 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).expect("check 1");
            let slot7 = cnode::resolve(&pool, cid, cgen, 7, 4).expect("check 7");
            assert!(matches!(slot1, CapSlot::Empty));
            assert!(matches!(slot7, CapSlot::Empty));
        }

        {
            let mut pool = POOL.lock();
            pool.dec_ref(ep_id, ep_gen);
            pool.dec_ref(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) = cnode::cnode_coords(pid, &ptable).expect("coords");

        (0..3u64).for_each(|i| {
            let (eid, egen) = POOL
                .lock_after(&ptable)
                .allocate(ObjectData::Endpoint(EndpointData::new()))
                .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, 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(cap.object_id(), 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).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.get(parent_pid).expect("get parent").root_cnode();

        let child_created = ptable.allocate_thread(parent_pid, &mut allocator).expect("alloc thread");
        let child_pid = child_created.pid();

        let child_cnode = ptable.get(child_pid).expect("get child").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) = cnode::cnode_coords(parent_pid, &ptable).expect("coords");
        let (ep_id, ep_gen) = POOL
            .lock_after(&ptable)
            .allocate(ObjectData::Endpoint(EndpointData::new()))
            .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, cap)
                .expect("insert via parent");
        }

        {
            let (child_cid, child_cgen, child_depth) =
                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)
                .expect("read via child");
            assert!(read.tag() == ObjectTag::Endpoint);
            assert!(read.object_id() == ep_id);
        }

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