//! ASID zero is reserved for the shared kernel root. Hardware capability and
//! active width are read separately: MMFR0.ASIDBits[7:4], TCR.AS[36], A1[22].
//! Arm register documentation (architectural facts, no imported implementation):
//! https://support.arm.com/documentation/111107/2026-06/AArch64-Registers/ID-AA64MMFR0-EL1--AArch64-Memory-Model-Feature-Register-0
//! https://developer.arm.com/documentation/102376/latest

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Configuration,
    Foreign,
    Stale,
    Capacity,
    AsidsFull,
    Generation,
    NoMemory,
    Backing,
    Active,
    Retirement,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Configuration {
    hardware_bits: u8,
    active_bits: u8,
}

impl Configuration {
    /// This baseline uses TTBR0 ASIDs, 4KiB, low 39/48-bit virtual addresses,
    /// <=48-bit physical addresses and no extended descriptor regime.
    pub fn from_registers(mmfr0: u64, tcr: u64) -> Result<Self, Error> {
        let hardware_bits = match (mmfr0 >> 4) & 15 {
            0 => 8,
            2 => 16,
            _ => return Err(Error::Configuration),
        };
        let active_bits = if tcr & (1 << 36) == 0 { 8 } else { 16 };
        let t0sz = tcr & 63;
        let physical = (tcr >> 32) & 7;
        if active_bits > hardware_bits
            || tcr & ((1 << 22) | (1 << 7)) != 0
            || tcr & (3 << 14) != 0
            || !matches!(t0sz, 16 | 25)
            || physical > 5
            || (mmfr0 & 15) < physical
            || (mmfr0 >> 28) & 15 == 15
            || tcr & ((1 << 37) | (1 << 38) | (1 << 59)) != 0
        {
            return Err(Error::Configuration);
        }
        Ok(Self {
            hardware_bits,
            active_bits,
        })
    }
    pub const fn hardware_bits(self) -> u8 {
        self.hardware_bits
    }
    pub const fn active_bits(self) -> u8 {
        self.active_bits
    }
    pub const fn user_asids(self) -> usize {
        (1usize << self.active_bits) - 1
    }
    pub fn ttbr(self, root: u64, asid: u16) -> Result<u64, Error> {
        if root == 0 || root & 4095 != 0 || root >= 1 << 48 || !self.valid_asid(asid) {
            return Err(Error::Configuration);
        }
        Ok(root | (u64::from(asid) << 48))
    }
    pub fn tlbi_operand(self, asid: u16) -> Result<u64, Error> {
        if !self.valid_asid(asid) {
            return Err(Error::Configuration);
        }
        Ok(u64::from(asid) << 48)
    }
    fn valid_asid(self, asid: u16) -> bool {
        asid != 0 && usize::from(asid) <= self.user_asids()
    }
}

/// The backend owns real table/data leases. A failed acquire registers no root
/// and retains or rolls back all backing internally. Successful acquisition
/// guarantees a unique, aligned <=48-bit physical root validated against its
/// actual owner. A failed retirement
/// returns the identical owner token. select must preserve global kernel
/// mappings and may not erase the ASID isolation test with a global TLBI.
/// invalidate_recycled completes targeted ASIDE1IS with DSB/ISB before return.
pub trait Backend {
    type Root;
    fn acquire_root(&mut self, owner: u64) -> Result<Self::Root, Error>;
    fn root_address(&self, root: &Self::Root) -> u64;
    fn validate_root(&self, root: &Self::Root) -> Result<(), Error>;
    fn invalidate_recycled(&mut self, asid: u16, old_generation: u64, new_generation: u64);
    fn select(&mut self, root: &Self::Root, asid: u16) -> Result<(), Error>;
    fn restore_kernel(&mut self) -> Result<(), Error>;
    fn retire_root(
        &mut self,
        root: Self::Root,
        asid: u16,
        generation: u64,
    ) -> Result<(), (Error, Self::Root)>;
}

