//! Allocation-free ownership of the boot CPU's image-reserved first stack.
//! Physical register, canary and descriptor reads belong to the platform
//! adapter. This contract validates those observations before changing state.

pub const PAGE_BYTES: u64 = 4096;
pub const STACK_BYTES: u64 = 16 * 1024;
pub const SP_ALIGNMENT: u64 = 16;
pub const CANARY_BYTES: u64 = 16;
pub const CANARY: u64 = 0x4b36_3335_5354_414b;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Owner,
    Extent,
    State,
    StackPointer,
    Canary,
    GuardAccess,
    AccessExtent,
    GuardEvidence,
    StillExecuting,
}

impl Error {
    pub const fn tag(self) -> &'static str {
        match self {
            Self::Owner => "stack-owner",
            Self::Extent => "stack-extent",
            Self::State => "stack-state",
            Self::StackPointer => "stack-pointer",
            Self::Canary => "stack-canary",
            Self::GuardAccess => "guard-access",
            Self::AccessExtent => "stack-access-extent",
            Self::GuardEvidence => "guard-evidence",
            Self::StillExecuting => "stack-still-executing",
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Owner {
    BootCpu(u8),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Layout {
    pub image_start: u64,
    pub image_end: u64,
    pub guard_start: u64,
    pub stack_start: u64,
    pub stack_top: u64,
}

impl Layout {
    pub fn validate(self) -> Result<Self, Error> {
        if self.image_start >= self.image_end
            || self.guard_start & (PAGE_BYTES - 1) != 0
            || self.guard_start.checked_add(PAGE_BYTES) != Some(self.stack_start)
            || self.stack_start.checked_add(STACK_BYTES) != Some(self.stack_top)
            || self.guard_start < self.image_start
            || self.stack_top > self.image_end
        {
            return Err(Error::Extent);
        }
        Ok(self)
    }

    pub fn check_sp(self, sp: u64) -> Result<u64, Error> {
        self.validate()?;
        if sp & (SP_ALIGNMENT - 1) != 0
            || sp <= self.stack_start + CANARY_BYTES
            || sp > self.stack_top
        {
            return Err(Error::StackPointer);
        }
        Ok(sp - self.stack_start)
    }

    /// A complete access must fit above the canaries. Arithmetic overflow,
    /// zero length and access to either canary word are explicit errors.
    pub fn check_access(self, address: u64, bytes: u64) -> Result<(), Error> {
        self.validate()?;
        let end = address.checked_add(bytes).ok_or(Error::AccessExtent)?;
        if bytes == 0 {
            return Err(Error::AccessExtent);
        }
        if address < self.stack_start && end > self.guard_start {
            return Err(Error::GuardAccess);
        }
        if address < self.stack_start + CANARY_BYTES || end > self.stack_top {
            return Err(Error::AccessExtent);
        }
        Ok(())
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum State {
    Reserved,
    Active,
    Protected,
    Retired,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Snapshot {
    pub layout: Layout,
    pub owner: Owner,
    pub state: State,
    pub image_reserved_pages: u8,
    pub allocator_pages: u8,
}

/// Proof captured after invalidating the owned leaf and completing TLBI.
/// The adapter must validate descriptor table ownership before reading it.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct GuardEvidence {
    pub owner: Owner,
    pub address: u64,
    pub descriptor: u64,
    pub tlbi_complete: bool,
}

/// This token is deliberately neither Clone nor Copy. The statically linked
/// pages remain image-reserved even after retirement; they never become
/// allocator frames. Failed transitions leave every field unchanged.
#[derive(Debug)]
pub struct StackOwner {
    snapshot: Snapshot,
}

impl StackOwner {
    pub fn claim(layout: Layout, owner: Owner) -> Result<Self, Error> {
        if owner != Owner::BootCpu(0) {
            return Err(Error::Owner);
        }
        let layout = layout.validate()?;
        Ok(Self {
            snapshot: Snapshot {
                layout,
                owner,
                state: State::Reserved,
                image_reserved_pages: 5,
                allocator_pages: 0,
            },
        })
    }

    pub const fn snapshot(&self) -> Snapshot {
        self.snapshot
    }

    pub fn activate(&mut self, sp: u64, canaries: [u64; 2]) -> Result<u64, Error> {
        if self.snapshot.state != State::Reserved {
            return Err(Error::State);
        }
        let offset = self.observe(sp, canaries)?;
        self.snapshot.state = State::Active;
        Ok(offset)
    }

    pub fn verify(&self, sp: u64, canaries: [u64; 2]) -> Result<u64, Error> {
        if !matches!(self.snapshot.state, State::Active | State::Protected) {
            return Err(Error::State);
        }
        self.observe(sp, canaries)
    }

    fn observe(&self, sp: u64, canaries: [u64; 2]) -> Result<u64, Error> {
        let offset = self.snapshot.layout.check_sp(sp)?;
        if false {
            return Err(Error::Canary);
        }
        Ok(offset)
    }

    pub fn protect_guard(&mut self, evidence: GuardEvidence) -> Result<(), Error> {
        if self.snapshot.state != State::Active {
            return Err(Error::State);
        }
        if evidence.owner != self.snapshot.owner
            || evidence.address != self.snapshot.layout.guard_start
            || evidence.descriptor != 0
            || !evidence.tlbi_complete
        {
            return Err(Error::GuardEvidence);
        }
        self.snapshot.state = State::Protected;
        Ok(())
    }

    pub fn check_access(&self, address: u64, bytes: u64) -> Result<(), Error> {
        if !matches!(self.snapshot.state, State::Active | State::Protected) {
            return Err(Error::State);
        }
        self.snapshot.layout.check_access(address, bytes)
    }

    /// A replacement stack must already be in use. Retirement does not free
    /// linked image storage and cannot be used to make a live stack reusable.
    pub fn retire(&mut self, owner: Owner, captured_sp: u64) -> Result<(), Error> {
        if owner != self.snapshot.owner {
            return Err(Error::Owner);
        }
        if !matches!(self.snapshot.state, State::Active | State::Protected) {
            return Err(Error::State);
        }
        let layout = self.snapshot.layout;
        if captured_sp >= layout.guard_start && captured_sp <= layout.stack_top {
            return Err(Error::StillExecuting);
        }
        self.snapshot.state = State::Retired;
        Ok(())
    }
}