use crate::cap::pool::POOL;
use crate::mem::phys::BitmapFrameAllocator;
use crate::proc::{PROCESSES, ProcessState};
use crate::types::Priority;
use lancer_core::header::KernelObjectHeader;
use lancer_core::object_layout::{KernelObject, SchedContextObject};
use lancer_core::object_tag::ObjectTag;

crate::kernel_test!(
    fn zombify_transitions_to_zombie() {
        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();

        assert!(ptable.zombify(pid), "zombify should succeed");
        assert_eq!(
            ptable[pid].state(),
            ProcessState::Zombie,
            "state should be Zombie after zombify"
        );

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

crate::kernel_test!(
    fn reap_transitions_zombie_to_free() {
        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();

        ptable.zombify(pid);
        assert_eq!(ptable[pid].state(), ProcessState::Zombie);

        ptable.reap(pid, &mut allocator);
        assert!(ptable.get(pid).is_none(), "slot should be freed after reap");
    }
);

crate::kernel_test!(
    fn destroy_equivalent_to_zombify_then_reap() {
        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();

        assert!(
            ptable.destroy(pid, &mut allocator),
            "destroy should succeed"
        );
        assert!(
            ptable.get(pid).is_none(),
            "slot should be freed after destroy"
        );
    }
);

crate::kernel_test!(
    fn zombify_free_process_returns_false() {
        let ptable = PROCESSES.lock();
        let free_pid = (1..crate::types::MAX_PIDS as u32)
            .filter_map(crate::types::Pid::try_new)
            .find(|&pid| ptable.get(pid).is_none());
        drop(ptable);

        let Some(pid) = free_pid else { return };

        let mut ptable = PROCESSES.lock();
        assert!(
            !ptable.zombify(pid),
            "zombify on Free process should return false"
        );
    }
);

crate::kernel_test!(
    fn reap_non_zombie_is_noop() {
        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();

        ptable.reap(pid, &mut allocator);
        assert_eq!(
            ptable[pid].state(),
            ProcessState::Ready,
            "reap on non-Zombie should be a noop"
        );

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

crate::kernel_test!(
    fn double_zombify_returns_false() {
        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();

        assert!(ptable.zombify(pid), "first zombify should succeed");
        assert!(
            !ptable.zombify(pid),
            "second zombify should return false (already Zombie)"
        );

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

crate::kernel_test!(
    fn zombify_blocked_process() {
        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();

        ptable.simulate_dispatch(pid);

        let (ep_id, ep_gen) =
            crate::tests::helpers::alloc_endpoint(&mut POOL.lock()).expect("alloc ep");

        let _ = ptable[pid]
            .block_on(crate::proc::BlockedReason::Sending(ep_id, ep_gen))
            .expect("block");

        assert_eq!(ptable[pid].state(), ProcessState::Blocked);
        assert!(ptable.zombify(pid), "zombify Blocked should succeed");
        assert_eq!(ptable[pid].state(), ProcessState::Zombie);

        ptable.reap(pid, &mut allocator);
        let _ = POOL.lock().dec_ref_phys(ep_id, ep_gen);
    }
);

crate::kernel_test!(
    fn zombify_preserves_generation() {
        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();

        let gen_before = ptable[pid].generation;
        ptable.zombify(pid);
        let gen_after = ptable[pid].generation;

        assert_eq!(
            gen_before, gen_after,
            "zombify should not change generation"
        );

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

crate::kernel_test!(
    fn zombie_not_runnable() {
        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();

        assert!(
            ptable[pid].is_runnable(),
            "Ready process should be runnable"
        );

        ptable.zombify(pid);
        assert!(
            !ptable[pid].is_runnable(),
            "Zombie process should not be runnable"
        );

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

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

        let a_created = ptable.allocate(&mut allocator).expect("alloc A");
        let b_created = ptable.allocate(&mut allocator).expect("alloc B");
        ptable.start(a_created).expect("start A");
        ptable.start(b_created).expect("start B");
        let a_pid = a_created.pid();
        let b_pid = b_created.pid();
        drop(ptable);

        let header_a = KernelObjectHeader::new(ObjectTag::SchedContext, 0, 64);
        let mut sc_a = SchedContextObject::init_default(header_a);
        sc_a.budget_us = 10_000;
        sc_a.period_us = 100_000;
        sc_a.priority = Priority::new(200).raw();
        sc_a.attached_pid = lancer_core::header::NONE_SENTINEL;
        let (sc_a_id, sc_a_gen) =
            crate::tests::helpers::alloc_typed(&mut POOL.lock(), ObjectTag::SchedContext, sc_a)
                .expect("alloc sc A");

        let header_b = KernelObjectHeader::new(ObjectTag::SchedContext, 0, 64);
        let mut sc_b = SchedContextObject::init_default(header_b);
        sc_b.budget_us = 10_000;
        sc_b.period_us = 100_000;
        sc_b.priority = Priority::new(50).raw();
        sc_b.attached_pid = lancer_core::header::NONE_SENTINEL;
        let (sc_b_id, sc_b_gen) =
            crate::tests::helpers::alloc_typed(&mut POOL.lock(), ObjectTag::SchedContext, sc_b)
                .expect("alloc sc B");

        let mut ptable = PROCESSES.lock();
        ptable[a_pid].attach_sched_context(sc_a_id, sc_a_gen, Priority::new(200));
        ptable[b_pid].attach_sched_context(sc_b_id, sc_b_gen, Priority::new(50));

        ptable.zombify(a_pid);

        let best = (0..crate::types::MAX_PIDS as u32).fold(
            None::<(crate::types::Pid, Priority)>,
            |best, idx| {
                let pid = crate::types::Pid::new(idx);
                let proc = match ptable.get(pid) {
                    Some(p) => p,
                    None => return best,
                };
                if !proc.is_runnable() {
                    return best;
                }
                let prio = proc.effective_priority();
                match best {
                    Some((_, best_prio)) if prio <= best_prio => best,
                    _ => Some((pid, prio)),
                }
            },
        );

        assert!(
            best.is_none_or(|(pid, _)| pid != a_pid),
            "zombie process A should not be selected by scheduler"
        );

        ptable.destroy(a_pid, &mut allocator);
        ptable.destroy(b_pid, &mut allocator);
    }
);

crate::kernel_test!(
    fn reap_with_sched_context_detaches() {
        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();

        let header = KernelObjectHeader::new(ObjectTag::SchedContext, 0, 64);
        let mut sc_obj = SchedContextObject::init_default(header);
        sc_obj.budget_us = 10_000;
        sc_obj.period_us = 100_000;
        sc_obj.priority = Priority::new(100).raw();
        sc_obj.attached_pid = lancer_core::header::NONE_SENTINEL;
        let (sc_id, sc_gen) =
            crate::tests::helpers::alloc_typed(&mut POOL.lock(), ObjectTag::SchedContext, sc_obj)
                .expect("alloc sc");
        ptable[pid].attach_sched_context(sc_id, sc_gen, Priority::new(100));

        ptable.zombify(pid);
        ptable.reap(pid, &mut allocator);

        let pool = POOL.lock();
        let sc = pool
            .read_as::<SchedContextObject>(sc_id, sc_gen)
            .expect("sc should still exist");
        assert_eq!(
            sc.attached_pid,
            lancer_core::header::NONE_SENTINEL,
            "sched context attached_pid should be None after reap"
        );
        drop(pool);

        let _ = POOL.lock().dec_ref_phys(sc_id, sc_gen);
    }
);