//! Complete range preflight over an exclusively owned root backend.
//! The backend stages a private inactive replacement root and preserves data
//! ownership across old-root retirement. Commit performs the actual synchronized
//! root switch and TLB retirement; this core alone proves no hardware behavior.
use core::sync::atomic::{AtomicU64, Ordering};
use turkos_os_page_tables::descriptor::{Memory, Permissions};

static NEXT_MAPPING_ID: AtomicU64 = AtomicU64::new(1);

pub const PAGE: u64 = 4096;
const MAX_PINS: usize = 8;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Invalid,
    Overflow,
    OutsideWindow,
    Foreign,
    Conflict,
    Unmapped,
    Guard,
    Permissions,
    Capacity,
    Generation,
    Busy,
    StalePin,
    Backing,
    NoMemory,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Owner {
    pub id: u64,
    pub generation: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Window {
    start: u64,
    end: u64,
}
impl Window {
    pub fn new(start: u64, end: u64, virtual_bits: u8) -> Result<Self, Error> {
        if !matches!(virtual_bits, 39 | 48)
            || start >= end
            || !start.is_multiple_of(PAGE)
            || !end.is_multiple_of(PAGE)
        {
            return Err(Error::Invalid);
        }
        if end > 1u64 << virtual_bits {
            return Err(Error::OutsideWindow);
        }
        Ok(Self { start, end })
    }
    fn addresses(self, start: u64, pages: usize) -> Result<(), Error> {
        if pages == 0 || !start.is_multiple_of(PAGE) {
            return Err(Error::Invalid);
        }
        let bytes = u64::try_from(pages)
            .map_err(|_| Error::Overflow)?
            .checked_mul(PAGE)
            .ok_or(Error::Overflow)?;
        let end = start.checked_add(bytes).ok_or(Error::Overflow)?;
        if start < self.start || end > self.end {
            return Err(Error::OutsideWindow);
        }
        Ok(())
    }
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Value<R: Copy + Eq> {
    pub frame: R,
    pub permissions: Permissions,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Change<R: Copy + Eq> {
    pub virtual_address: u64,
    pub before: Option<Value<R>>,
    pub after: Option<Value<R>>,
}
/// FrameRef is an inert reference to a retained private data-page owner.
/// Backend must validate its owner/generation and physical address before use.
/// Table and data ownership are separate: retirement of an old root does not
/// return backing still referenced by a replacement root.
pub trait Backend {
    type FrameRef: Copy + Eq;
    type Prepared;
    fn owner(&self) -> Owner;
    fn validate_frame(&self, owner: Owner, frame: Self::FrameRef) -> Result<(), Error>;
    fn current(&self, virtual_address: u64) -> Result<Option<Value<Self::FrameRef>>, Error>;
    /// Preparation acquires every draft table before any visible prefix. Any
    /// failure leaves current PTEs intact and retains failed rollback handles.
    fn prepare(&mut self, changes: &[Change<Self::FrameRef>]) -> Result<Self::Prepared, Error>;
    /// Infallible while the trusted adapter owns root quiescence and all table
    /// backing. It switches the complete draft and finishes required BBM/TLBI
    /// before returning old table resources or permitting data reclamation.
    fn commit(&mut self, prepared: Self::Prepared);
}
#[derive(Clone, Copy)]
struct Entry<R: Copy + Eq> {
    address: u64,
    value: Option<Value<R>>,
    generation: u64,
    pins: u8,
}
pub struct Pin<const N: usize> {
    owner: Owner,
    scope: u64,
    slot: usize,
    generation: u64,
    entries: [usize; N],
    count: usize,
    consumed: bool,
}
pub struct Mapping<B: Backend, const N: usize> {
    backend: B,
    owner: Owner,
    scope: u64,
    identity: u64,
    window: Window,
    entries: [Option<Entry<B::FrameRef>>; N],
    next_generation: u64,
    pin_generations: [u64; MAX_PINS],
    active_pins: [bool; MAX_PINS],
}
impl<B: Backend, const N: usize> Mapping<B, N> {
    pub fn new(backend: B, scope: u64, window: Window) -> Result<Self, Error> {
        let owner = backend.owner();
        if scope == 0 || N == 0 || owner.id == 0 || owner.generation == 0 {
            return Err(Error::Invalid);
        }
        let identity = NEXT_MAPPING_ID
            .try_update(Ordering::Relaxed, Ordering::Relaxed, |next| {
                next.checked_add(1)
            })
            .map_err(|_| Error::Generation)?;
        Ok(Self {
            backend,
            owner,
            scope,
            identity,
            window,
            entries: [None; N],
            next_generation: 1,
            pin_generations: [0; MAX_PINS],
            active_pins: [false; MAX_PINS],
        })
    }
    pub fn with_backend_observation<R>(&self, callback: impl FnOnce(&B) -> R) -> R {
        callback(&self.backend)
    }
    pub const fn scope(&self) -> u64 {
        self.scope
    }
    fn owner(&self, owner: Owner) -> Result<(), Error> {
        if owner == self.owner && owner == self.backend.owner() {
            Ok(())
        } else {
            Err(Error::Foreign)
        }
    }
    fn range(&self, start: u64, pages: usize) -> Result<(), Error> {
        self.window.addresses(start, pages)?;
        if pages > N {
            Err(Error::Capacity)
        } else {
            Ok(())
        }
    }
    fn slot(&self, address: u64) -> Option<usize> {
        self.entries
            .iter()
            .position(|row| row.is_some_and(|row| row.address == address))
    }
    fn permission(permissions: Permissions) -> Result<(), Error> {
        if !permissions.user {
            return Err(Error::Permissions);
        }
        permissions
            .validate(Memory::Normal)
            .map_err(|_| Error::Permissions)
    }
    fn generation(&self) -> Result<u64, Error> {
        self.next_generation.checked_add(1).ok_or(Error::Generation)
    }
    fn aliases(&self, changes: &[Change<B::FrameRef>]) -> Result<(), Error> {
        let incompatible = |a: Value<B::FrameRef>, b: Value<B::FrameRef>| {
            a.frame == b.frame
                && (a.permissions.write && b.permissions.execute
                    || b.permissions.write && a.permissions.execute)
        };
        for (index, change) in changes.iter().enumerate() {
            if let Some(value) = change.after {
                for other in self.entries.iter().flatten().filter(|entry| {
                    !changes
                        .iter()
                        .any(|changed| changed.virtual_address == entry.address)
                }) {
                    if other.value.is_some_and(|other| incompatible(value, other)) {
                        return Err(Error::Permissions);
                    }
                }
                for other in &changes[index + 1..] {
                    if other.after.is_some_and(|other| incompatible(value, other)) {
                        return Err(Error::Permissions);
                    }
                }
            }
        }
        Ok(())
    }
    pub fn map(
        &mut self,
        owner: Owner,
        start: u64,
        frames: &[B::FrameRef],
        permissions: Permissions,
    ) -> Result<(), Error> {
        self.owner(owner)?;
        self.range(start, frames.len())?;
        Self::permission(permissions)?;
        let next = self.generation()?;
        let mut free = [0usize; N];
        let mut free_len = 0;
        for (index, row) in self.entries.iter().enumerate() {
            if row.is_none() {
                free[free_len] = index;
                free_len += 1;
            }
        }
        if free_len < frames.len() {
            return Err(Error::Capacity);
        }
        let empty = Change {
            virtual_address: 0,
            before: None,
            after: None,
        };
        let mut changes = [empty; N];
        for (index, frame) in frames.iter().copied().enumerate() {
            let address = start + index as u64 * PAGE;
            if self.slot(address).is_some() || self.backend.current(address)?.is_some() {
                return Err(Error::Conflict);
            }
            self.backend.validate_frame(owner, frame)?;
            changes[index] = Change {
                virtual_address: address,
                before: None,
                after: Some(Value { frame, permissions }),
            };
        }
        // Omitted prospective backing alias policy.
        let prepared = self.backend.prepare(&changes[..frames.len()])?;
        self.backend.commit(prepared);
        for (index, change) in changes[..frames.len()].iter().enumerate() {
            self.entries[free[index]] = Some(Entry {
                address: change.virtual_address,
                value: change.after,
                generation: self.next_generation,
                pins: 0,
            });
        }
        self.next_generation = next;
        Ok(())
    }
    pub fn guard(&mut self, owner: Owner, start: u64, pages: usize) -> Result<(), Error> {
        self.owner(owner)?;
        self.range(start, pages)?;
        let next = self.generation()?;
        let mut free = [0usize; N];
        let mut len = 0;
        for (index, row) in self.entries.iter().enumerate() {
            if row.is_none() {
                free[len] = index;
                len += 1;
            }
        }
        if len < pages {
            return Err(Error::Capacity);
        }
        for index in 0..pages {
            let address = start + index as u64 * PAGE;
            if self.slot(address).is_some() || self.backend.current(address)?.is_some() {
                return Err(Error::Conflict);
            }
        }
        // Guard reservations deliberately never prepare a physical table/data
        // transaction; absent PTEs remain absent with zero backing acquisitions.
        for (index, slot) in free[..pages].iter().enumerate() {
            self.entries[*slot] = Some(Entry {
                address: start + index as u64 * PAGE,
                value: None,
                generation: self.next_generation,
                pins: 0,
            });
        }
        self.next_generation = next;
        Ok(())
    }
    fn modify(
        &mut self,
        owner: Owner,
        start: u64,
        pages: usize,
        permissions: Option<Permissions>,
    ) -> Result<(), Error> {
        self.owner(owner)?;
        self.range(start, pages)?;
        if let Some(permissions) = permissions {
            Self::permission(permissions)?;
        }
        let next = self.generation()?;
        let mut slots = [0usize; N];
        let mut changes = [Change {
            virtual_address: 0,
            before: None,
            after: None,
        }; N];
        let mut physical = false;
        for index in 0..pages {
            let address = start + index as u64 * PAGE;
            let slot = self.slot(address).ok_or(Error::Unmapped)?;
            let entry = self.entries[slot].ok_or(Error::Unmapped)?;
            if entry.pins != 0 {
                return Err(Error::Busy);
            }
            if self.backend.current(address)? != entry.value {
                return Err(Error::Backing);
            }
            let after = if let Some(permissions) = permissions {
                let value = entry.value.ok_or(Error::Guard)?;
                self.backend.validate_frame(owner, value.frame)?;
                Some(Value {
                    permissions,
                    ..value
                })
            } else {
                None
            };
            physical |= entry.value.is_some();
            slots[index] = slot;
            changes[index] = Change {
                virtual_address: address,
                before: entry.value,
                after,
            };
        }
        if physical {
            self.aliases(&changes[..pages])?;
            let prepared = self.backend.prepare(&changes[..pages])?;
            self.backend.commit(prepared);
        }
        for index in 0..pages {
            self.entries[slots[index]] = changes[index].after.map(|value| Entry {
                address: changes[index].virtual_address,
                value: Some(value),
                generation: self.next_generation,
                pins: 0,
            });
        }
        self.next_generation = next;
        Ok(())
    }
    pub fn unmap(&mut self, owner: Owner, start: u64, pages: usize) -> Result<(), Error> {
        self.modify(owner, start, pages, None)
    }
    pub fn protect(
        &mut self,
        owner: Owner,
        start: u64,
        pages: usize,
        permissions: Permissions,
    ) -> Result<(), Error> {
        self.modify(owner, start, pages, Some(permissions))
    }
    pub fn lookup(&self, owner: Owner, address: u64) -> Result<Option<Value<B::FrameRef>>, Error> {
        self.owner(owner)?;
        let entry =
            self.entries[self.slot(address).ok_or(Error::Unmapped)?].ok_or(Error::Unmapped)?;
        Ok(entry.value)
    }
    pub fn pin(
        &mut self,
        owner: Owner,
        start: u64,
        pages: usize,
        write: bool,
    ) -> Result<Pin<N>, Error> {
        self.owner(owner)?;
        self.range(start, pages)?;
        let slot = self
            .active_pins
            .iter()
            .enumerate()
            .position(|(index, active)| !*active && self.pin_generations[index] != u64::MAX)
            .ok_or_else(|| {
                if self.active_pins.iter().any(|active| !*active) {
                    Error::Generation
                } else {
                    Error::Capacity
                }
            })?;
        let generation = self.pin_generations[slot]
            .checked_add(1)
            .ok_or(Error::Generation)?;
        let mut entries = [0usize; N];
        for (index, item) in entries[..pages].iter_mut().enumerate() {
            let address = start + index as u64 * PAGE;
            let at = self.slot(address).ok_or(Error::Unmapped)?;
            let entry = self.entries[at].ok_or(Error::Unmapped)?;
            let value = entry.value.ok_or(Error::Guard)?;
            if write && !value.permissions.write {
                return Err(Error::Permissions);
            }
            if entry.pins == u8::MAX {
                return Err(Error::Capacity);
            }
            self.backend.validate_frame(owner, value.frame)?;
            if self.backend.current(address)? != entry.value {
                return Err(Error::Backing);
            }
            *item = at;
        }
        for at in &entries[..pages] {
            self.entries[*at].as_mut().ok_or(Error::Unmapped)?.pins += 1;
        }
        self.active_pins[slot] = true;
        self.pin_generations[slot] = generation;
        Ok(Pin {
            owner,
            scope: self.identity,
            slot,
            generation,
            entries,
            count: pages,
            consumed: false,
        })
    }
    pub fn unpin(&mut self, pin: &mut Pin<N>) -> Result<(), Error> {
        if pin.consumed
            || pin.owner != self.owner
            || pin.scope != self.identity
            || pin.slot >= MAX_PINS
            || !self.active_pins[pin.slot]
            || self.pin_generations[pin.slot] != pin.generation
        {
            return Err(Error::StalePin);
        }
        for at in &pin.entries[..pin.count] {
            if self.entries[*at].is_none_or(|row| row.pins == 0) {
                return Err(Error::Backing);
            }
        }
        for at in &pin.entries[..pin.count] {
            self.entries[*at].as_mut().ok_or(Error::Backing)?.pins -= 1;
        }
        self.active_pins[pin.slot] = false;
        pin.consumed = true;
        Ok(())
    }
    pub fn records(&self) -> impl Iterator<Item = (u64, Option<Value<B::FrameRef>>, u64, u8)> + '_ {
        self.entries
            .iter()
            .flatten()
            .map(|row| (row.address, row.value, row.generation, row.pins))
    }
    // A refused shutdown returns the owned root and every live record; boxing
    // would require the allocator this bounded core deliberately avoids.
    #[allow(clippy::result_large_err)]
    pub fn into_backend(self) -> Result<B, (Error, Self)> {
        if self.active_pins.iter().any(|active| *active) || self.entries.iter().any(Option::is_some)
        {
            Err((Error::Busy, self))
        } else {
            Ok(self.backend)
        }
    }
}

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

    const OWNER: Owner = Owner {
        id: 1,
        generation: 1,
    };
    const RW: Permissions = Permissions {
        user: true,
        write: true,
        execute: false,
    };
    struct Memory {
        entries: [Option<Value<u8>>; 4],
        prepares: usize,
    }
    impl Backend for Memory {
        type FrameRef = u8;
        type Prepared = [Option<Value<u8>>; 4];
        fn owner(&self) -> Owner {
            OWNER
        }
        fn validate_frame(&self, owner: Owner, frame: u8) -> Result<(), Error> {
            if owner == OWNER && frame < 4 {
                Ok(())
            } else {
                Err(Error::Backing)
            }
        }
        fn current(&self, address: u64) -> Result<Option<Value<u8>>, Error> {
            Ok(self.entries[(address / PAGE) as usize])
        }
        fn prepare(&mut self, changes: &[Change<u8>]) -> Result<Self::Prepared, Error> {
            self.prepares += 1;
            let mut draft = self.entries;
            for change in changes {
                draft[(change.virtual_address / PAGE) as usize] = change.after;
            }
            Ok(draft)
        }
        fn commit(&mut self, draft: Self::Prepared) {
            self.entries = draft;
        }
    }
    fn mapper() -> Mapping<Memory, 4> {
        Mapping::new(
            Memory {
                entries: [None; 4],
                prepares: 0,
            },
            841,
            Window::new(0, 4 * PAGE, 39).unwrap(),
        )
        .unwrap()
    }
    #[test]
    fn map_generation_exhaustion_precedes_physical_preparation() {
        let mut mapper = mapper();
        mapper.next_generation = u64::MAX;
        assert_eq!(mapper.map(OWNER, 0, &[0], RW), Err(Error::Generation));
        assert_eq!(mapper.backend.prepares, 0);
        assert_eq!(mapper.records().count(), 0);
    }
    #[test]
    fn guard_generation_exhaustion_preserves_absent_backing() {
        let mut mapper = mapper();
        mapper.next_generation = u64::MAX;
        assert_eq!(mapper.guard(OWNER, 0, 1), Err(Error::Generation));
        assert_eq!(mapper.backend.prepares, 0);
        assert_eq!(mapper.records().count(), 0);
    }
    #[test]
    fn retired_pin_slot_is_skipped_without_generation_wrap() {
        let mut mapper = mapper();
        mapper.map(OWNER, 0, &[0], RW).unwrap();
        mapper.pin_generations[0] = u64::MAX;
        let mut pin = mapper.pin(OWNER, 0, 1, true).unwrap();
        assert_eq!(pin.slot, 1);
        assert_eq!(pin.generation, 1);
        mapper.unpin(&mut pin).unwrap();
        assert_eq!(mapper.pin_generations[0], u64::MAX);
    }
    #[test]
    fn all_pin_generations_exhausted_preserve_every_mapping() {
        let mut mapper = mapper();
        mapper.map(OWNER, 0, &[0], RW).unwrap();
        mapper.pin_generations.fill(u64::MAX);
        assert!(matches!(
            mapper.pin(OWNER, 0, 1, true),
            Err(Error::Generation)
        ));
        assert!(!mapper.active_pins.iter().any(|value| *value));
        assert!(mapper.records().all(|(_, _, _, pins)| pins == 0));
        assert_eq!(mapper.backend.prepares, 1);
        assert_eq!(mapper.lookup(OWNER, 0).unwrap().unwrap().frame, 0);
    }
}