//! Bounded thread identities and custody of private, guarded kernel stacks.
//! Context saving and instruction-level switching belong to the next module.
//! A trusted backend supplies actual SP, mappings, canaries and resource release.

pub const PAGE: u64 = 4096;
pub const LOW_CANARY: u64 = 0x4b313032304c4f57;
pub const HIGH_CANARY: u64 = 0x4b31303230484947;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Capacity,
    Exhausted,
    Context,
    Layout,
    SharedStack,
    ActiveStack,
    Entry,
    Guard,
    Canary,
    NotZeroed,
    Stale,
    State,
    Backend,
    Quarantined,
    Busy,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ThreadId {
    table: u64,
    slot: usize,
    generation: u64,
}
impl ThreadId {
    pub fn table(self) -> u64 {
        self.table
    }
    pub fn slot(self) -> usize {
        self.slot
    }
    pub fn generation(self) -> u64 {
        self.generation
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum State {
    Vacant,
    Ready,
    Running,
    Sleeping,
    Dead,
    Quarantined,
}

/// The backend retains all physical pages and the unmapped guard below `base`.
/// This first thread module deliberately requires contiguous physical backing.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct StackLayout {
    pub base: u64,
    pub physical: u64,
    pub bytes: u64,
}
impl StackLayout {
    fn valid(self) -> bool {
        self.base >= PAGE
            && self.base.is_multiple_of(PAGE)
            && self.physical != 0
            && self.physical.is_multiple_of(PAGE)
            && self.bytes >= 4 * PAGE
            && self.bytes <= 16 * PAGE
            && self.bytes.is_multiple_of(PAGE)
            && self.base.checked_add(self.bytes).is_some()
            && self.physical.checked_add(self.bytes).is_some()
    }
    pub fn guard(self) -> u64 {
        self.base - PAGE
    }
    pub fn usable_bottom(self) -> u64 {
        self.base + 16
    }
    pub fn usable_top(self) -> u64 {
        self.base + self.bytes - 16
    }
    fn contains(self, sp: u64) -> bool {
        sp >= self.guard() && sp <= self.base + self.bytes
    }
    fn usable(self, sp: u64) -> bool {
        sp.is_multiple_of(16) && sp >= self.usable_bottom() && sp <= self.usable_top()
    }
    fn overlaps(self, other: Self) -> bool {
        self.guard() < other.base + other.bytes && other.guard() < self.base + self.bytes
            || self.physical < other.physical + other.bytes
                && other.physical < self.physical + self.bytes
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Execution {
    pub cpu: u32,
    pub irq_masked: bool,
    pub in_irq: bool,
    pub sp: u64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct StackObservation {
    pub layout: StackLayout,
    pub guard_descriptor: u64,
    pub low_canary: u64,
    pub high_canary: u64,
    /// Required only for initial publication; excludes the two canary words.
    pub initially_zeroed: bool,
}

/// # Safety
/// Stack values own exclusive typed backing leases and mapping custody. Dropping
/// a Stack must not release its pages. `prepare` zeros the entire backing and
/// writes both canaries; `observe` reads real mapping and canary state. Execution
/// observes the calling CPU's real SP and IRQ context. Entries are executable
/// kernel mappings. Successful `release` first excludes active use, clears
/// mappings with barriers, erases every owned page and releases its leases;
/// failure returns the same actual resource, including after partial teardown.
/// No method allocates an untracked resource, sleeps or unmasks IRQs. A backend
/// must not return aliases of resources retained by another live table.
pub unsafe trait Backend {
    type Stack;
    fn execution(&self) -> Execution;
    fn layout(&self, stack: &Self::Stack) -> Result<StackLayout, Error>;
    fn entry_executable(&self, entry: u64) -> bool;
    fn prepare(&mut self, stack: &mut Self::Stack) -> Result<(), Error>;
    fn observe(&self, stack: &Self::Stack) -> Result<StackObservation, Error>;
    fn release(&mut self, stack: Self::Stack) -> Result<(), (Error, Self::Stack)>;
}

/// Single-use wake authority for one sleep episode of one thread generation.
#[derive(Debug)]
pub struct Wait {
    id: ThreadId,
    epoch: u64,
}
impl Wait {
    pub fn thread(&self) -> ThreadId {
        self.id
    }
    pub fn epoch(&self) -> u64 {
        self.epoch
    }
}

/// Read-only initial register inputs; the context-switch service must validate
/// the live Ready thread again immediately before selecting its retained stack.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ReadyContext {
    pub id: ThreadId,
    pub entry: u64,
    pub argument: u64,
    pub initial_sp: u64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ThreadRecord {
    pub id: ThreadId,
    pub state: State,
    pub layout: StackLayout,
    pub entry: u64,
    pub argument: u64,
    pub wait_epoch: u64,
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Snapshot {
    pub live: usize,
    pub ready: usize,
    pub running: usize,
    pub sleeping: usize,
    pub dead: usize,
    pub quarantined: usize,
    pub retired: usize,
    pub pages: u64,
}

struct Slot<S> {
    generation: u64,
    wait_epoch: u64,
    state: State,
    stack: Option<S>,
    layout: StackLayout,
    entry: u64,
    argument: u64,
}
impl<S> Slot<S> {
    fn empty() -> Self {
        Self {
            generation: 0,
            wait_epoch: 0,
            state: State::Vacant,
            stack: None,
            layout: StackLayout {
                base: 0,
                physical: 0,
                bytes: 0,
            },
            entry: 0,
            argument: 0,
        }
    }
}

pub struct Table<B: Backend, const N: usize> {
    backend: B,
    identity: u64,
    cpu: u32,
    slots: [Slot<B::Stack>; N],
}

#[cfg(not(target_os = "none"))]
fn identity() -> Result<u64, Error> {
    use core::sync::atomic::{AtomicU64, Ordering};
    static NEXT: AtomicU64 = AtomicU64::new(1);
    let mut current = NEXT.load(Ordering::Relaxed);
    loop {
        let next = current.checked_add(1).ok_or(Error::Exhausted)?;
        match NEXT.compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed) {
            Ok(value) => return Ok(value),
            Err(value) => current = value,
        }
    }
}
#[cfg(target_os = "none")]
fn identity() -> Result<u64, Error> {
    unsafe extern "C" {
        fn __os_k1020_threads_identity() -> u64;
    }
    let value = unsafe { __os_k1020_threads_identity() };
    if value == 0 {
        Err(Error::Exhausted)
    } else {
        Ok(value)
    }
}

impl<B: Backend, const N: usize> Table<B, N> {
    pub fn new(backend: B) -> Result<Self, (Error, B)> {
        if N == 0 || N > 32 {
            return Err((Error::Capacity, backend));
        }
        let context = backend.execution();
        if !context.irq_masked || context.in_irq {
            return Err((Error::Context, backend));
        }
        let identity = match identity() {
            Ok(v) => v,
            Err(e) => return Err((e, backend)),
        };
        Ok(Self {
            backend,
            identity,
            cpu: context.cpu,
            slots: core::array::from_fn(|_| Slot::empty()),
        })
    }
    fn context(&self) -> Result<Execution, Error> {
        let context = self.backend.execution();
        if !context.irq_masked || context.in_irq || context.cpu != self.cpu {
            Err(Error::Context)
        } else {
            Ok(context)
        }
    }
    fn index(&self, id: ThreadId) -> Result<usize, Error> {
        if id.table != self.identity || id.slot >= N {
            return Err(Error::Stale);
        }
        let slot = &self.slots[id.slot];
        if slot.state == State::Vacant || slot.generation != id.generation {
            return Err(Error::Stale);
        }
        Ok(id.slot)
    }
    fn id(&self, index: usize) -> ThreadId {
        ThreadId {
            table: self.identity,
            slot: index,
            generation: self.slots[index].generation,
        }
    }
    fn verify_slot(&self, index: usize, initial: bool) -> Result<(), Error> {
        let slot = &self.slots[index];
        if slot.state == State::Quarantined {
            return Err(Error::Quarantined);
        }
        let actual = self
            .backend
            .observe(slot.stack.as_ref().ok_or(Error::State)?)?;
        Self::verify_observation(actual, slot.layout, initial)
    }
    fn verify_observation(
        actual: StackObservation,
        expected: StackLayout,
        initial: bool,
    ) -> Result<(), Error> {
        if actual.layout != expected {
            return Err(Error::Layout);
        }
        if actual.guard_descriptor != 0 {
            return Err(Error::Guard);
        }
        if actual.low_canary != LOW_CANARY || actual.high_canary != HIGH_CANARY {
            return Err(Error::Canary);
        }
        if initial && !actual.initially_zeroed {
            return Err(Error::NotZeroed);
        }
        Ok(())
    }
    pub fn create(
        &mut self,
        mut stack: B::Stack,
        entry: u64,
        argument: u64,
    ) -> Result<ThreadId, (Error, B::Stack)> {
        let context = match self.context() {
            Ok(v) => v,
            Err(e) => return Err((e, stack)),
        };
        let index = match self
            .slots
            .iter()
            .position(|s| s.state == State::Vacant && s.generation != u64::MAX)
        {
            Some(i) => i,
            None => return Err((Error::Capacity, stack)),
        };
        let generation = match self.slots[index].generation.checked_add(1) {
            Some(v) => v,
            None => return Err((Error::Exhausted, stack)),
        };
        let layout = match self.backend.layout(&stack) {
            Ok(v) => v,
            Err(e) => return Err((e, stack)),
        };
        if !layout.valid() {
            return Err((Error::Layout, stack));
        }
        if self
            .slots
            .iter()
            .any(|s| s.state != State::Vacant && s.layout.overlaps(layout))
        {
            return Err((Error::SharedStack, stack));
        }
        if layout.contains(context.sp) {
            return Err((Error::ActiveStack, stack));
        }
        if entry == 0 || !entry.is_multiple_of(4) || !self.backend.entry_executable(entry) {
            return Err((Error::Entry, stack));
        }
        if let Err(e) = self.backend.prepare(&mut stack) {
            return Err((e, stack));
        }
        let observation = match self.backend.observe(&stack) {
            Ok(v) => v,
            Err(e) => return Err((e, stack)),
        };
        if let Err(e) = Self::verify_observation(observation, layout, true) {
            return Err((e, stack));
        }
        self.slots[index] = Slot {
            generation,
            wait_epoch: 0,
            state: State::Ready,
            stack: Some(stack),
            layout,
            entry,
            argument,
        };
        Ok(self.id(index))
    }
    pub fn record(&self, id: ThreadId) -> Result<ThreadRecord, Error> {
        let index = self.index(id)?;
        let slot = &self.slots[index];
        Ok(ThreadRecord {
            id,
            state: slot.state,
            layout: slot.layout,
            entry: slot.entry,
            argument: slot.argument,
            wait_epoch: slot.wait_epoch,
        })
    }
    pub fn ready_context(&self, id: ThreadId) -> Result<ReadyContext, Error> {
        self.context()?;
        let index = self.index(id)?;
        if self.slots[index].state != State::Ready {
            return Err(Error::State);
        }
        self.verify_slot(index, false)?;
        let slot = &self.slots[index];
        Ok(ReadyContext {
            id,
            entry: slot.entry,
            argument: slot.argument,
            initial_sp: slot.layout.usable_top(),
        })
    }
    /// Called after a real switch; actual SP must be within this retained stack.
    pub fn activate(&mut self, id: ThreadId) -> Result<(), Error> {
        let context = self.context()?;
        let index = self.index(id)?;
        if self.slots[index].state != State::Ready
            || self.slots.iter().any(|s| s.state == State::Running)
        {
            return Err(Error::State);
        }
        self.verify_slot(index, false)?;
        if !self.slots[index].layout.usable(context.sp) {
            return Err(Error::ActiveStack);
        }
        self.slots[index].state = State::Running;
        Ok(())
    }
    fn current(&self, id: ThreadId) -> Result<usize, Error> {
        let context = self.context()?;
        let index = self.index(id)?;
        if self.slots[index].state != State::Running {
            return Err(Error::State);
        }
        self.verify_slot(index, false)?;
        if !self.slots[index].layout.usable(context.sp) {
            return Err(Error::ActiveStack);
        }
        Ok(index)
    }
    pub fn yield_current(&mut self, id: ThreadId) -> Result<(), Error> {
        let index = self.current(id)?;
        self.slots[index].state = State::Ready;
        Ok(())
    }
    pub fn sleep_current(&mut self, id: ThreadId) -> Result<Wait, Error> {
        let index = self.current(id)?;
        let epoch = self.slots[index]
            .wait_epoch
            .checked_add(1)
            .ok_or(Error::Exhausted)?;
        self.slots[index].wait_epoch = epoch;
        self.slots[index].state = State::Sleeping;
        Ok(Wait { id, epoch })
    }
    pub fn wake(&mut self, wait: Wait) -> Result<ThreadId, (Error, Wait)> {
        if let Err(e) = self.context() {
            return Err((e, wait));
        }
        let index = match self.index(wait.id) {
            Ok(v) => v,
            Err(e) => return Err((e, wait)),
        };
        if self.slots[index].state != State::Sleeping || self.slots[index].wait_epoch != wait.epoch
        {
            return Err((Error::State, wait));
        }
        if let Err(e) = self.verify_slot(index, false) {
            return Err((e, wait));
        }
        self.slots[index].state = State::Ready;
        Ok(wait.id)
    }
    pub fn exit_current(&mut self, id: ThreadId) -> Result<(), Error> {
        let index = self.current(id)?;
        self.slots[index].state = State::Dead;
        Ok(())
    }
    /// Retained Ready or Sleeping threads may be cancelled by the task caller.
    pub fn cancel(&mut self, id: ThreadId) -> Result<(), Error> {
        self.context()?;
        let index = self.index(id)?;
        if !matches!(self.slots[index].state, State::Ready | State::Sleeping) {
            return Err(Error::State);
        }
        self.verify_slot(index, false)?;
        self.slots[index].state = State::Dead;
        Ok(())
    }
    pub fn reap(&mut self, id: ThreadId) -> Result<(), Error> {
        let context = self.context()?;
        let index = self.index(id)?;
        if self.slots[index].state == State::Quarantined {
            return Err(Error::Quarantined);
        }
        if self.slots[index].state != State::Dead {
            return Err(Error::State);
        }
        if self.slots[index].layout.contains(context.sp) {
            return Err(Error::ActiveStack);
        }
        self.verify_slot(index, false)?;
        let stack = self.slots[index].stack.take().ok_or(Error::State)?;
        if let Err((_error, stack)) = self.backend.release(stack) {
            self.slots[index].stack = Some(stack);
            self.slots[index].state = State::Quarantined;
            return Err(Error::Backend);
        }
        let generation = self.slots[index].generation;
        self.slots[index] = Slot::empty();
        let _ = generation;
        Ok(())
    }
    pub fn snapshot(&self) -> Snapshot {
        let mut result = Snapshot::default();
        for slot in &self.slots {
            if slot.state == State::Vacant {
                result.retired += usize::from(slot.generation == u64::MAX);
                continue;
            }
            result.live += 1;
            result.pages += slot.layout.bytes / PAGE;
            match slot.state {
                State::Ready => result.ready += 1,
                State::Running => result.running += 1,
                State::Sleeping => result.sleeping += 1,
                State::Dead => result.dead += 1,
                State::Quarantined => result.quarantined += 1,
                State::Vacant => {}
            }
        }
        result
    }
    pub fn into_backend(self) -> Result<B, (Error, Self)> {
        if let Err(e) = self.context() {
            return Err((e, self));
        }
        if self.slots.iter().any(|s| s.state != State::Vacant) {
            return Err((Error::Busy, self));
        }
        Ok(self.backend)
    }
}

#[cfg(test)]
mod exhaustion_tests {
    use super::*;
    struct Stack(u64);
    struct TestBackend {
        sp: u64,
        prepares: usize,
    }
    unsafe impl Backend for TestBackend {
        type Stack = Stack;
        fn execution(&self) -> Execution {
            Execution {
                cpu: 0,
                irq_masked: true,
                in_irq: false,
                sp: self.sp,
            }
        }
        fn layout(&self, stack: &Stack) -> Result<StackLayout, Error> {
            Ok(StackLayout {
                base: 0x80001000 + stack.0 * 0x20000,
                physical: 0x41000000 + stack.0 * 0x20000,
                bytes: 4 * PAGE,
            })
        }
        fn entry_executable(&self, _: u64) -> bool {
            true
        }
        fn prepare(&mut self, _: &mut Stack) -> Result<(), Error> {
            self.prepares += 1;
            Ok(())
        }
        fn observe(&self, stack: &Stack) -> Result<StackObservation, Error> {
            Ok(StackObservation {
                layout: self.layout(stack)?,
                guard_descriptor: 0,
                low_canary: LOW_CANARY,
                high_canary: HIGH_CANARY,
                initially_zeroed: true,
            })
        }
        fn release(&mut self, _: Stack) -> Result<(), (Error, Stack)> {
            Ok(())
        }
    }
    fn table<const N: usize>() -> Table<TestBackend, N> {
        Table::new(TestBackend { sp: 0, prepares: 0 }).ok().unwrap()
    }
    #[test]
    fn retired_generation_is_skipped_while_other_slot_progresses() {
        let mut t = table::<2>();
        t.slots[0].generation = u64::MAX;
        let id = t.create(Stack(1), 4, 0).ok().unwrap();
        assert_eq!(id.slot(), 1);
        assert_eq!(id.generation(), 1);
        assert_eq!(t.snapshot().retired, 1);
        assert_eq!(t.backend.prepares, 1);
    }
    #[test]
    fn all_retired_generations_refuse_before_preparing_candidate() {
        let mut t = table::<1>();
        t.slots[0].generation = u64::MAX;
        let (error, stack) = t.create(Stack(9), 4, 0).err().unwrap();
        assert_eq!(error, Error::Capacity);
        assert_eq!(stack.0, 9);
        assert_eq!(t.backend.prepares, 0);
        assert_eq!(t.snapshot().retired, 1);
    }
    #[test]
    fn sleep_epoch_exhaustion_keeps_running_thread_and_backing() {
        let mut t = table::<1>();
        let id = t.create(Stack(0), 4, 0).ok().unwrap();
        t.backend.sp = t.ready_context(id).unwrap().initial_sp;
        t.activate(id).unwrap();
        t.slots[0].wait_epoch = u64::MAX;
        let before = t.snapshot();
        assert_eq!(t.sleep_current(id).err(), Some(Error::Exhausted));
        assert_eq!(t.snapshot(), before);
        assert!(t.slots[0].stack.is_some());
        assert_eq!(t.slots[0].wait_epoch, u64::MAX);
    }
    #[test]
    fn stale_sleep_epoch_cannot_wake_matching_thread_generation() {
        let mut t = table::<1>();
        let id = t.create(Stack(0), 4, 0).ok().unwrap();
        t.backend.sp = t.ready_context(id).unwrap().initial_sp;
        t.activate(id).unwrap();
        let current = t.sleep_current(id).unwrap();
        t.backend.sp = 0;
        let stale = Wait {
            id,
            epoch: current.epoch - 1,
        };
        let before = t.snapshot();
        assert_eq!(t.wake(stale).err().unwrap().0, Error::State);
        assert_eq!(t.snapshot(), before);
        t.wake(current).unwrap();
        assert_eq!(t.snapshot().ready, 1);
    }
}