use crate::cap::pool::POOL;
use crate::mem::phys::BitmapFrameAllocator;
use crate::proc::PROCESSES;
use crate::tests::helpers::{dequeue_ours, destroy_batch_and_verify, spawn_batch_with_sched};
use crate::types::Pid;
use crate::wcet::tsc;
use lancer_core::object_layout::SchedContextObject;

crate::kernel_test!(
    fn budget_lifecycle_under_sustained_load() {
        let baseline = BitmapFrameAllocator::free_frames();
        let thread_count = 500usize;
        let budget_us = 200u64;
        let period_us = 2000u64;
        let tick_step = 1000u64;
        let total_ticks = 200u32;
        let dispatch_per_tick = 100usize;

        let batch = spawn_batch_with_sched(thread_count, budget_us, period_us, |_| 128);

        let mut cycle_counts = [0u32; 1024];

        let mut clock = 1_000_000u64;
        let mut total_exhaustions = 0u64;
        let mut total_replenishments = 0u64;
        let mut per_tick_ns = crate::static_vec::StaticVec::<u64, 256>::new();

        {
            let mut ptable = PROCESSES.lock();
            ptable.timer_seed(clock);
        }

        (0..total_ticks).for_each(|_| {
            let tick_start = tsc::read_tsc_fenced();
            let mut ptable = PROCESSES.lock();

            let mut dispatched = crate::static_vec::StaticVec::<Pid, 128>::new();
            dequeue_ours(&mut ptable, &batch, dispatch_per_tick, &mut dispatched);

            {
                let mut pool = POOL.lock_after(&ptable);
                dispatched.iter().for_each(|&pid| {
                    let (sc_id, sc_gen) = ptable[pid].sched_context().expect("has sched context");
                    let sc = pool
                        .write_as::<SchedContextObject>(sc_id, sc_gen)
                        .expect("write sc");
                    crate::sched::context::consume(sc, tick_step, clock);
                    match crate::sched::context::is_exhausted(sc) {
                        true => {
                            let replenish_at = clock + sc.period_us;
                            ptable.timer_insert(pid, replenish_at);
                            total_exhaustions += 1;
                        }
                        false => {
                            ptable.enqueue_ready(pid);
                        }
                    }
                });
            }

            clock += tick_step;

            let mut fired = crate::static_vec::StaticVec::<u32, 128>::new();
            ptable.timer_advance(clock, |pid_raw| {
                fired.push(pid_raw).expect("fired overflow");
            });

            {
                let mut pool = POOL.lock_after(&ptable);
                fired.iter().for_each(|&pid_raw| {
                    let pid = Pid::new(pid_raw);
                    let (sc_id, sc_gen) = ptable[pid].sched_context().expect("has sched context");
                    let sc = pool
                        .write_as::<SchedContextObject>(sc_id, sc_gen)
                        .expect("write sc");
                    crate::sched::context::replenish(sc, clock);
                    assert!(
                        sc.remaining_us == budget_us,
                        "replenishment didn't restore full budget: got {}",
                        sc.remaining_us,
                    );
                    ptable.enqueue_ready(pid);

                    let idx = batch
                        .pids
                        .iter()
                        .position(|&p| p == pid)
                        .expect("fired pid not in batch");
                    cycle_counts[idx] += 1;
                    total_replenishments += 1;
                });
            }

            drop(ptable);
            let tick_end = tsc::read_tsc_fenced();
            let _ = per_tick_ns.push(tsc::cycles_to_ns(tick_end.saturating_sub(tick_start)));
        });

        let in_flight = total_exhaustions.saturating_sub(total_replenishments);

        let threads_that_cycled = cycle_counts[..thread_count]
            .iter()
            .filter(|&&c| c > 0)
            .count();

        let min_cycles = cycle_counts[..thread_count]
            .iter()
            .copied()
            .fold(u32::MAX, u32::min);
        let max_cycles = cycle_counts[..thread_count]
            .iter()
            .copied()
            .fold(0u32, u32::max);
        let avg_cycles = total_replenishments as u32 / thread_count as u32;

        let max_tick_ns = per_tick_ns.iter().copied().fold(0u64, u64::max);
        let avg_tick_ns = per_tick_ns.iter().sum::<u64>() / total_ticks as u64;
        let half = total_ticks as usize / 2;
        let first_half = per_tick_ns.as_slice()[..half].iter().sum::<u64>() / half as u64;
        let second_half = per_tick_ns.as_slice()[half..].iter().sum::<u64>() / half as u64;

        crate::kprintln!(
            "[timer_stress] lifecycle: {}t, {} ticks, exhaust={}, replenish={}, in_flight={}",
            thread_count,
            total_ticks,
            total_exhaustions,
            total_replenishments,
            in_flight,
        );
        crate::kprintln!(
            "[timer_stress]   cycles: min={}, max={}, avg={}, coverage={}/{}",
            min_cycles,
            max_cycles,
            avg_cycles,
            threads_that_cycled,
            thread_count,
        );
        crate::kprintln!(
            "[timer_stress]   timing: avg={}ns, max={}ns, 1st_half={}ns, 2nd_half={}ns",
            avg_tick_ns,
            max_tick_ns,
            first_half,
            second_half,
        );

        assert!(
            threads_that_cycled == thread_count,
            "only {}/{} threads completed a full exhaust/replenish cycle",
            threads_that_cycled,
            thread_count,
        );

        assert!(
            max_cycles <= min_cycles * 4 + 1,
            "unfair scheduling: min {} cycles, max {} cycles ({}x spread)",
            min_cycles,
            max_cycles,
            max_cycles / min_cycles.max(1),
        );

        assert!(
            second_half < first_half * 3 + 1000,
            "per-tick cost growing: first half {}ns, second half {}ns",
            first_half,
            second_half,
        );

        assert!(
            max_tick_ns < 500_000,
            "worst-case tick {}ns exceeds 500us",
            max_tick_ns,
        );

        destroy_batch_and_verify(&batch, baseline);
    }
);

