//! Transactional 4 KiB mapping *plan*, not hardware page-table descriptors.
//! No TTBR/SCTLR writes occur here. The owned frames model future table pages
//! and return to the allocator when the plan is dropped.

use crate::memory::{FrameAllocator, FreeMap, PAGE_SIZE};

const WINDOW_SIZE: u64 = 2 * 1024 * 1024;
const MAX_TABLES: usize = 32;
const MAX_MAPPINGS: usize = 64;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Kind {
    Ram,
    Mmio,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Access {
    ReadWrite,
    ReadExecute,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Mapping {
    pub virtual_start: u64,
    pub physical_start: u64,
    pub pages: u64,
    pub kind: Kind,
    pub access: Access,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Invalid,
    Overlap,
    UnsafeAttribute,
    TooManyMappings,
    TooManyTables,
    OutOfFrames,
}

pub struct Plan<'a> {
    allocator: &'a mut FrameAllocator,
    allowed_ram: FreeMap,
    root: u64,
    table_frames: [u64; MAX_TABLES],
    table_windows: [u64; MAX_TABLES],
    tables: usize,
    mappings: [Option<Mapping>; MAX_MAPPINGS],
    len: usize,
}

impl<'a> Plan<'a> {
    pub fn new(allocator: &'a mut FrameAllocator, allowed_ram: &FreeMap) -> Result<Self, Error> {
        let root = allocator.allocate().ok_or(Error::OutOfFrames)?;
        Ok(Self {
            allocator,
            allowed_ram: allowed_ram.clone(),
            root,
            table_frames: [0; MAX_TABLES],
            table_windows: [0; MAX_TABLES],
            tables: 0,
            mappings: [None; MAX_MAPPINGS],
            len: 0,
        })
    }

    pub fn root_frame(&self) -> u64 {
        self.root
    }
    pub fn table_count(&self) -> usize {
        self.tables
    }
    pub fn mappings(&self) -> &[Option<Mapping>] {
        &self.mappings[..self.len]
    }
    pub fn remaining_pages(&self) -> u64 {
        self.allocator.remaining_pages()
    }

    pub fn map(&mut self, mapping: Mapping) -> Result<(), Error> {
        if mapping.pages == 0
            || mapping.virtual_start % PAGE_SIZE != 0
            || mapping.physical_start % PAGE_SIZE != 0
        {
            return Err(Error::Invalid);
        }
        if mapping.kind == Kind::Mmio && mapping.access == Access::ReadExecute {
            return Err(Error::UnsafeAttribute);
        }
        let size = mapping.pages.checked_mul(PAGE_SIZE).ok_or(Error::Invalid)?;
        let virtual_end = mapping
            .virtual_start
            .checked_add(size)
            .ok_or(Error::Invalid)?;
        let physical_end = mapping
            .physical_start
            .checked_add(size)
            .ok_or(Error::Invalid)?;
        if mapping.kind == Kind::Ram
            && !self
                .allowed_ram
                .regions()
                .iter()
                .any(|region| region.start <= mapping.physical_start && physical_end <= region.end)
        {
            return Err(Error::UnsafeAttribute);
        }
        if mapping.kind == Kind::Ram
            && (overlaps_frame(mapping.physical_start, physical_end, self.root)
                || self.table_frames[..self.tables]
                    .iter()
                    .any(|frame| overlaps_frame(mapping.physical_start, physical_end, *frame)))
        {
            return Err(Error::UnsafeAttribute);
        }
        if self.len == MAX_MAPPINGS {
            return Err(Error::TooManyMappings);
        }
        for prior in self.mappings().iter().flatten() {
            let prior_end = prior.virtual_start + prior.pages * PAGE_SIZE;
            if mapping.virtual_start < prior_end && prior.virtual_start < virtual_end {
                return Err(Error::Overlap);
            }
        }
        let first = mapping.virtual_start / WINDOW_SIZE;
        let last = (virtual_end - 1) / WINDOW_SIZE;
        let mut needed = [0; MAX_TABLES];
        let mut count = 0;
        for window in first..=last {
            if !self.table_windows[..self.tables].contains(&window) {
                if count == MAX_TABLES {
                    return Err(Error::TooManyTables);
                }
                needed[count] = window;
                count += 1;
            }
        }
        if self.tables + count > MAX_TABLES {
            return Err(Error::TooManyTables);
        }
        if self.allocator.remaining_pages() < count as u64 {
            return Err(Error::OutOfFrames);
        }
        // All fallible checks precede mutation. If the allocator ever violates
        // remaining_pages, roll back each newly owned frame before returning.
        let mut fresh = [0; MAX_TABLES];
        for index in 0..count {
            let Some(frame) = self.allocator.allocate() else {
                for address in fresh[..index].iter().copied() {
                    self.allocator.release(address).expect("fresh table frame");
                }
                return Err(Error::OutOfFrames);
            };
            fresh[index] = frame;
        }
        if fresh[..count].iter().any(|frame| {
            (mapping.kind == Kind::Ram
                && overlaps_frame(mapping.physical_start, physical_end, *frame))
                || self.mappings().iter().flatten().any(|prior| {
                    prior.kind == Kind::Ram
                        && overlaps_frame(
                            prior.physical_start,
                            prior.physical_start + prior.pages * PAGE_SIZE,
                            *frame,
                        )
                })
        }) {
            for address in fresh[..count].iter().rev().copied() {
                self.allocator.release(address).expect("fresh table frame");
            }
            return Err(Error::UnsafeAttribute);
        }
        for index in 0..count {
            self.table_frames[self.tables] = fresh[index];
            self.table_windows[self.tables] = needed[index];
            self.tables += 1;
        }
        self.mappings[self.len] = Some(mapping);
        self.len += 1;
        Ok(())
    }
}

fn overlaps_frame(start: u64, end: u64, frame: u64) -> bool {
    start < frame + PAGE_SIZE && frame < end
}

impl Drop for Plan<'_> {
    fn drop(&mut self) {
        for frame in self.table_frames[..self.tables].iter().copied() {
            self.allocator.release(frame).expect("owned table frame");
        }
        self.allocator.release(self.root).expect("owned root frame");
    }
}

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

    fn allocator(pages: u64) -> (FrameAllocator, FreeMap) {
        let allowed = FreeMap::from_ram(0x4000_0000, pages * PAGE_SIZE).unwrap();
        (FrameAllocator::new(allowed.clone()), allowed)
    }
    fn ram(virtual_start: u64, pages: u64) -> Mapping {
        Mapping {
            virtual_start,
            physical_start: 0x4000_4000,
            pages,
            kind: Kind::Ram,
            access: Access::ReadWrite,
        }
    }

    #[test]
    fn owns_and_releases_all_table_frames() {
        let (mut frames, allowed) = allocator(8);
        let initial = frames.remaining_pages();
        {
            let mut plan = Plan::new(&mut frames, &allowed).unwrap();
            assert_eq!(plan.root_frame(), 0x4000_0000);
            plan.map(ram(0x4000_0000, 1)).unwrap();
            plan.map(ram(0x4020_0000, 1)).unwrap();
            assert_eq!(plan.table_count(), 2);
            assert_eq!(plan.remaining_pages(), initial - 3);
        }
        assert_eq!(frames.remaining_pages(), initial);
    }

    #[test]
    fn ram_and_mmio_have_distinct_attributes() {
        let (mut frames, allowed) = allocator(8);
        let mut plan = Plan::new(&mut frames, &allowed).unwrap();
        plan.map(Mapping {
            access: Access::ReadExecute,
            ..ram(0x4000_0000, 1)
        })
        .unwrap();
        let device = Mapping {
            virtual_start: 0x0900_0000,
            physical_start: 0x0900_0000,
            pages: 1,
            kind: Kind::Mmio,
            access: Access::ReadWrite,
        };
        plan.map(device).unwrap();
        assert_eq!(plan.mappings()[1], Some(device));
        assert_eq!(
            plan.map(Mapping {
                access: Access::ReadExecute,
                ..device
            }),
            Err(Error::UnsafeAttribute)
        );
    }

    #[test]
    fn invalid_or_overlapping_map_does_not_change_ownership() {
        let (mut frames, allowed) = allocator(4);
        let mut plan = Plan::new(&mut frames, &allowed).unwrap();
        plan.map(Mapping {
            physical_start: 0x4000_2000,
            ..ram(0x4000_0000, 2)
        })
        .unwrap();
        let before = (
            plan.remaining_pages(),
            plan.table_count(),
            plan.mappings().len(),
        );
        assert_eq!(
            plan.map(Mapping {
                physical_start: 0x4000_2000,
                ..ram(0x4000_1000, 1)
            }),
            Err(Error::Overlap)
        );
        assert_eq!(plan.map(ram(0x4000_0001, 1)), Err(Error::Invalid));
        assert_eq!(
            plan.map(ram(u64::MAX - PAGE_SIZE + 1, 2)),
            Err(Error::Invalid)
        );
        assert_eq!(
            (
                plan.remaining_pages(),
                plan.table_count(),
                plan.mappings().len()
            ),
            before
        );
    }

    #[test]
    fn out_of_frames_does_not_add_partial_mapping() {
        let (mut frames, allowed) = allocator(2);
        let mut plan = Plan::new(&mut frames, &allowed).unwrap();
        assert_eq!(
            plan.map(Mapping {
                kind: Kind::Mmio,
                physical_start: 0x0900_0000,
                ..ram(0x4000_0000, 513)
            }),
            Err(Error::OutOfFrames)
        );
        assert_eq!(plan.table_count(), 0);
        assert!(plan.mappings().is_empty());
        assert_eq!(plan.remaining_pages(), 1);
    }

    #[test]
    fn all_ram_reserved_returns_out_of_frames() {
        const BASE: u64 = 0x4000_0000;
        let mut allowed = FreeMap::from_ram(BASE, 4 * PAGE_SIZE).unwrap();
        allowed.reserve(BASE, 4 * PAGE_SIZE).unwrap();
        assert!(allowed.regions().is_empty());
        let mut frames = FrameAllocator::new(allowed.clone());
        assert_eq!(frames.remaining_pages(), 0);
        assert!(matches!(
            Plan::new(&mut frames, &allowed),
            Err(Error::OutOfFrames)
        ));
        assert_eq!(frames.remaining_pages(), 0);
    }

    #[test]
    fn protected_ram_target_is_rejected_without_frame_loss() {
        let mut allowed = FreeMap::from_ram(0x4000_0000, 8 * PAGE_SIZE).unwrap();
        allowed.reserve(0x4000_2000, PAGE_SIZE).unwrap();
        let mut frames = FrameAllocator::new(allowed.clone());
        let mut plan = Plan::new(&mut frames, &allowed).unwrap();
        let before = plan.remaining_pages();
        assert_eq!(
            plan.map(Mapping {
                physical_start: 0x4000_2000,
                ..ram(0x5000_0000, 1)
            }),
            Err(Error::UnsafeAttribute)
        );
        assert_eq!(plan.remaining_pages(), before);
        assert!(plan.mappings().is_empty());
    }

    #[test]
    fn rejects_table_frame_as_data_target_without_leak() {
        let (mut frames, allowed) = allocator(4);
        let mut plan = Plan::new(&mut frames, &allowed).unwrap();
        let before = plan.remaining_pages();
        assert_eq!(
            plan.map(Mapping {
                physical_start: 0x4000_1000,
                ..ram(0x5000_0000, 1)
            }),
            Err(Error::UnsafeAttribute)
        );
        assert_eq!(plan.remaining_pages(), before);
        assert_eq!(plan.table_count(), 0);
    }

    #[test]
    fn new_table_frame_cannot_alias_prior_data_mapping() {
        let (mut frames, allowed) = allocator(8);
        let mut plan = Plan::new(&mut frames, &allowed).unwrap();
        plan.map(Mapping {
            physical_start: 0x4000_2000,
            ..ram(0x4000_0000, 1)
        })
        .unwrap();
        let before = (
            plan.remaining_pages(),
            plan.table_count(),
            plan.mappings().len(),
        );
        assert_eq!(
            plan.map(Mapping {
                physical_start: 0x4000_4000,
                ..ram(0x4020_0000, 1)
            }),
            Err(Error::UnsafeAttribute)
        );
        assert_eq!(
            (
                plan.remaining_pages(),
                plan.table_count(),
                plan.mappings().len()
            ),
            before
        );
    }
}