use crate::{capability::CapRef, generations::SlotGenerations};
use types::{
    address::{PhysicalAddress, VirtualAddress},
    bounded::PageCount,
    generation::Generation,
    id::{ProcessId, RegionId},
    mapping::PagePerms,
};

pub mod frame;
pub mod quota;

#[cfg(test)]
#[path = "quota_tests.rs"]
mod quota_tests;

#[cfg(test)]
#[path = "tests.rs"]
mod tests;

pub const MAX_REGIONS: usize = 64;
pub const MAX_REGION_MAPPINGS: usize = 4;

/// Where a region's frames come from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RegionBacking {
    /// Frames from the user pool, freed with the region.
    Pool,
    /// Device MMIO; never freed.
    Device,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RegionMapping {
    pub process: ProcessId,
    pub va: VirtualAddress,
    pub perms: PagePerms,
    /// The capability the mapping was made through. Deleting that
    /// capability removes the mapping; `None` for kernel-made mappings.
    pub via: Option<CapRef>,
}

/// Physically contiguous memory that may be mapped into several address
/// spaces. `owner` is charged for pool-backed frames.
pub struct Region {
    owner: ProcessId,
    backing: RegionBacking,
    base: PhysicalAddress,
    pages: PageCount,
    mappings: [Option<RegionMapping>; MAX_REGION_MAPPINGS],
}

impl Region {
    pub fn new(
        owner: ProcessId,
        backing: RegionBacking,
        base: PhysicalAddress,
        pages: PageCount,
    ) -> Self {
        Self {
            owner,
            backing,
            base,
            pages,
            mappings: [None; MAX_REGION_MAPPINGS],
        }
    }

    pub fn owner(&self) -> ProcessId {
        self.owner
    }

    pub fn backing(&self) -> RegionBacking {
        self.backing
    }

    pub fn base(&self) -> PhysicalAddress {
        self.base
    }

    pub fn pages(&self) -> PageCount {
        self.pages
    }

    pub fn mappings(&self) -> impl Iterator<Item = &RegionMapping> {
        self.mappings.iter().flatten()
    }

    /// Records that `mapping.process` maps this region. Each process maps a
    /// region at most once.
    pub fn add_mapping(&mut self, mapping: RegionMapping) -> bool {
        if self.mappings().any(|m| m.process == mapping.process) {
            return false;
        }
        match self.mappings.iter_mut().find(|m| m.is_none()) {
            Some(slot) => {
                *slot = Some(mapping);
                true
            }
            None => false,
        }
    }

    pub fn remove_mapping(&mut self, process: ProcessId) -> Option<RegionMapping> {
        self.mappings
            .iter_mut()
            .find(|m| m.is_some_and(|m| m.process == process))?
            .take()
    }

    /// Whether `va` falls inside this region's mapping in `process`.
    pub fn covers(&self, process: ProcessId, va: VirtualAddress) -> bool {
        let len = self.pages.value() as usize * frame::PAGE_SIZE;
        self.mappings()
            .any(|m| m.process == process && va.raw() >= m.va.raw() && va.raw() - m.va.raw() < len)
    }
}

pub struct RegionTable {
    regions: [Option<Region>; MAX_REGIONS],
    generations: SlotGenerations<MAX_REGIONS>,
}

impl RegionTable {
    pub fn new() -> Self {
        Self {
            regions: [const { None }; MAX_REGIONS],
            generations: SlotGenerations::new(),
        }
    }

    /// Generation capabilities to a live region must carry.
    pub fn generation(&self, id: RegionId) -> Option<Generation> {
        self.get(id)?;
        self.generations.current(id.raw() as usize)
    }

    fn free_slot(&self) -> Option<usize> {
        (0..MAX_REGIONS).find(|&i| self.regions[i].is_none() && self.generations.is_usable(i))
    }

    /// The id the next `insert` will return.
    pub fn next_free_id(&self) -> Option<RegionId> {
        let idx = self.free_slot()?;
        Some(RegionId::from_raw(idx as u16))
    }

    pub fn insert(&mut self, region: Region) -> Option<RegionId> {
        let idx = self.free_slot()?;
        self.regions[idx] = Some(region);
        Some(RegionId::from_raw(idx as u16))
    }

    pub fn get(&self, id: RegionId) -> Option<&Region> {
        self.regions.get(id.raw() as usize)?.as_ref()
    }

    pub fn get_mut(&mut self, id: RegionId) -> Option<&mut Region> {
        self.regions.get_mut(id.raw() as usize)?.as_mut()
    }

    /// Removes the region; capabilities to it become invalid even if the
    /// slot is later reused.
    pub fn remove(&mut self, id: RegionId) -> Option<Region> {
        let region = self.regions.get_mut(id.raw() as usize)?.take()?;
        self.generations.advance(id.raw() as usize);
        Some(region)
    }

    pub fn ids_owned_by(&self, owner: ProcessId) -> impl Iterator<Item = RegionId> + '_ {
        self.regions.iter().enumerate().filter_map(move |(i, r)| {
            r.as_ref()
                .filter(|r| r.owner == owner)
                .map(|_| RegionId::from_raw(i as u16))
        })
    }

    pub fn iter(&self) -> impl Iterator<Item = (RegionId, &Region)> {
        self.regions
            .iter()
            .enumerate()
            .filter_map(|(i, r)| r.as_ref().map(|r| (RegionId::from_raw(i as u16), r)))
    }

    /// Every region, by id, so callers can drop a process's mappings.
    pub fn iter_mut(&mut self) -> impl Iterator<Item = (RegionId, &mut Region)> {
        self.regions
            .iter_mut()
            .enumerate()
            .filter_map(|(i, r)| r.as_mut().map(|r| (RegionId::from_raw(i as u16), r)))
    }

    /// The region whose mapping in `process` covers `va`.
    pub fn covering(&self, process: ProcessId, va: VirtualAddress) -> Option<RegionId> {
        self.regions.iter().enumerate().find_map(|(i, r)| {
            r.as_ref()
                .filter(|r| r.covers(process, va))
                .map(|_| RegionId::from_raw(i as u16))
        })
    }
}

impl Default for RegionTable {
    fn default() -> Self {
        Self::new()
    }
}