//! Bounded single-CPU task-state core. This module owns IDs, ready-queue
//! membership, sleep/exit transitions and the cooperative context layout.
//! The live EL1 switch is in task_switch.S and TaskSwitchImage; this module
//! does not implement timer preemption or user-mode tasks.

use core::cell::UnsafeCell;
use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};

pub const MAX_TASKS: usize = 4;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TaskId {
    pub slot: u8,
    pub generation: u64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum State {
    Empty,
    Ready,
    Running,
    Sleeping {
        until: u64,
    },
    Blocked {
        channel: u64,
        kind: WaitKind,
        deadline: Option<u64>,
    },
    WaitingChild {
        child: TaskId,
    },
    Exited {
        code: i32,
    },
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum WaitKind {
    Send,
    Receive,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum WakeCause {
    SpaceAvailable,
    MessageAvailable,
    ChannelClosed,
    TimedOut,
    ChildExited,
    WaitCancelled,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChildWait {
    Blocked(Option<TaskId>),
    Exited(i32),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Capacity,
    IdExhausted,
    InvalidTask,
    InvalidTransition,
    InvalidDeadline,
    InInterrupt,
    QueueCorrupt,
    InvalidContext,
    NotChild,
}

/// Callee-saved AAPCS64 state used only by the cooperative EL1 switch path.
/// IRQ preemption needs a full interrupted-register frame instead.
#[repr(C, align(16))]
#[derive(Clone, Copy)]
pub struct CpuContext {
    pub sp: u64,
    pub callee: [u64; 12], // x19..x30; initial x30 is the task entry PC.
    pub vector: [u128; 8], // q8..q15
}

impl CpuContext {
    pub const EMPTY: Self = Self {
        sp: 0,
        callee: [0; 12],
        vector: [0; 8],
    };

    pub fn initial(stack_top: u64, entry_pc: u64) -> Result<Self, Error> {
        if stack_top == 0 || stack_top & 15 != 0 || entry_pc == 0 || entry_pc & 3 != 0 {
            return Err(Error::InvalidContext);
        }
        let mut context = Self::EMPTY;
        context.sp = stack_top;
        context.callee[11] = entry_pc;
        Ok(context)
    }
}

const _: () = {
    assert!(core::mem::offset_of!(CpuContext, sp) == 0);
    assert!(core::mem::offset_of!(CpuContext, callee) == 8);
    assert!(core::mem::offset_of!(CpuContext, vector) == 112);
    assert!(core::mem::size_of::<CpuContext>() == 240);
};

#[derive(Clone, Copy)]
struct Task {
    generation: u64,
    state: State,
    wake: Option<WakeCause>,
    parent: Option<TaskId>,
}

impl Task {
    const EMPTY: Self = Self {
        generation: 0,
        state: State::Empty,
        wake: None,
        parent: None,
    };
}

pub struct Scheduler {
    tasks: [Task; MAX_TASKS],
    ready: [Option<TaskId>; MAX_TASKS],
    ready_len: usize,
    running: Option<TaskId>,
    next_generation: u64,
}

/// K467's separate, two-CPU ready-token experiment. Unlike the historical
/// single-CPU Scheduler, it records one running owner per CPU. A claimed token
/// is not yet a switched CpuContext; the guest test executes bounded work
/// under each claim and checks that no token is active twice.
pub struct SmpReadyQueue {
    owned: AtomicBool,
    contentions: AtomicU64,
    state: UnsafeCell<SmpReadyState>,
}

// SAFETY: all state reads and writes are under the acquire/release lock. The
// caller must place the object in Inner-Shareable Normal RAM with both MMUs on.
unsafe impl Sync for SmpReadyQueue {}

struct SmpReadyState {
    ready: [Option<TaskId>; MAX_TASKS],
    ready_len: usize,
    running: [Option<TaskId>; 2],
    claims: u64,
    finishes: u64,
    per_task: [u64; MAX_TASKS],
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SmpReadySnapshot {
    pub ready_len: usize,
    pub running: [Option<TaskId>; 2],
    pub claims: u64,
    pub finishes: u64,
    pub per_task: [u64; MAX_TASKS],
    pub contentions: u64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SmpReadyError {
    BadCpu,
    AlreadyRunning,
    WrongOwner,
    QueueCorrupt,
}

struct SmpReadyGuard<'a>(&'a SmpReadyQueue);

impl Drop for SmpReadyGuard<'_> {
    fn drop(&mut self) {
        self.0.owned.store(false, Ordering::Release);
    }
}

impl SmpReadyQueue {
    pub const fn new() -> Self {
        Self {
            owned: AtomicBool::new(false),
            contentions: AtomicU64::new(0),
            state: UnsafeCell::new(SmpReadyState {
                ready: [
                    Some(TaskId {
                        slot: 0,
                        generation: 1,
                    }),
                    Some(TaskId {
                        slot: 1,
                        generation: 1,
                    }),
                    Some(TaskId {
                        slot: 2,
                        generation: 1,
                    }),
                    Some(TaskId {
                        slot: 3,
                        generation: 1,
                    }),
                ],
                ready_len: MAX_TASKS,
                running: [None, None],
                claims: 0,
                finishes: 0,
                per_task: [0; MAX_TASKS],
            }),
        }
    }

    fn with<R>(&self, operation: impl FnOnce(&mut SmpReadyState) -> R) -> R {
        while self
            .owned
            .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            self.contentions.fetch_add(1, Ordering::Relaxed);
            core::hint::spin_loop();
        }
        let guard = SmpReadyGuard(self);
        // SAFETY: the lock grants unique access until the release guard drops.
        let result = operation(unsafe { &mut *self.state.get() });
        drop(guard);
        result
    }

    pub fn claim(&self, cpu: usize) -> Result<Option<TaskId>, SmpReadyError> {
        if cpu >= 2 {
            return Err(SmpReadyError::BadCpu);
        }
        self.with(|state| {
            for _ in 0..512 {
                core::hint::spin_loop();
            }
            if state.running[cpu].is_some() {
                return Err(SmpReadyError::AlreadyRunning);
            }
            if state.ready_len == 0 {
                return Ok(None);
            }
            let id = state.ready[0].ok_or(SmpReadyError::QueueCorrupt)?;
            if id.slot as usize >= MAX_TASKS
                || id.generation != 1
                || state.running.contains(&Some(id))
            {
                return Err(SmpReadyError::QueueCorrupt);
            }
            state.ready.copy_within(1..state.ready_len, 0);
            state.ready_len -= 1;
            state.ready[state.ready_len] = None;
            state.running[cpu] = Some(id);
            state.claims += 1;
            Ok(Some(id))
        })
    }

    pub fn finish(&self, cpu: usize, id: TaskId) -> Result<(), SmpReadyError> {
        if cpu >= 2 {
            return Err(SmpReadyError::BadCpu);
        }
        self.with(|state| {
            for _ in 0..512 {
                core::hint::spin_loop();
            }
            if state.running[cpu] != Some(id) {
                return Err(SmpReadyError::WrongOwner);
            }
            if state.ready_len == MAX_TASKS
                || state.ready[..state.ready_len].contains(&Some(id))
                || id.slot as usize >= MAX_TASKS
                || id.generation != 1
            {
                return Err(SmpReadyError::QueueCorrupt);
            }
            state.running[cpu] = None;
            state.ready[state.ready_len] = Some(id);
            state.ready_len += 1;
            state.finishes += 1;
            state.per_task[id.slot as usize] += 1;
            Ok(())
        })
    }

    pub fn snapshot(&self) -> SmpReadySnapshot {
        self.with(|state| SmpReadySnapshot {
            ready_len: state.ready_len,
            running: state.running,
            claims: state.claims,
            finishes: state.finishes,
            per_task: state.per_task,
            contentions: self.contentions.load(Ordering::Relaxed),
        })
    }
}

/// Bounded timer-wakeup ownership, independent of CPU register-frame storage.
/// A slot can occur at most once. A wake removes it before the caller may
/// select that task's saved frame for execution.
#[derive(Clone, Copy)]
pub struct SleepQueue {
    entries: [Option<(usize, u64)>; MAX_TASKS],
}

/// Ownership for one task stack. A released stack may be claimed only by a
/// newer generation, never by the exited task whose context used to own it.
#[derive(Clone, Copy)]
pub struct StackLease {
    owner: Option<TaskId>,
    last_generation: u64,
}

/// Two runnable EL1 tasks with fixed base priority and bounded aging. A
/// waiting task gains one point per quantum; a selected task loses its age.
/// Equal scores keep the current task, avoiding gratuitous frame swaps.
#[derive(Clone, Copy)]
pub struct AgingPair {
    age: [u8; 2],
}

impl AgingPair {
    pub const fn new() -> Self {
        Self { age: [0; 2] }
    }

    pub fn choose(&mut self, current: usize) -> Result<usize, Error> {
        if current >= 2 {
            return Err(Error::InvalidTask);
        }
        let waiting = 1 - current;
        self.age[waiting] = self.age[waiting].saturating_add(1);
        let scores = [2u16 + self.age[0] as u16, self.age[1] as u16];
        let next = if scores[0] > scores[1] {
            0
        } else if scores[1] > scores[0] {
            1
        } else {
            current
        };
        self.age[next] = 0;
        Ok(next)
    }
}

impl Default for AgingPair {
    fn default() -> Self {
        Self::new()
    }
}

impl StackLease {
    pub const fn new() -> Self {
        Self {
            owner: None,
            last_generation: 0,
        }
    }

    pub fn owner(&self) -> Option<TaskId> {
        self.owner
    }

    pub fn claim(&mut self, id: TaskId) -> Result<(), Error> {
        if id.generation <= self.last_generation {
            return Err(Error::InvalidTask);
        }
        if self.owner.is_some() {
            return Err(Error::Capacity);
        }
        self.owner = Some(id);
        self.last_generation = id.generation;
        Ok(())
    }

    pub fn release(&mut self, id: TaskId) -> Result<(), Error> {
        if self.owner != Some(id) {
            return Err(Error::InvalidTask);
        }
        self.owner = None;
        Ok(())
    }
}

impl Default for StackLease {
    fn default() -> Self {
        Self::new()
    }
}

impl SleepQueue {
    pub const fn new() -> Self {
        Self {
            entries: [None; MAX_TASKS],
        }
    }

    pub fn contains(&self, slot: usize) -> bool {
        self.entries
            .iter()
            .any(|entry| entry.is_some_and(|(id, _)| id == slot))
    }

    pub fn is_empty(&self) -> bool {
        self.entries.iter().all(Option::is_none)
    }

    pub fn sleep(&mut self, slot: usize, until: u64, now: u64) -> Result<(), Error> {
        if slot >= MAX_TASKS {
            return Err(Error::InvalidTask);
        }
        if until <= now {
            return Err(Error::InvalidDeadline);
        }
        if self.contains(slot) {
            return Err(Error::InvalidTransition);
        }
        let vacant = self
            .entries
            .iter_mut()
            .find(|entry| entry.is_none())
            .ok_or(Error::Capacity)?;
        *vacant = Some((slot, until));
        Ok(())
    }

    pub fn wake_due(&mut self, now: u64) -> Option<usize> {
        let index = self
            .entries
            .iter()
            .position(|entry| entry.is_some_and(|(_, until)| until <= now))?;
        self.entries[index].take().map(|(slot, _)| slot)
    }
}

impl Default for SleepQueue {
    fn default() -> Self {
        Self::new()
    }
}

impl Default for Scheduler {
    fn default() -> Self {
        Self::new()
    }
}

impl Scheduler {
    pub const fn new() -> Self {
        Self {
            tasks: [Task::EMPTY; MAX_TASKS],
            ready: [None; MAX_TASKS],
            ready_len: 0,
            running: None,
            next_generation: 1,
        }
    }

    pub fn running(&self) -> Option<TaskId> {
        self.running
    }

    pub fn ready_len(&self) -> usize {
        self.ready_len
    }

    pub fn state(&self, id: TaskId) -> Result<State, Error> {
        let Some(task) = self.tasks.get(id.slot as usize) else {
            return Err(Error::InvalidTask);
        };
        if task.generation != id.generation || task.state == State::Empty {
            return Err(Error::InvalidTask);
        }
        Ok(task.state)
    }

    pub fn spawn(&mut self) -> Result<TaskId, Error> {
        self.spawn_with_parent(None)
    }

    pub fn spawn_child(&mut self, parent: TaskId) -> Result<TaskId, Error> {
        if self.running != Some(parent) || self.state(parent) != Ok(State::Running) {
            return Err(Error::InvalidTransition);
        }
        self.spawn_with_parent(Some(parent))
    }

    fn spawn_with_parent(&mut self, parent: Option<TaskId>) -> Result<TaskId, Error> {
        let slot = self
            .tasks
            .iter()
            .position(|task| task.state == State::Empty)
            .ok_or(Error::Capacity)?;
        let generation = self.next_generation;
        self.next_generation = generation.checked_add(1).ok_or(Error::IdExhausted)?;
        let id = TaskId {
            slot: slot as u8,
            generation,
        };
        self.tasks[slot] = Task {
            generation,
            state: State::Ready,
            wake: None,
            parent,
        };
        self.enqueue(id)?;
        Ok(id)
    }

    /// A wait is authorized by the full parent/child TaskId pair, including
    /// both generations. One successful collection empties the child slot.
    pub fn wait_child(&mut self, parent: TaskId, child: TaskId) -> Result<ChildWait, Error> {
        if self.running != Some(parent) || self.state(parent) != Ok(State::Running) {
            return Err(Error::InvalidTransition);
        }
        let Some(task) = self.tasks.get(child.slot as usize) else {
            return Err(Error::NotChild);
        };
        if task.generation != child.generation
            || task.parent != Some(parent)
            || task.state == State::Empty
        {
            return Err(Error::NotChild);
        }
        if let State::Exited { code } = task.state {
            self.tasks[child.slot as usize] = Task::EMPTY;
            return Ok(ChildWait::Exited(code));
        }
        self.tasks[parent.slot as usize].state = State::WaitingChild { child };
        self.tasks[parent.slot as usize].wake = None;
        self.running = None;
        Ok(ChildWait::Blocked(self.run_next()?))
    }

    /// EL1 cancellation of one exact parent/child wait. The child remains
    /// alive and may still be collected later; a repeated cancellation cannot
    /// enqueue the parent twice.
    pub fn cancel_child_wait(&mut self, parent: TaskId, child: TaskId) -> Result<(), Error> {
        if self.state(parent)? != (State::WaitingChild { child }) {
            return Err(Error::InvalidTransition);
        }
        let task = self.tasks.get(child.slot as usize).ok_or(Error::NotChild)?;
        if task.generation != child.generation
            || task.parent != Some(parent)
            || task.state == State::Empty
        {
            return Err(Error::NotChild);
        }
        self.tasks[parent.slot as usize].state = State::Ready;
        self.tasks[parent.slot as usize].wake = Some(WakeCause::WaitCancelled);
        self.enqueue(parent)
    }

    pub fn run_next(&mut self) -> Result<Option<TaskId>, Error> {
        if self.running.is_some() {
            return Err(Error::InvalidTransition);
        }
        if self.ready_len == 0 {
            return Ok(None);
        }
        let id = self.ready[0].ok_or(Error::QueueCorrupt)?;
        if self.state(id)? != State::Ready {
            return Err(Error::QueueCorrupt);
        }
        self.ready.copy_within(1..self.ready_len, 0);
        self.ready_len -= 1;
        self.ready[self.ready_len] = None;
        self.tasks[id.slot as usize].state = State::Running;
        self.running = Some(id);
        Ok(Some(id))
    }

    pub fn yield_current(&mut self, in_irq: bool) -> Result<Option<TaskId>, Error> {
        if in_irq {
            return Err(Error::InInterrupt);
        }
        let id = self.running.ok_or(Error::InvalidTransition)?;
        self.tasks[id.slot as usize].state = State::Ready;
        self.running = None;
        self.enqueue(id)?;
        self.run_next()
    }

    pub fn sleep_current(&mut self, until: u64, now: u64) -> Result<Option<TaskId>, Error> {
        if until <= now {
            return Err(Error::InvalidDeadline);
        }
        let id = self.running.ok_or(Error::InvalidTransition)?;
        self.tasks[id.slot as usize].state = State::Sleeping { until };
        self.running = None;
        self.run_next()
    }

    /// Remove the running task from the ready set. On a single CPU, callers
    /// register its channel waiter before letting the selected peer run.
    pub fn block_current(
        &mut self,
        channel: u64,
        kind: WaitKind,
        deadline: Option<u64>,
        now: u64,
    ) -> Result<Option<TaskId>, Error> {
        if channel == 0 || deadline.is_some_and(|until| until <= now) {
            return Err(Error::InvalidDeadline);
        }
        let id = self.running.ok_or(Error::InvalidTransition)?;
        self.tasks[id.slot as usize].state = State::Blocked {
            channel,
            kind,
            deadline,
        };
        self.tasks[id.slot as usize].wake = None;
        self.running = None;
        self.run_next()
    }

    pub fn wake_blocked(
        &mut self,
        id: TaskId,
        channel: u64,
        kind: WaitKind,
        cause: WakeCause,
    ) -> Result<(), Error> {
        if !matches!(self.state(id)?, State::Blocked { channel: actual, kind: actual_kind, .. } if actual == channel && actual_kind == kind)
        {
            return Err(Error::InvalidTransition);
        }
        self.tasks[id.slot as usize].state = State::Ready;
        self.tasks[id.slot as usize].wake = Some(cause);
        self.enqueue(id)
    }

    pub fn take_wake_cause(&mut self, id: TaskId) -> Result<Option<WakeCause>, Error> {
        if self.running != Some(id) || self.state(id)? != State::Running {
            return Err(Error::InvalidTransition);
        }
        Ok(self.tasks[id.slot as usize].wake.take())
    }

    pub fn wake_expired(&mut self, now: u64) -> Result<usize, Error> {
        let mut count = 0;
        for slot in 0..MAX_TASKS {
            if let State::Blocked {
                deadline: Some(until),
                ..
            } = self.tasks[slot].state
            {
                if until <= now {
                    let id = TaskId {
                        slot: slot as u8,
                        generation: self.tasks[slot].generation,
                    };
                    self.tasks[slot].state = State::Ready;
                    self.tasks[slot].wake = Some(WakeCause::TimedOut);
                    self.enqueue(id)?;
                    count += 1;
                }
            }
        }
        Ok(count)
    }

    pub fn wake_due(&mut self, now: u64) -> Result<usize, Error> {
        let mut count = 0;
        for slot in 0..MAX_TASKS {
            if let State::Sleeping { until } = self.tasks[slot].state {
                if until <= now {
                    let id = TaskId {
                        slot: slot as u8,
                        generation: self.tasks[slot].generation,
                    };
                    self.tasks[slot].state = State::Ready;
                    self.enqueue(id)?;
                    count += 1;
                }
            }
        }
        Ok(count)
    }

    pub fn exit_current(&mut self, code: i32) -> Result<Option<TaskId>, Error> {
        let id = self.running.ok_or(Error::InvalidTransition)?;
        let parent = self.tasks[id.slot as usize].parent;
        let waiting_parent =
            parent.filter(|parent| self.state(*parent) == Ok(State::WaitingChild { child: id }));
        self.tasks[id.slot as usize].state = State::Exited { code };
        self.running = None;
        if let Some(parent) = waiting_parent {
            self.tasks[parent.slot as usize].state = State::Ready;
            self.tasks[parent.slot as usize].wake = Some(WakeCause::ChildExited);
            self.enqueue(parent)?;
        }
        self.run_next()
    }

    /// Reuse is explicit: an exited slot remains occupied until reaped. A new
    /// occupant receives a new generation, so an old ID never aliases it.
    pub fn reap(&mut self, id: TaskId) -> Result<i32, Error> {
        let State::Exited { code } = self.state(id)? else {
            return Err(Error::InvalidTransition);
        };
        self.tasks[id.slot as usize].state = State::Empty;
        self.tasks[id.slot as usize].parent = None;
        Ok(code)
    }

    pub fn deadline(now: u64, quantum: u64) -> Result<u64, Error> {
        if quantum == 0 {
            return Err(Error::InvalidDeadline);
        }
        now.checked_add(quantum).ok_or(Error::InvalidDeadline)
    }

    fn enqueue(&mut self, id: TaskId) -> Result<(), Error> {
        if self.state(id)? != State::Ready
            || self.ready[..self.ready_len].contains(&Some(id))
            || self.ready_len == MAX_TASKS
        {
            return Err(Error::QueueCorrupt);
        }
        self.ready[self.ready_len] = Some(id);
        self.ready_len += 1;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    extern crate std;

    #[test]
    fn k467_two_threads_never_claim_one_ready_task_together() {
        use core::sync::atomic::AtomicUsize;
        use std::sync::{Arc, Barrier};
        use std::thread;

        let queue = Arc::new(SmpReadyQueue::new());
        let active = Arc::new(core::array::from_fn::<_, 4, _>(|_| AtomicUsize::new(0)));
        let barrier = Arc::new(Barrier::new(2));
        let mut workers = std::vec::Vec::new();
        for cpu in 0..2 {
            let queue = Arc::clone(&queue);
            let active = Arc::clone(&active);
            let barrier = Arc::clone(&barrier);
            workers.push(thread::spawn(move || {
                barrier.wait();
                for _ in 0..128 {
                    let id = queue.claim(cpu).unwrap().unwrap();
                    assert_eq!(active[id.slot as usize].swap(cpu + 1, Ordering::AcqRel), 0);
                    for _ in 0..512 {
                        core::hint::spin_loop();
                    }
                    assert_eq!(active[id.slot as usize].swap(0, Ordering::Release), cpu + 1);
                    queue.finish(cpu, id).unwrap();
                }
            }));
        }
        for worker in workers {
            worker.join().unwrap();
        }
        let state = queue.snapshot();
        assert_eq!(state.ready_len, 4);
        assert_eq!(state.running, [None, None]);
        assert_eq!((state.claims, state.finishes), (256, 256));
        assert_eq!(state.per_task.iter().sum::<u64>(), 256);
        assert!(state.per_task.iter().all(|count| *count > 0));
        assert!(state.contentions > 0);
        let id = queue.claim(0).unwrap().unwrap();
        let before = queue.snapshot();
        assert_eq!(queue.claim(0), Err(SmpReadyError::AlreadyRunning));
        assert_eq!(queue.finish(1, id), Err(SmpReadyError::WrongOwner));
        assert_eq!(queue.finish(2, id), Err(SmpReadyError::BadCpu));
        assert_eq!(queue.snapshot(), before);
        queue.finish(0, id).unwrap();
    }

    #[test]
    fn channel_sender_block_is_not_runnable_until_one_wake() {
        let mut scheduler = Scheduler::new();
        let sender = scheduler.spawn().unwrap();
        let receiver = scheduler.spawn().unwrap();
        assert_eq!(scheduler.run_next(), Ok(Some(sender)));
        assert_eq!(
            scheduler.block_current(7, WaitKind::Send, None, 0),
            Ok(Some(receiver))
        );
        assert_eq!(
            scheduler.state(sender),
            Ok(State::Blocked {
                channel: 7,
                kind: WaitKind::Send,
                deadline: None
            })
        );
        assert_eq!(scheduler.ready_len(), 0);
        assert_eq!(
            scheduler.wake_blocked(sender, 8, WaitKind::Send, WakeCause::SpaceAvailable),
            Err(Error::InvalidTransition)
        );
        assert_eq!(
            scheduler.wake_blocked(sender, 7, WaitKind::Send, WakeCause::SpaceAvailable),
            Ok(())
        );
        assert_eq!(
            scheduler.wake_blocked(sender, 7, WaitKind::Send, WakeCause::SpaceAvailable),
            Err(Error::InvalidTransition)
        );
        assert_eq!(scheduler.yield_current(false), Ok(Some(sender)));
        assert_eq!(
            scheduler.take_wake_cause(sender),
            Ok(Some(WakeCause::SpaceAvailable))
        );
        assert_eq!(scheduler.take_wake_cause(sender), Ok(None));
    }

    #[test]
    fn channel_receiver_blocks_with_empty_ready_queue_then_wakes_once() {
        let mut scheduler = Scheduler::new();
        let receiver = scheduler.spawn().unwrap();
        assert_eq!(scheduler.run_next(), Ok(Some(receiver)));
        assert_eq!(
            scheduler.block_current(9, WaitKind::Receive, None, 0),
            Ok(None)
        );
        assert_eq!(scheduler.ready_len(), 0);
        assert_eq!(scheduler.run_next(), Ok(None));
        assert_eq!(
            scheduler.wake_blocked(receiver, 9, WaitKind::Receive, WakeCause::MessageAvailable),
            Ok(())
        );
        assert_eq!(scheduler.run_next(), Ok(Some(receiver)));
        assert_eq!(
            scheduler.take_wake_cause(receiver),
            Ok(Some(WakeCause::MessageAvailable))
        );
    }

    #[test]
    fn channel_receive_deadline_wakes_without_duplicate_ready_entry() {
        let mut scheduler = Scheduler::new();
        let receiver = scheduler.spawn().unwrap();
        assert_eq!(scheduler.run_next(), Ok(Some(receiver)));
        assert_eq!(
            scheduler.block_current(9, WaitKind::Receive, Some(10), 10),
            Err(Error::InvalidDeadline)
        );
        assert_eq!(scheduler.running(), Some(receiver));
        assert_eq!(
            scheduler.block_current(9, WaitKind::Receive, Some(10), 5),
            Ok(None)
        );
        assert_eq!(scheduler.wake_expired(9), Ok(0));
        assert_eq!(scheduler.ready_len(), 0);
        assert_eq!(scheduler.wake_expired(10), Ok(1));
        assert_eq!(scheduler.wake_expired(10), Ok(0));
        assert_eq!(scheduler.ready_len(), 1);
        assert_eq!(scheduler.run_next(), Ok(Some(receiver)));
        assert_eq!(
            scheduler.take_wake_cause(receiver),
            Ok(Some(WakeCause::TimedOut))
        );
    }

    #[test]
    fn state_transitions_and_unique_generation_after_explicit_reap() {
        let mut s = Scheduler::new();
        let a = s.spawn().unwrap();
        assert_eq!(s.state(a), Ok(State::Ready));
        assert_eq!(s.run_next(), Ok(Some(a)));
        assert_eq!(s.state(a), Ok(State::Running));
        assert_eq!(s.reap(a), Err(Error::InvalidTransition));
        assert_eq!(s.exit_current(42), Ok(None));
        assert_eq!(s.state(a), Ok(State::Exited { code: 42 }));
        assert_eq!(s.reap(a), Ok(42));
        let b = s.spawn().unwrap();
        assert_eq!(a.slot, b.slot);
        assert_ne!(a.generation, b.generation);
        assert_eq!(s.state(a), Err(Error::InvalidTask));
    }

    #[test]
    fn stack_lease_requires_release_and_new_generation_before_reuse() {
        let old = TaskId {
            slot: 0,
            generation: 1,
        };
        let other = TaskId {
            slot: 1,
            generation: 2,
        };
        let new = TaskId {
            slot: 0,
            generation: 3,
        };
        let mut lease = StackLease::new();
        assert_eq!(lease.claim(old), Ok(()));
        assert_eq!(lease.claim(other), Err(Error::Capacity));
        assert_eq!(lease.release(other), Err(Error::InvalidTask));
        assert_eq!(lease.owner(), Some(old));
        assert_eq!(lease.release(old), Ok(()));
        assert_eq!(lease.claim(old), Err(Error::InvalidTask));
        assert_eq!(lease.claim(new), Ok(()));
        assert_eq!(lease.release(old), Err(Error::InvalidTask));
        assert_eq!(lease.owner(), Some(new));
    }

    #[test]
    fn aging_priority_prevents_low_priority_starvation() {
        let mut policy = AgingPair::new();
        assert_eq!(policy.choose(2), Err(Error::InvalidTask));
        let mut current = 0;
        let mut low_runs = 0;
        let mut last_low = 0;
        for quantum in 1..=10_000 {
            current = policy.choose(current).unwrap();
            if current == 1 {
                assert!(quantum - last_low <= 4);
                last_low = quantum;
                low_runs += 1;
            }
        }
        assert!(low_runs >= 2_000);
        assert_eq!(policy.choose(usize::MAX), Err(Error::InvalidTask));
    }

    #[test]
    fn peer_reads_retained_exit_code_after_yield_and_before_reap() {
        let mut scheduler = Scheduler::new();
        let a = scheduler.spawn().unwrap();
        let b = scheduler.spawn().unwrap();
        assert_eq!(scheduler.run_next(), Ok(Some(a)));
        assert_eq!(scheduler.yield_current(false), Ok(Some(b)));
        assert_eq!(scheduler.yield_current(false), Ok(Some(a)));
        assert_eq!(scheduler.exit_current(0x0a), Ok(Some(b)));
        assert_eq!(scheduler.running(), Some(b));
        assert_eq!(scheduler.state(a), Ok(State::Exited { code: 0x0a }));
        assert_eq!(scheduler.reap(a), Ok(0x0a));
        assert_eq!(scheduler.state(a), Err(Error::InvalidTask));
    }

    #[test]
    fn ready_queue_is_single_owned_and_round_robin() {
        let mut s = Scheduler::new();
        let a = s.spawn().unwrap();
        let b = s.spawn().unwrap();
        assert_eq!(s.ready_len(), 2);
        assert_eq!(s.run_next(), Ok(Some(a)));
        assert_eq!(s.run_next(), Err(Error::InvalidTransition));
        assert_eq!(s.yield_current(false), Ok(Some(b)));
        assert_eq!(s.ready_len(), 1);
        assert_eq!(s.yield_current(false), Ok(Some(a)));
        assert_eq!(s.ready_len(), 1);
    }

    #[test]
    fn sleep_waits_for_deadline_then_rejoins_tail() {
        let mut s = Scheduler::new();
        let a = s.spawn().unwrap();
        let b = s.spawn().unwrap();
        assert_eq!(s.run_next(), Ok(Some(a)));
        assert_eq!(s.sleep_current(10, 5), Ok(Some(b)));
        assert_eq!(s.state(a), Ok(State::Sleeping { until: 10 }));
        assert_eq!(s.wake_due(9), Ok(0));
        assert_eq!(s.ready_len(), 0);
        assert_eq!(s.wake_due(10), Ok(1));
        assert_eq!(s.wake_due(10), Ok(0));
        assert_eq!(s.yield_current(false), Ok(Some(a)));
    }

    #[test]
    fn sleep_queue_rejects_duplicate_and_early_wake_without_losing_owner() {
        let mut queue = SleepQueue::new();
        assert_eq!(queue.sleep(1, 4, 2), Ok(()));
        assert_eq!(queue.sleep(1, 5, 2), Err(Error::InvalidTransition));
        assert_eq!(queue.sleep(2, 2, 2), Err(Error::InvalidDeadline));
        assert_eq!(queue.sleep(MAX_TASKS, 5, 2), Err(Error::InvalidTask));
        assert_eq!(queue.wake_due(3), None);
        assert!(queue.contains(1));
        assert_eq!(queue.wake_due(4), Some(1));
        assert!(!queue.contains(1));
        assert_eq!(queue.wake_due(4), None);
        assert!(queue.is_empty());
    }

    #[test]
    fn exit_is_not_scheduled_and_idle_is_explicit() {
        let mut s = Scheduler::new();
        assert_eq!(s.run_next(), Ok(None));
        let a = s.spawn().unwrap();
        s.run_next().unwrap();
        assert_eq!(s.exit_current(-7), Ok(None));
        assert_eq!(s.run_next(), Ok(None));
        assert_eq!(s.state(a), Ok(State::Exited { code: -7 }));
        assert_eq!(s.reap(a), Ok(-7));
    }

    #[test]
    fn invalid_transition_irq_and_deadline_fail_closed() {
        let mut s = Scheduler::new();
        assert_eq!(s.yield_current(false), Err(Error::InvalidTransition));
        assert_eq!(s.sleep_current(1, 0), Err(Error::InvalidTransition));
        let a = s.spawn().unwrap();
        s.run_next().unwrap();
        assert_eq!(s.yield_current(true), Err(Error::InInterrupt));
        assert_eq!(s.running(), Some(a));
        assert_eq!(s.sleep_current(5, 5), Err(Error::InvalidDeadline));
        assert_eq!(s.state(a), Ok(State::Running));
        assert_eq!(Scheduler::deadline(10, 0), Err(Error::InvalidDeadline));
        assert_eq!(
            Scheduler::deadline(u64::MAX, 1),
            Err(Error::InvalidDeadline)
        );
        assert_eq!(Scheduler::deadline(10, 4), Ok(14));
    }

    #[test]
    fn irq_yield_rejection_preserves_running_task_and_ready_queue() {
        let mut scheduler = Scheduler::new();
        let a = scheduler.spawn().unwrap();
        let b = scheduler.spawn().unwrap();
        assert_eq!(scheduler.run_next(), Ok(Some(a)));
        assert_eq!(scheduler.yield_current(true), Err(Error::InInterrupt));
        assert_eq!(scheduler.running(), Some(a));
        assert_eq!(scheduler.state(a), Ok(State::Running));
        assert_eq!(scheduler.state(b), Ok(State::Ready));
        assert_eq!(scheduler.ready_len(), 1);
        assert_eq!(scheduler.yield_current(false), Ok(Some(b)));
        assert_eq!(scheduler.running(), Some(b));
    }

    #[test]
    fn quantum_deadline_boundary_never_wraps_or_arms_zero() {
        assert_eq!(Scheduler::deadline(u64::MAX - 1, 1), Ok(u64::MAX));
        assert_eq!(
            Scheduler::deadline(u64::MAX - 1, 2),
            Err(Error::InvalidDeadline)
        );
        assert_eq!(
            Scheduler::deadline(u64::MAX, 1),
            Err(Error::InvalidDeadline)
        );
        assert_eq!(Scheduler::deadline(1, 0), Err(Error::InvalidDeadline));
        assert_eq!(crate::irq::next_timer_deadline(u64::MAX - 1, 2, 0), None);
        assert_eq!(
            crate::irq::next_timer_deadline(u64::MAX - 2, 1, u64::MAX),
            None
        );
    }

    #[test]
    fn bounded_table_and_generation_overflow_do_not_mutate_state() {
        let mut s = Scheduler::new();
        let mut ids = [TaskId {
            slot: 0,
            generation: 0,
        }; MAX_TASKS];
        for id in &mut ids {
            *id = s.spawn().unwrap();
        }
        assert_eq!(s.spawn(), Err(Error::Capacity));
        assert_eq!(s.ready_len(), MAX_TASKS);
        let mut other = Scheduler::new();
        other.next_generation = u64::MAX;
        assert_eq!(other.spawn(), Err(Error::IdExhausted));
        assert_eq!(other.ready_len(), 0);
    }

    #[test]
    fn initial_context_has_aligned_stack_and_entry_return_address() {
        let context = CpuContext::initial(0x4010_4000, 0x4008_1234).unwrap();
        assert_eq!(context.sp, 0x4010_4000);
        assert_eq!(context.callee[11], 0x4008_1234);
        assert!(context.callee[..11].iter().all(|value| *value == 0));
        assert_eq!(
            CpuContext::initial(0x4010_4008, 0x4008_1234).err(),
            Some(Error::InvalidContext)
        );
        assert_eq!(
            CpuContext::initial(0x4010_4000, 0x4008_1235).err(),
            Some(Error::InvalidContext)
        );
        assert_eq!(core::mem::size_of::<CpuContext>(), 240);
    }

    #[test]
    fn k316_parent_wait_rejects_self_foreign_stale_and_wrong_caller() {
        let mut s = Scheduler::new();
        let parent = s.spawn().unwrap();
        assert_eq!(s.spawn_child(parent), Err(Error::InvalidTransition));
        assert_eq!(s.run_next(), Ok(Some(parent)));
        let child = s.spawn_child(parent).unwrap();
        let foreign = s.spawn().unwrap();
        let stale = TaskId {
            slot: child.slot,
            generation: child.generation + 1,
        };
        for target in [parent, foreign, stale] {
            assert_eq!(s.wait_child(parent, target), Err(Error::NotChild));
        }
        assert_eq!(s.wait_child(foreign, child), Err(Error::InvalidTransition));
        assert_eq!(s.running(), Some(parent));
        assert_eq!(s.ready_len(), 2);
        assert_eq!(s.state(child), Ok(State::Ready));
        assert_eq!(s.state(foreign), Ok(State::Ready));
    }

    #[test]
    fn k316_valid_child_wait_blocks_parent_and_only_child_exit_wakes_it() {
        let mut s = Scheduler::new();
        let parent = s.spawn().unwrap();
        assert_eq!(s.run_next(), Ok(Some(parent)));
        let child = s.spawn_child(parent).unwrap();
        assert_eq!(
            s.wait_child(parent, child),
            Ok(ChildWait::Blocked(Some(child)))
        );
        assert_eq!(s.state(parent), Ok(State::WaitingChild { child }));
        assert_eq!(s.ready_len(), 0);
        assert_eq!(s.exit_current(0x2a), Ok(Some(parent)));
        assert_eq!(s.take_wake_cause(parent), Ok(Some(WakeCause::ChildExited)));
        assert_eq!(s.take_wake_cause(parent), Ok(None));
        assert_eq!(s.wait_child(parent, child), Ok(ChildWait::Exited(0x2a)));
        assert_eq!(s.ready_len(), 0);
    }

    #[test]
    fn k317_exit_code_is_read_once_and_recycled_slot_rejects_old_generation() {
        let mut s = Scheduler::new();
        let parent = s.spawn().unwrap();
        assert_eq!(s.run_next(), Ok(Some(parent)));
        let child = s.spawn_child(parent).unwrap();
        assert_eq!(
            s.wait_child(parent, child),
            Ok(ChildWait::Blocked(Some(child)))
        );
        assert_eq!(s.exit_current(-7), Ok(Some(parent)));
        assert_eq!(s.wait_child(parent, child), Ok(ChildWait::Exited(-7)));
        assert_eq!(s.wait_child(parent, child), Err(Error::NotChild));
        let replacement = s.spawn_child(parent).unwrap();
        assert_eq!(replacement.slot, child.slot);
        assert!(replacement.generation > child.generation);
        assert_eq!(s.wait_child(parent, child), Err(Error::NotChild));
        assert_eq!(s.state(replacement), Ok(State::Ready));
    }

    #[test]
    fn k317_child_exited_before_wait_returns_signed_code_without_blocking() {
        let mut s = Scheduler::new();
        let parent = s.spawn().unwrap();
        assert_eq!(s.run_next(), Ok(Some(parent)));
        let child = s.spawn_child(parent).unwrap();
        assert_eq!(s.yield_current(false), Ok(Some(child)));
        assert_eq!(s.exit_current(-19), Ok(Some(parent)));
        assert_eq!(s.take_wake_cause(parent), Ok(None));
        assert_eq!(s.wait_child(parent, child), Ok(ChildWait::Exited(-19)));
        assert_eq!(s.wait_child(parent, child), Err(Error::NotChild));
        assert_eq!(s.running(), Some(parent));
        assert_eq!(s.ready_len(), 0);
    }

    #[test]
    fn k318_cancelled_wait_enqueues_parent_once_and_preserves_child() {
        let mut s = Scheduler::new();
        let parent = s.spawn().unwrap();
        assert_eq!(s.run_next(), Ok(Some(parent)));
        let child = s.spawn_child(parent).unwrap();
        assert_eq!(
            s.wait_child(parent, child),
            Ok(ChildWait::Blocked(Some(child)))
        );
        assert_eq!(s.cancel_child_wait(parent, child), Ok(()));
        assert_eq!(s.ready_len(), 1);
        assert_eq!(s.state(child), Ok(State::Running));
        assert_eq!(
            s.cancel_child_wait(parent, child),
            Err(Error::InvalidTransition)
        );
        assert_eq!(s.ready_len(), 1);
        assert_eq!(s.exit_current(42), Ok(Some(parent)));
        assert_eq!(s.ready_len(), 0);
        assert_eq!(
            s.take_wake_cause(parent),
            Ok(Some(WakeCause::WaitCancelled))
        );
        assert_eq!(s.take_wake_cause(parent), Ok(None));
        assert_eq!(s.wait_child(parent, child), Ok(ChildWait::Exited(42)));
        assert_eq!(s.wait_child(parent, child), Err(Error::NotChild));
    }

    #[test]
    fn k318_wrong_child_cancel_does_not_change_wait_or_ready_queue() {
        let mut s = Scheduler::new();
        let parent = s.spawn().unwrap();
        assert_eq!(s.run_next(), Ok(Some(parent)));
        let target = s.spawn_child(parent).unwrap();
        let sibling = s.spawn_child(parent).unwrap();
        assert_eq!(
            s.wait_child(parent, target),
            Ok(ChildWait::Blocked(Some(target)))
        );
        let stale = TaskId {
            slot: target.slot,
            generation: target.generation + 1,
        };
        assert_eq!(
            s.cancel_child_wait(parent, sibling),
            Err(Error::InvalidTransition)
        );
        assert_eq!(
            s.cancel_child_wait(parent, stale),
            Err(Error::InvalidTransition)
        );
        assert_eq!(s.state(parent), Ok(State::WaitingChild { child: target }));
        assert_eq!(s.ready_len(), 1);
        assert_eq!(s.yield_current(false), Ok(Some(sibling)));
        assert_eq!(s.exit_current(3), Ok(Some(target)));
        assert_eq!(s.state(parent), Ok(State::WaitingChild { child: target }));
        assert_eq!(s.ready_len(), 0);
        assert_eq!(s.exit_current(4), Ok(Some(parent)));
        assert_eq!(s.take_wake_cause(parent), Ok(Some(WakeCause::ChildExited)));
        assert_eq!(s.wait_child(parent, target), Ok(ChildWait::Exited(4)));
        assert_eq!(s.wait_child(parent, sibling), Ok(ChildWait::Exited(3)));
    }
}