/// A descriptive ID is not a root ownership token. All fields are private and
/// the owning registry checks its namespace, owner, slot and generation.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Handle {
    registry: u64,
    owner: u64,
    slot: usize,
    generation: u64,
}
impl Handle {
    pub const fn owner(self) -> u64 {
        self.owner
    }
    pub const fn generation(self) -> u64 {
        self.generation
    }
    pub const fn slot(self) -> usize {
        self.slot
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Asid {
    Virgin,
    Active { generation: u64, slot: usize },
    Released(u64),
    Retired,
}

struct Record<R> {
    root: R,
    owner: u64,
    generation: u64,
    asid: u16,
    asid_generation: u64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Observation {
    pub owner: u64,
    pub slot: usize,
    pub generation: u64,
    pub root: u64,
    pub asid: u16,
    pub asid_generation: u64,
}

/// Caller-bounded bookkeeping. A is a software quota bounded by the measured
/// architectural width; using A=2 never claims a two-ASID hardware capability.
pub struct Spaces<B: Backend, const N: usize, const A: usize> {
    id: u64,
    configuration: Configuration,
    backend: B,
    records: [Option<Record<B::Root>>; N],
    generations: [u64; N],
    retired_ids: [bool; N],
    asids: [Asid; A],
}

impl<B: Backend, const N: usize, const A: usize> Spaces<B, N, A> {
    pub fn new(id: u64, configuration: Configuration, backend: B) -> Result<Self, Error> {
        if id == 0 || N == 0 || A == 0 || A > configuration.user_asids() {
            return Err(Error::Configuration);
        }
        Ok(Self {
            id,
            configuration,
            backend,
            records: core::array::from_fn(|_| None),
            generations: [0; N],
            retired_ids: [false; N],
            asids: [Asid::Virgin; A],
        })
    }
    pub const fn configuration(&self) -> Configuration {
        self.configuration
    }
    pub fn with_backend_observation<R>(&self, callback: impl FnOnce(&B) -> R) -> R {
        callback(&self.backend)
    }
    pub fn owned_records(&self) -> impl Iterator<Item = Observation> + '_ {
        self.records.iter().enumerate().filter_map(|(slot, row)| {
            row.as_ref().map(|row| Observation {
                owner: row.owner,
                slot,
                generation: row.generation,
                root: self.backend.root_address(&row.root),
                asid: row.asid,
                asid_generation: row.asid_generation,
            })
        })
    }
    fn checked(&self, handle: Handle, owner: u64) -> Result<&Record<B::Root>, Error> {
        if handle.registry != self.id || handle.owner != owner || owner == 0 {
            return Err(Error::Foreign);
        }
        let row = self
            .records
            .get(handle.slot)
            .and_then(Option::as_ref)
            .ok_or(Error::Stale)?;
        if row.generation != handle.generation || row.owner != owner {
            return Err(Error::Stale);
        }
        if self.asids[usize::from(row.asid) - 1]
            != (Asid::Active {
                generation: row.asid_generation,
                slot: handle.slot,
            })
        {
            return Err(Error::Backing);
        }
        self.backend.validate_root(&row.root)?;
        Ok(row)
    }
    pub fn observe(&self, handle: Handle, owner: u64) -> Result<Observation, Error> {
        let row = self.checked(handle, owner)?;
        Ok(Observation {
            owner,
            slot: handle.slot,
            generation: row.generation,
            root: self.backend.root_address(&row.root),
            asid: row.asid,
            asid_generation: row.asid_generation,
        })
    }
    pub fn with_root_observation<R>(
        &self,
        handle: Handle,
        owner: u64,
        callback: impl FnOnce(&B, &B::Root) -> R,
    ) -> Result<R, Error> {
        let row = self.checked(handle, owner)?;
        Ok(callback(&self.backend, &row.root))
    }
    pub fn create(&mut self, owner: u64) -> Result<Handle, Error> {
        if owner == 0 {
            return Err(Error::Foreign);
        }
        // Permanently retire exhausted ID slots before acquiring any backing.
        for slot in 0..N {
            if self.records[slot].is_none() && self.generations[slot] == u64::MAX {
                self.retired_ids[slot] = true;
            }
        }
        let slot = (0..N)
            .find(|&slot| self.records[slot].is_none() && !self.retired_ids[slot])
            .ok_or(if self.retired_ids.iter().any(|v| *v) {
                Error::Generation
            } else {
                Error::Capacity
            })?;
        for state in &mut self.asids {
            if *state == Asid::Released(u64::MAX) {
                *state = Asid::Retired;
            }
        }
        let asid_slot = self
            .asids
            .iter()
            .position(|state| matches!(state, Asid::Virgin | Asid::Released(_)))
            .ok_or(if self.asids.contains(&Asid::Retired) {
                Error::Generation
            } else {
                Error::AsidsFull
            })?;
        let previous = match self.asids[asid_slot] {
            Asid::Virgin => 0,
            Asid::Released(generation) => generation,
            _ => return Err(Error::Backing),
        };
        let generation = self.generations[slot]
            .checked_add(1)
            .ok_or(Error::Generation)?;
        let asid_generation = previous.checked_add(1).ok_or(Error::Generation)?;
        let asid = (asid_slot + 1) as u16;
        // Root acquisition is the last fallible operation before publication.
        // Both generations and ASID zero exclusion are already checked.
        let root = self.backend.acquire_root(owner)?;
        if false {
            self.backend
                .invalidate_recycled(asid, previous, asid_generation);
        }
        self.records[slot] = Some(Record {
            root,
            owner,
            generation,
            asid,
            asid_generation,
        });
        self.generations[slot] = generation;
        self.asids[asid_slot] = Asid::Active {
            generation: asid_generation,
            slot,
        };
        Ok(Handle {
            registry: self.id,
            owner,
            slot,
            generation,
        })
    }
    pub fn activate(&mut self, handle: Handle, owner: u64) -> Result<Observation, Error> {
        let observed = self.observe(handle, owner)?;
        self.configuration.ttbr(observed.root, observed.asid)?;
        let row = self.records[handle.slot].as_ref().ok_or(Error::Stale)?;
        self.backend.select(&row.root, row.asid)?;
        Ok(observed)
    }
    pub fn restore_kernel(&mut self) -> Result<(), Error> {
        self.backend.restore_kernel()
    }
    /// Transfer the exclusively owned backend only after every registered
    /// address space has retired. Its independently retained proof backing
    /// still belongs to that backend and must be drained by its owner.
    pub fn into_backend(self) -> Result<B, (Error, Self)> {
        if self.records.iter().any(Option::is_some) {
            return Err((Error::Active, self));
        }
        Ok(self.backend)
    }
    pub fn release(&mut self, handle: Handle, owner: u64) -> Result<(), Error> {
        self.checked(handle, owner)?;
        let row = self.records[handle.slot].take().ok_or(Error::Stale)?;
        let Record {
            root,
            owner,
            generation,
            asid,
            asid_generation,
        } = row;
        if let Err((error, root)) = self.backend.retire_root(root, asid, asid_generation) {
            self.records[handle.slot] = Some(Record {
                root,
                owner,
                generation,
                asid,
                asid_generation,
            });
            return Err(error);
        }
        self.asids[usize::from(asid) - 1] = Asid::Released(asid_generation);
        Ok(())
    }
    pub fn audit(&self) -> Result<(), Error> {
        for observed in self.owned_records() {
            let handle = Handle {
                registry: self.id,
                owner: observed.owner,
                slot: observed.slot,
                generation: observed.generation,
            };
            self.checked(handle, observed.owner)?;
            self.configuration.ttbr(observed.root, observed.asid)?;
            if self.owned_records().any(|other| {
                other.slot != observed.slot
                    && (other.root == observed.root || other.asid == observed.asid)
            }) {
                return Err(Error::Backing);
            }
        }
        for (asid, state) in self.asids.iter().enumerate() {
            if let Asid::Active { generation, slot } = *state {
                let row = self
                    .records
                    .get(slot)
                    .and_then(Option::as_ref)
                    .ok_or(Error::Backing)?;
                if usize::from(row.asid) != asid + 1 || row.asid_generation != generation {
                    return Err(Error::Backing);
                }
            }
        }
        Ok(())
    }
}

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

    #[derive(Debug)]
    struct Root(u64);
    #[derive(Default)]
    struct Model {
        acquires: u64,
        invalidations: u64,
        fail_next: bool,
    }
    impl Backend for Model {
        type Root = Root;
        fn acquire_root(&mut self, _: u64) -> Result<Root, Error> {
            self.acquires += 1;
            if self.fail_next {
                self.fail_next = false;
                return Err(Error::NoMemory);
            }
            Ok(Root(0x4000_0000 + 4096 * self.acquires))
        }
        fn root_address(&self, root: &Root) -> u64 {
            root.0
        }
        fn validate_root(&self, _: &Root) -> Result<(), Error> {
            Ok(())
        }
        fn invalidate_recycled(&mut self, _: u16, _: u64, _: u64) {
            self.invalidations += 1;
        }
        fn select(&mut self, _: &Root, _: u16) -> Result<(), Error> {
            Ok(())
        }
        fn restore_kernel(&mut self) -> Result<(), Error> {
            Ok(())
        }
        fn retire_root(&mut self, _: Root, _: u16, _: u64) -> Result<(), (Error, Root)> {
            Ok(())
        }
    }
    fn configuration() -> Configuration {
        Configuration::from_registers(2, 25).unwrap()
    }
    #[test]
    fn id_generation_exhaustion_retires_before_acquisition_and_uses_next_slot() {
        let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), Model::default()).unwrap();
        spaces.generations[0] = u64::MAX;
        let b = spaces.create(22).unwrap();
        assert_eq!(b.slot(), 1);
        assert!(spaces.retired_ids[0]);
        assert_eq!(spaces.generations[0], u64::MAX);
        assert_eq!(spaces.backend.acquires, 1);
        spaces.release(b, 22).unwrap();
        spaces.generations[1] = u64::MAX;
        assert_eq!(spaces.create(33), Err(Error::Generation));
        assert_eq!(spaces.backend.acquires, 1);
        assert!(spaces.retired_ids.iter().all(|retired| *retired));
        assert_eq!(spaces.owned_records().count(), 0);
        spaces.audit().unwrap();
    }
    #[test]
    fn asid_generation_exhaustion_retires_without_tlbi_or_acquisition() {
        let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), Model::default()).unwrap();
        spaces.asids[0] = Asid::Released(u64::MAX);
        let b = spaces.create(22).unwrap();
        assert_eq!(spaces.observe(b, 22).unwrap().asid, 2);
        assert_eq!(spaces.asids[0], Asid::Retired);
        assert_eq!(spaces.backend.invalidations, 0);
        spaces.release(b, 22).unwrap();
        spaces.asids[1] = Asid::Released(u64::MAX);
        assert_eq!(spaces.create(33), Err(Error::Generation));
        assert_eq!(spaces.backend.acquires, 1);
        assert_eq!(spaces.backend.invalidations, 0);
        assert_eq!(spaces.owned_records().count(), 0);
        spaces.audit().unwrap();
    }
    #[test]
    fn failed_recycled_root_keeps_both_generations_and_retirement_history() {
        let mut spaces = Spaces::<_, 2, 1>::new(1, configuration(), Model::default()).unwrap();
        let a = spaces.create(11).unwrap();
        spaces.release(a, 11).unwrap();
        spaces.backend.fail_next = true;
        assert_eq!(spaces.create(22), Err(Error::NoMemory));
        assert_eq!(spaces.asids[0], Asid::Released(1));
        assert_eq!(spaces.generations[0], 1);
        assert_eq!(spaces.backend.invalidations, 0);
        assert_eq!(spaces.owned_records().count(), 0);
        let b = spaces.create(22).unwrap();
        let observed = spaces.observe(b, 22).unwrap();
        assert_eq!((observed.generation, observed.asid_generation), (2, 2));
        assert_eq!(spaces.backend.invalidations, 1);
        assert_eq!(spaces.generations[1], 0);
        spaces.audit().unwrap();
    }
}