use crate::cap::cnode;
use crate::cap::pool::POOL;
use crate::cap::retype;
use crate::mem::phys::BitmapFrameAllocator;
use crate::proc::address_space;
use crate::proc::{PROCESSES, ProcessState};
use lancer_core::object_tag::ObjectTag;

fn retype_thread(
    ptable: &mut crate::sync::IrqMutexGuard<'_, crate::proc::ProcessManager, 0>,
    ut_id: crate::types::ObjPhys,
    ut_gen: crate::types::Generation,
    parent_pid: crate::types::Pid,
    dest_slot: u64,
) -> Result<crate::types::Pid, crate::error::KernelError> {
    let (cnode_id, cnode_gen, depth, gv, gb) = cnode::cnode_coords(parent_pid, ptable)?;
    let mut pool = POOL.lock_after(ptable);
    retype::retype_thread_from_untyped(
        &mut pool, ptable, ut_id, ut_gen, parent_pid, cnode_id, cnode_gen, dest_slot, depth, gv, gb,
    )
}

struct ThreadPair {
    owner_pid: crate::types::Pid,
    child_pid: crate::types::Pid,
}

fn with_thread_pair(
    dest_slot: u64,
    body: fn(&mut crate::sync::IrqMutexGuard<'_, crate::proc::ProcessManager, 0>, &ThreadPair),
) {
    let mut allocator = BitmapFrameAllocator;
    let mut ptable = PROCESSES.lock();

    let created = ptable.allocate(&mut allocator).expect("alloc owner");
    ptable.start(created).expect("start");
    let owner_pid = created.pid();
    crate::tests::helpers::bootstrap_test_cnode(owner_pid, &mut ptable);

    let (ut_id, ut_gen, _) = crate::tests::helpers::allocate_small_untyped(&ptable, 16);
    let child_pid =
        retype_thread(&mut ptable, ut_id, ut_gen, owner_pid, dest_slot).expect("retype thread");

    let pair = ThreadPair {
        owner_pid,
        child_pid,
    };
    body(&mut ptable, &pair);

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

crate::kernel_test!(
    fn thread_create_shares_pml4() {
        with_thread_pair(200, |ptable, p| {
            let parent_pml4 = ptable.exec(p.owner_pid).unwrap().pml4_phys;
            assert!(
                ptable.exec(p.child_pid).unwrap().pml4_phys == parent_pml4,
                "thread must share parent PML4"
            );

            let refcount = address_space::pml4_ref_get(parent_pml4.raw());
            assert!(refcount == 2, "PML4 refcount must be 2, got {}", refcount);
        });
    }
);

crate::kernel_test!(
    fn thread_create_independent_context() {
        with_thread_pair(201, |ptable, p| {
            assert!(
                ptable[p.child_pid].state() == ProcessState::Created,
                "thread initial state must be Created"
            );
            assert!(
                p.child_pid != p.owner_pid,
                "thread must have a different PID"
            );
        });
    }
);

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

        let created = ptable.allocate(&mut allocator).expect("alloc owner");
        ptable.start(created).expect("start");
        let owner_pid = created.pid();
        crate::tests::helpers::bootstrap_test_cnode(owner_pid, &mut ptable);

        let (ut_id, ut_gen, _) = crate::tests::helpers::allocate_small_untyped(&ptable, 16);
        let child_pid =
            retype_thread(&mut ptable, ut_id, ut_gen, owner_pid, 202).expect("retype thread");
        let pml4 = ptable.exec(owner_pid).unwrap().pml4_phys;

        assert!(
            address_space::pml4_ref_get(pml4.raw()) == 2,
            "refcount should be 2"
        );

        ptable.destroy(child_pid, &mut allocator);

        assert!(
            address_space::pml4_ref_get(pml4.raw()) == 1,
            "refcount should drop to 1 after thread destroy"
        );
        assert!(ptable.get(owner_pid).is_some(), "parent must still exist");
        assert!(
            ptable.exec(owner_pid).unwrap().pml4_phys == pml4,
            "parent PML4 must be unchanged"
        );

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

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

        let created = ptable.allocate(&mut allocator).expect("alloc owner");
        ptable.start(created).expect("start");
        let owner_pid = created.pid();
        crate::tests::helpers::bootstrap_test_cnode(owner_pid, &mut ptable);

        let (ut_id, ut_gen, _) = crate::tests::helpers::allocate_small_untyped(&ptable, 16);
        let child_pid =
            retype_thread(&mut ptable, ut_id, ut_gen, owner_pid, 203).expect("retype thread");
        let pml4 = ptable.exec(owner_pid).unwrap().pml4_phys;

        ptable.destroy(owner_pid, &mut allocator);
        assert!(
            address_space::pml4_ref_get(pml4.raw()) == 1,
            "refcount should be 1 after parent destroy"
        );

        ptable.destroy(child_pid, &mut allocator);
        assert!(
            address_space::pml4_ref_get(pml4.raw()) == 0,
            "refcount should be 0 after last thread destroy"
        );
    }
);

