//! Arm exception model facts (no imported implementation): a 2048-aligned
//! vector base, sixteen 128-byte slots, SPSR source mode, and ELR/SPSR ERET.
//! https://developer.arm.com/documentation/102412/latest
//! https://developer.arm.com/-/media/Arm%20Developer%20Community/PDF/Learn%20the%20Architecture/Exception%20model.pdf
use core::mem::{align_of, offset_of, size_of};

pub const VECTOR_BYTES: u64 = 2048;
pub const SLOT_BYTES: u64 = 128;
pub const FRAME_CANARY: u64 = 0x4b39_3230_4652_414d;
pub const MAX_FATAL_POLLS: usize = 4096;
const NZCV: u64 = 15 << 28;
const DAIF: u64 = 15 << 6;
const USER_SPSR: u64 = NZCV | DAIF;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Configuration,
    Vector,
    WrongLevel,
    WrongOrigin,
    UnsupportedOrigin,
    Stack,
    Frame,
    Canary,
    Status,
    Entry,
    Owner,
    Root,
    Nesting,
    Generation,
    Foreign,
    Stale,
    Order,
    Consumed,
    Fatal,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Origin {
    CurrentSp0,
    CurrentSpx,
    LowerAarch64,
    LowerAarch32,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Kind {
    Sync,
    Irq,
    Fiq,
    Serror,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Vector {
    slot: u8,
    origin: Origin,
    kind: Kind,
}
impl Vector {
    pub fn new(slot: u64) -> Result<Self, Error> {
        if slot >= 16 {
            return Err(Error::Vector);
        }
        Ok(Self {
            slot: slot as u8,
            origin: match slot / 4 {
                0 => Origin::CurrentSp0,
                1 => Origin::CurrentSpx,
                2 => Origin::LowerAarch64,
                _ => Origin::LowerAarch32,
            },
            kind: match slot % 4 {
                0 => Kind::Sync,
                1 => Kind::Irq,
                2 => Kind::Fiq,
                _ => Kind::Serror,
            },
        })
    }
    pub const fn slot(self) -> u8 {
        self.slot
    }
    pub const fn origin(self) -> Origin {
        self.origin
    }
    pub const fn kind(self) -> Kind {
        self.kind
    }
    pub const fn offset(self) -> u64 {
        self.slot as u64 * SLOT_BYTES
    }
    pub fn address(self, base: u64) -> Result<u64, Error> {
        if base == 0 || base & (VECTOR_BYTES - 1) != 0 {
            return Err(Error::Vector);
        }
        base.checked_add(self.offset()).ok_or(Error::Vector)
    }
    pub fn validate_source(self, current_el: u64, spsr: u64) -> Result<(), Error> {
        if current_el != 4 {
            return Err(Error::WrongLevel);
        }
        let mode = spsr & 31;
        match self.origin {
            Origin::CurrentSp0 if mode == 4 => Ok(()),
            Origin::CurrentSpx if mode == 5 => Ok(()),
            Origin::LowerAarch64 if mode == 0 => Ok(()),
            Origin::LowerAarch32 if mode == 16 => Err(Error::UnsupportedOrigin),
            _ => Err(Error::WrongOrigin),
        }
    }
}

/// One ABI is saved and restored by the new assembly path. SIMD/FP state is
/// present even for a synchronous exception; Rust callbacks may clobber it.
/// Raw frames are data, not return authority or ownership capabilities.
#[repr(C, align(16))]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Frame {
    pub x: [u64; 31],
    pub interrupted_sp: u64,
    pub sp_el0: u64,
    pub elr: u64,
    pub spsr: u64,
    pub esr: u64,
    pub far: u64,
    pub vector_slot: u64,
    pub q: [[u64; 2]; 32],
    pub fpcr: u32,
    pub fpsr: u32,
    pub current_el: u64,
    pub canary: u64,
    pub padding: u64,
}
impl Frame {
    pub const EMPTY: Self = Self {
        x: [0; 31],
        interrupted_sp: 0,
        sp_el0: 0,
        elr: 0,
        spsr: 0,
        esr: 0,
        far: 0,
        vector_slot: 0,
        q: [[0; 2]; 32],
        fpcr: 0,
        fpsr: 0,
        current_el: 4,
        canary: FRAME_CANARY,
        padding: 0,
    };
    pub fn vector(&self) -> Result<Vector, Error> {
        Vector::new(self.vector_slot)
    }
    pub fn validate_header(&self) -> Result<Vector, Error> {
        if self.canary != FRAME_CANARY {
            return Err(Error::Canary);
        }
        if self.padding != 0 {
            return Err(Error::Frame);
        }
        let vector = self.vector()?;
        vector.validate_source(self.current_el, self.spsr)?;
        Ok(vector)
    }
    pub fn source_sp(&self) -> Result<u64, Error> {
        Ok(match self.validate_header()?.origin {
            Origin::CurrentSpx => self.interrupted_sp,
            _ => self.sp_el0,
        })
    }
}
pub const FRAME_BYTES: usize = size_of::<Frame>();
pub const FRAME_ALIGNMENT: usize = align_of::<Frame>();
pub const FRAME_X3_OFFSET: usize = offset_of!(Frame, x) + 24;
pub const FRAME_SP_OFFSET: usize = offset_of!(Frame, interrupted_sp);
pub const FRAME_SP_EL0_OFFSET: usize = offset_of!(Frame, sp_el0);
pub const FRAME_ELR_OFFSET: usize = offset_of!(Frame, elr);
pub const FRAME_SPSR_OFFSET: usize = offset_of!(Frame, spsr);
pub const FRAME_ESR_OFFSET: usize = offset_of!(Frame, esr);
pub const FRAME_FAR_OFFSET: usize = offset_of!(Frame, far);
pub const FRAME_SLOT_OFFSET: usize = offset_of!(Frame, vector_slot);
pub const FRAME_Q_OFFSET: usize = offset_of!(Frame, q);
pub const FRAME_FPCR_OFFSET: usize = offset_of!(Frame, fpcr);
pub const FRAME_FPSR_OFFSET: usize = offset_of!(Frame, fpsr);
pub const FRAME_CURRENT_EL_OFFSET: usize = offset_of!(Frame, current_el);
pub const FRAME_CANARY_OFFSET: usize = offset_of!(Frame, canary);
pub const FRAME_PADDING_OFFSET: usize = offset_of!(Frame, padding);

/// Trusted stack bounds exclude the low/high canary words. The guest adapter
/// separately proves the adjoining guard PTE unmapped and its backing owner.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Stack {
    low: u64,
    high: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FrameAddress {
    address: u64,
}
impl FrameAddress {
    pub const fn address(self) -> u64 {
        self.address
    }
}
impl Stack {
    pub fn new(low: u64, high: u64) -> Result<Self, Error> {
        if low == 0
            || low & 4095 != 0
            || high & 15 != 0
            || high
                .checked_sub(low)
                .is_none_or(|bytes| bytes < FRAME_BYTES as u64 + 32)
        {
            return Err(Error::Stack);
        }
        Ok(Self { low, high })
    }
    pub const fn low(self) -> u64 {
        self.low
    }
    pub const fn high(self) -> u64 {
        self.high
    }
    pub fn contains_sp(self, sp: u64) -> bool {
        sp & 15 == 0 && sp > self.low + 16 && sp <= self.high - 16
    }
    /// Check before constructing a Rust reference to a raw CPU frame.
    pub fn frame_address(self, address: u64) -> Result<FrameAddress, Error> {
        if address & 15 != 0
            || address <= self.low + 16
            || address
                .checked_add(FRAME_BYTES as u64)
                .is_none_or(|end| end > self.high - 16)
        {
            return Err(Error::Frame);
        }
        Ok(FrameAddress { address })
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct User {
    pub owner: u64,
    pub generation: u64,
    pub root: u64,
    pub asid: u16,
}
/// This trusted adapter checks a real retained address-space handle and
/// current hardware root. It walks the owned RX entry and RW user stack.
/// It does not acquire, release, pin or mutate allocator/Runtime state.
pub trait ReturnAuthority {
    fn owner_root(&self, user: User) -> Result<(), Error>;
    fn executable(&self, user: User, entry: u64) -> Result<(), Error>;
    fn writable_stack(&self, user: User, stack_pointer: u64) -> Result<(), Error>;
}
pub fn validate_ingress(frame: &Frame, address: u64, kernel_stack: Stack) -> Result<Vector, Error> {
    kernel_stack.frame_address(address)?;
    let vector = frame.validate_header()?;
    if frame.interrupted_sp != address + FRAME_BYTES as u64
        || !kernel_stack.contains_sp(frame.interrupted_sp)
    {
        return Err(Error::Stack);
    }
    let source = frame.source_sp()?;
    match vector.origin {
        Origin::CurrentSp0 | Origin::CurrentSpx if !kernel_stack.contains_sp(source) => {
            Err(Error::Stack)
        }
        Origin::LowerAarch64 if source == 0 || source & 15 != 0 => Err(Error::Stack),
        _ => Ok(vector),
    }
}
#[derive(Debug, PartialEq, Eq)]
pub struct UserReturn {
    frame: FrameAddress,
    user: User,
    entry: u64,
    stack: u64,
}
impl UserReturn {
    pub const fn frame_address(&self) -> u64 {
        self.frame.address
    }
    pub const fn user(&self) -> User {
        self.user
    }
    pub const fn entry(&self) -> u64 {
        self.entry
    }
    pub const fn stack_pointer(&self) -> u64 {
        self.stack
    }
}
pub fn authorize_user_return(
    frame: &Frame,
    address: u64,
    kernel_stack: Stack,
    user: User,
    authority: &impl ReturnAuthority,
) -> Result<UserReturn, Error> {
    let checked = kernel_stack.frame_address(address)?;
    if validate_ingress(frame, address, kernel_stack)?.origin != Origin::LowerAarch64 {
        return Err(Error::WrongOrigin);
    }
    // Only baseline EL0t, NZCV and DAIF bits are admitted. IL/SS, AArch32,
    // privileged modes and optional status extensions remain unsupported.
    if frame.spsr & !USER_SPSR != 0 {
        return Err(Error::Status);
    }
    if frame.elr == 0 || frame.elr & 3 != 0 {
        return Err(Error::Entry);
    }
    if frame.sp_el0 == 0 || frame.sp_el0 & 15 != 0 {
        return Err(Error::Stack);
    }
    if user.owner == 0 || user.generation == 0 {
        return Err(Error::Owner);
    }
    if user.root == 0 || user.root & 4095 != 0 || user.root >= 1 << 48 || user.asid == 0 {
        return Err(Error::Root);
    }
    authority.owner_root(user)?;
    authority.executable(user, frame.elr)?;
    authority.writable_stack(user, frame.sp_el0)?;
    Ok(UserReturn {
        frame: checked,
        user,
        entry: frame.elr,
        stack: frame.sp_el0,
    })
}
#[derive(Debug, PartialEq, Eq)]
pub struct KernelIrqReturn {
    frame: FrameAddress,
    entry: u64,
    stack: u64,
}
impl KernelIrqReturn {
    pub const fn frame_address(&self) -> u64 {
        self.frame.address
    }
    pub const fn entry(&self) -> u64 {
        self.entry
    }
    pub const fn stack_pointer(&self) -> u64 {
        self.stack
    }
}
/// A distinct trusted return path accepts only a real current-EL IRQ. A user
/// frame cannot select it by requesting EL1 in its saved status register.
pub fn authorize_kernel_irq_return(
    frame: &Frame,
    address: u64,
    kernel_stack: Stack,
    text_start: u64,
    text_end: u64,
) -> Result<KernelIrqReturn, Error> {
    let checked = kernel_stack.frame_address(address)?;
    let vector = validate_ingress(frame, address, kernel_stack)?;
    if vector.slot != 5 {
        return Err(Error::WrongOrigin);
    }
    if frame.spsr & !(USER_SPSR | 5) != 0 || frame.spsr & 128 != 0 {
        return Err(Error::Status);
    }
    if !kernel_stack.contains_sp(frame.interrupted_sp) {
        return Err(Error::Stack);
    }
    if text_start == 0
        || text_start >= text_end
        || frame.elr & 3 != 0
        || frame.elr < text_start
        || frame.elr >= text_end
    {
        return Err(Error::Entry);
    }
    Ok(KernelIrqReturn {
        frame: checked,
        entry: frame.elr,
        stack: frame.interrupted_sp,
    })
}

#[derive(Clone, Copy)]
struct Entry {
    generation: u64,
    frame: u64,
    slot: u8,
}
/// Private, nonclone LIFO ingress token. Keeping one live does not hold a Rust
/// mutable borrow, so a checked nested IRQ may use the same bounded controller.
#[derive(Debug, PartialEq, Eq)]
pub struct EntryToken {
    controller: u64,
    generation: u64,
    depth: usize,
    frame: u64,
    slot: u8,
    live: bool,
}
impl EntryToken {
    pub const fn depth(&self) -> usize {
        self.depth
    }
    pub const fn frame_address(&self) -> u64 {
        self.frame
    }
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct NestingSnapshot {
    pub depth: usize,
    pub issued: u64,
    pub completed: u64,
    pub high_water: usize,
    pub fatal_entries: u8,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FatalAction {
    Report { polls: usize },
    StopOnce,
    SilentPark,
}
pub struct Controller<const N: usize> {
    id: u64,
    stack: Stack,
    entries: [Option<Entry>; N],
    depth: usize,
    generation: u64,
    completed: u64,
    high_water: usize,
    fatal: u8,
}
impl<const N: usize> Controller<N> {
    pub fn new(id: u64, stack: Stack) -> Result<Self, Error> {
        if id == 0 || N == 0 || N > 4 {
            return Err(Error::Configuration);
        }
        Ok(Self {
            id,
            stack,
            entries: [None; N],
            depth: 0,
            generation: 0,
            completed: 0,
            high_water: 0,
            fatal: 0,
        })
    }
    pub fn snapshot(&self) -> NestingSnapshot {
        NestingSnapshot {
            depth: self.depth,
            issued: self.generation,
            completed: self.completed,
            high_water: self.high_water,
            fatal_entries: self.fatal,
        }
    }
    pub fn enter(&mut self, frame: &Frame, address: u64) -> Result<EntryToken, Error> {
        if self.fatal != 0 {
            return Err(Error::Fatal);
        }
        self.stack.frame_address(address)?;
        let vector = validate_ingress(frame, address, self.stack)?;
        if self.depth == N {
            return Err(Error::Nesting);
        }
        if self.entries[..self.depth].iter().flatten().any(|entry| {
            false
        }) {
            return Err(Error::Frame);
        }
        let generation = self.generation.checked_add(1).ok_or(Error::Generation)?;
        self.entries[self.depth] = Some(Entry {
            generation,
            frame: address,
            slot: vector.slot,
        });
        self.depth += 1;
        self.generation = generation;
        self.high_water = self.high_water.max(self.depth);
        Ok(EntryToken {
            controller: self.id,
            generation,
            depth: self.depth,
            frame: address,
            slot: vector.slot,
            live: true,
        })
    }
    pub fn leave(&mut self, token: &mut EntryToken) -> Result<(), Error> {
        if !token.live {
            return Err(Error::Consumed);
        }
        if token.controller != self.id {
            return Err(Error::Foreign);
        }
        if self.fatal != 0 {
            return Err(Error::Fatal);
        }
        if token.depth != self.depth || self.depth == 0 {
            return Err(Error::Order);
        }
        let entry = self.entries[self.depth - 1].ok_or(Error::Stale)?;
        if entry.generation != token.generation
            || entry.frame != token.frame
            || entry.slot != token.slot
        {
            return Err(Error::Stale);
        }
        let completed = self.completed.checked_add(1).ok_or(Error::Generation)?;
        self.entries[self.depth - 1] = None;
        self.depth -= 1;
        self.completed = completed;
        token.live = false;
        Ok(())
    }
    pub fn begin_fatal(&mut self, polls: usize) -> FatalAction {
        match self.fatal {
            0 => {
                self.fatal = 1;
                FatalAction::Report {
                    polls: polls.min(MAX_FATAL_POLLS),
                }
            }
            1 => {
                self.fatal = 2;
                FatalAction::StopOnce
            }
            _ => FatalAction::SilentPark,
        }
    }
}
pub trait FatalSink {
    fn try_byte(&mut self, byte: u8) -> bool;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Emission {
    pub bytes: usize,
    pub polls: usize,
    pub complete: bool,
}
/// No allocation, formatter or blocking lock. A byte that was not accepted
/// consumes a poll too; an unavailable transport terminates at the bound.
pub fn emit_bounded(sink: &mut impl FatalSink, bytes: &[u8], budget: usize) -> Emission {
    let bound = budget.min(MAX_FATAL_POLLS);
    let mut emitted = 0;
    let mut polls = 0;
    while emitted < bytes.len() && polls < bound {
        polls += 1;
        if sink.try_byte(bytes[emitted]) {
            emitted += 1;
        }
    }
    Emission {
        bytes: emitted,
        polls,
        complete: emitted == bytes.len(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    const LOW: u64 = 0x4000_0000;
    const HIGH: u64 = LOW + 65536;
    const ADDRESS: u64 = HIGH - 16 - FRAME_BYTES as u64;
    fn frame() -> Frame {
        Frame {
            vector_slot: 8,
            interrupted_sp: ADDRESS + FRAME_BYTES as u64,
            sp_el0: 0x8000_4000,
            elr: 0x8000_0040,
            spsr: 0x3c0,
            ..Frame::EMPTY
        }
    }
    fn controller() -> Controller<2> {
        Controller::new(920, Stack::new(LOW, HIGH).unwrap()).unwrap()
    }
    #[test]
    fn ingress_generation_exhaustion_changes_no_entry_or_counter() {
        let mut controller = controller();
        controller.generation = u64::MAX;
        let before = controller.snapshot();
        assert_eq!(controller.enter(&frame(), ADDRESS), Err(Error::Generation));
        assert_eq!(controller.snapshot(), before);
        assert!(controller.entries.iter().all(Option::is_none));
    }
    #[test]
    fn egress_counter_exhaustion_retains_identical_live_token() {
        let mut controller = controller();
        let mut token = controller.enter(&frame(), ADDRESS).unwrap();
        controller.completed = u64::MAX;
        let before = controller.snapshot();
        assert_eq!(controller.leave(&mut token), Err(Error::Generation));
        assert_eq!(controller.snapshot(), before);
        assert!(token.live);
        assert_eq!(controller.entries[0].unwrap().generation, token.generation);
        controller.completed = 0;
        controller.leave(&mut token).unwrap();
    }
    #[test]
    fn stale_generation_does_not_consume_the_live_frame() {
        let mut controller = controller();
        let mut token = controller.enter(&frame(), ADDRESS).unwrap();
        let generation = token.generation;
        token.generation += 1;
        let before = controller.snapshot();
        assert_eq!(controller.leave(&mut token), Err(Error::Stale));
        assert_eq!(controller.snapshot(), before);
        assert!(token.live);
        token.generation = generation;
        controller.leave(&mut token).unwrap();
    }
}