crate::kernel_test!(
    fn timer_fire_timing_accuracy() {
        let baseline = BitmapFrameAllocator::free_frames();
        let thread_count = 500usize;
        let batch = spawn_batch_with_sched(thread_count, 50, 500, |_| 128);

        let mut ptable = PROCESSES.lock();
        let base_us = 1_000_000u64;
        ptable.timer_seed(base_us);

        let spacing = 20u64;
        let tick_step = 10u64;

        let mut our_pids = crate::static_vec::StaticVec::<Pid, 512>::new();
        let mut foreign = crate::static_vec::StaticVec::<Pid, 64>::new();

        core::iter::from_fn(|| ptable.dequeue_highest())
            .for_each(|pid| match batch.pids.iter().any(|&p| p == pid) {
                true => our_pids.push(pid).expect("our_pids overflow"),
                false => foreign.push(pid).expect("foreign overflow"),
            });

        let mut deadlines = crate::static_vec::StaticVec::<(u32, u64), 512>::new();
        our_pids.iter().enumerate().for_each(|(i, &pid)| {
            let deadline = base_us + (i as u64 + 1) * spacing;
            ptable.timer_insert(pid, deadline);
            deadlines
                .push((pid.raw(), deadline))
                .expect("deadline vec overflow");
        });

        assert!(
            deadlines.len() == thread_count,
            "expected {} timers inserted, got {}",
            thread_count,
            deadlines.len(),
        );

        let end_us = base_us + (thread_count as u64 + 2) * spacing;
        let mut fires = crate::static_vec::StaticVec::<(u32, u64), 512>::new();

        core::iter::successors(Some(base_us + tick_step), |&t| {
            (t < end_us).then_some(t + tick_step)
        })
        .for_each(|tick| {
            ptable.timer_advance(tick, |pid_raw| {
                fires.push((pid_raw, tick)).expect("fire vec overflow");
            });
        });

        assert!(
            fires.len() == thread_count,
            "lost timers: inserted {}, fired {}",
            thread_count,
            fires.len(),
        );

        let mut max_drift = 0u64;
        fires.iter().for_each(|&(pid_raw, actual_tick)| {
            let expected = deadlines
                .iter()
                .find(|&&(id, _)| id == pid_raw)
                .map(|&(_, d)| d)
                .expect("fired pid not in deadline set");

            let delta_at_insert = expected.saturating_sub(base_us);
            let level = match delta_at_insert {
                0..64 => 0,
                64..4096 => 1,
                4096..262144 => 2,
                _ => 3,
            };
            let granularity = [1u64, 64, 4096, 262144][level];

            let drift = actual_tick.saturating_sub(expected);
            assert!(
                drift <= granularity + tick_step,
                "pid {} drift {}us exceeds level-{} bound {}us (expected {}, actual {})",
                pid_raw,
                drift,
                level,
                granularity + tick_step,
                expected,
                actual_tick,
            );
            max_drift = max_drift.max(drift);
        });

        crate::kprintln!(
            "[timer_stress] timing: {} entries, spacing={}us, tick={}us, max_drift={}us",
            thread_count,
            spacing,
            tick_step,
            max_drift,
        );

        foreign
            .iter()
            .for_each(|&pid| ptable.enqueue_ready(pid));
        batch
            .pids
            .iter()
            .for_each(|&pid| ptable.enqueue_ready(pid));
        drop(ptable);

        destroy_batch_and_verify(&batch, baseline);
    }
);