crate::kernel_test!(
    fn fs_base_default_zero() {
        with_thread_pair(204, |ptable, p| {
            assert!(
                ptable.exec(p.child_pid).unwrap().fs_base == 0,
                "thread fs_base must default to 0"
            );
            assert!(
                ptable.exec(p.owner_pid).unwrap().fs_base == 0,
                "parent fs_base must default to 0"
            );
        });
    }
);

crate::kernel_test!(
    fn fs_base_stored_per_process() {
        with_thread_pair(205, |ptable, p| {
            ptable.exec_mut(p.owner_pid).unwrap().fs_base = 0x7FFF_0000_1000;
            ptable.exec_mut(p.child_pid).unwrap().fs_base = 0x7FFF_0000_2000;

            assert!(
                ptable.exec(p.owner_pid).unwrap().fs_base == 0x7FFF_0000_1000,
                "parent fs_base mismatch"
            );
            assert!(
                ptable.exec(p.child_pid).unwrap().fs_base == 0x7FFF_0000_2000,
                "child fs_base mismatch"
            );
        });
    }
);

crate::kernel_test!(
    fn thread_inherits_parent_cnode() {
        with_thread_pair(206, |ptable, p| {
            assert!(
                ptable.exec(p.child_pid).unwrap().root_cnode()
                    == ptable.exec(p.owner_pid).unwrap().root_cnode(),
                "thread must inherit parent's root cnode"
            );
        });
    }
);

crate::kernel_test!(
    fn thread_retype_from_untyped_marks_from_untyped() {
        with_thread_pair(207, |ptable, p| {
            assert!(
                ptable.exec(p.child_pid).unwrap().from_untyped,
                "thread created via retype must have from_untyped = true"
            );
        });
    }
);

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

        let created = ptable.allocate(&mut allocator).expect("alloc owner");
        ptable.start(created).expect("start");
        let owner_pid = created.pid();
        crate::tests::helpers::bootstrap_test_cnode(owner_pid, &mut ptable);

        let (dev_id, dev_gen, _) = crate::tests::helpers::allocate_untyped_inner(&ptable, 16, true);
        let result = retype_thread(&mut ptable, dev_id, dev_gen, owner_pid, 208);

        assert!(
            result.is_err(),
            "thread retype from device untyped must fail"
        );

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

crate::kernel_test!(
    fn fs_base_msr_round_trip() {
        let original = crate::sched::switch::read_fs_base();

        crate::sched::switch::write_fs_base(0x7FFF_DEAD_BEEF);
        let readback = crate::sched::switch::read_fs_base();
        assert!(
            readback == 0x7FFF_DEAD_BEEF,
            "fs_base readback mismatch: expected 0x7FFF_DEAD_BEEF, got {:#x}",
            readback
        );

        crate::sched::switch::write_fs_base(0x0000_CAFE_0000);
        let readback2 = crate::sched::switch::read_fs_base();
        assert!(
            readback2 == 0x0000_CAFE_0000,
            "fs_base second readback mismatch: expected 0x0000_CAFE_0000, got {:#x}",
            readback2
        );

        crate::sched::switch::write_fs_base(original);
    }
);

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

        let created = ptable.allocate(&mut allocator).expect("alloc owner");
        ptable.start(created).expect("start");
        let owner_pid = created.pid();
        crate::tests::helpers::bootstrap_test_cnode(owner_pid, &mut ptable);

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

        let (cnode_id, cnode_gen, depth, gv, gb) =
            cnode::cnode_coords(owner_pid, &ptable).expect("coords");
        {
            let mut pool = POOL.lock_after(&ptable);
            crate::cap::retype::kernel_retype(
                &mut pool,
                Some(&mut *ptable),
                ut_id,
                ut_gen,
                ObjectTag::Endpoint,
                0,
                cnode_id,
                cnode_gen,
                250,
                depth,
                gv,
                gb,
                1,
            )
            .expect("consume untyped space");
        }

        let result = retype_thread(&mut ptable, ut_id, ut_gen, owner_pid, 209);
        assert!(
            result.is_err(),
            "thread retype from too-small untyped must fail"
        );

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