//! Physical frame ownership. The kernel's single source of truth for what
//! every frame of the user pool is used for and who is charged for it.

use types::{
    address::{PhysicalAddress, VirtualAddress},
    id::{ProcessId, RegionId},
    mapping::PagePerms,
};

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

#[cfg(kani)]
#[path = "proofs.rs"]
mod proofs;

pub const PAGE_SIZE: usize = 4096;

/// Frames the kernel can track. The board's user pool must fit.
pub const MAX_FRAMES: usize = if cfg!(kani) { 6 } else { 4096 };

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FrameUse {
    /// Mapped exactly once, at `page` in `owner`'s address space.
    User {
        owner: ProcessId,
        page: u32,
        perms: PagePerms,
    },
    /// Holds part of `owner`'s page tables.
    PageTable { owner: ProcessId },
    /// Backs a shared memory region.
    Region { region: RegionId },
}

impl FrameUse {
    /// A user mapping at `va`, which must be page aligned.
    pub fn user(owner: ProcessId, va: VirtualAddress, perms: PagePerms) -> Option<Self> {
        if !va.raw().is_multiple_of(PAGE_SIZE) {
            return None;
        }
        let page = u32::try_from(va.raw() / PAGE_SIZE).ok()?;
        Some(Self::User { owner, page, perms })
    }
}

pub struct FrameTable<const N: usize = MAX_FRAMES> {
    base: PhysicalAddress,
    total: usize,
    frames: [Option<FrameUse>; N],
    free: usize,
    /// Where the next allocation starts looking: just past the last one.
    cursor: usize,
}

impl<const N: usize> FrameTable<N> {
    pub const fn new(base: PhysicalAddress, total: usize) -> Self {
        let total = if total < N { total } else { N };
        Self {
            base,
            total,
            frames: [None; N],
            free: total,
            cursor: 0,
        }
    }

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

    pub fn capacity(&self) -> usize {
        self.total
    }

    pub fn free_count(&self) -> usize {
        self.free
    }

    fn index_of(&self, pa: PhysicalAddress) -> Option<usize> {
        let offset = pa.raw().checked_sub(self.base.raw())?;
        if !offset.is_multiple_of(PAGE_SIZE) {
            return None;
        }
        let idx = offset / PAGE_SIZE;
        (idx < self.total).then_some(idx)
    }

    fn address_of(&self, idx: usize) -> PhysicalAddress {
        self.base.offset(idx)
    }

    /// The frames that exist (`total` is clamped to `N` at construction).
    fn live(&self) -> &[Option<FrameUse>] {
        self.frames.get(..self.total).unwrap_or(&[])
    }

    fn live_mut(&mut self) -> &mut [Option<FrameUse>] {
        self.frames.get_mut(..self.total).unwrap_or(&mut [])
    }

    pub fn get(&self, pa: PhysicalAddress) -> Option<&FrameUse> {
        self.frames.get(self.index_of(pa)?)?.as_ref()
    }

    /// Next-fit: searches from just past the previous allocation, wrapping once.
    pub fn alloc(&mut self, usage: FrameUse) -> Option<PhysicalAddress> {
        if self.free == 0 {
            return None;
        }
        let total = self.total;
        let start = self.cursor.min(total);
        let idx = (start..total)
            .chain(0..start)
            .find(|&i| self.frames.get(i).is_some_and(Option::is_none))?;
        *self.frames.get_mut(idx)? = Some(usage);
        self.free -= 1;
        self.cursor = idx + 1;
        Some(self.address_of(idx))
    }

    /// Allocates `count` physically adjacent frames; frame `i` gets `usage(i)`.
    pub fn alloc_contiguous(
        &mut self,
        count: usize,
        usage: impl Fn(usize) -> FrameUse,
    ) -> Option<PhysicalAddress> {
        if count == 0 {
            return None;
        }
        let mut run = 0;
        let mut start = None;
        for (idx, f) in self.live().iter().enumerate() {
            if f.is_some() {
                run = 0;
                continue;
            }
            run += 1;
            if run == count {
                start = Some(idx + 1 - count);
                break;
            }
        }
        let start = start?;
        let run = self.live_mut().get_mut(start..start + count)?;
        for (i, f) in run.iter_mut().enumerate() {
            *f = Some(usage(i));
        }
        self.free -= count;
        Some(self.address_of(start))
    }

    /// Frees an allocated frame and reports what it was used for.
    pub fn free(&mut self, pa: PhysicalAddress) -> Option<FrameUse> {
        let idx = self.index_of(pa)?;
        let usage = self.frames.get_mut(idx)?.take()?;
        self.free += 1;
        Some(usage)
    }

    /// Changes the use of an allocated frame (e.g. moving it to another
    /// process). Fails on free frames.
    pub fn retag(&mut self, pa: PhysicalAddress, usage: FrameUse) -> bool {
        match self
            .index_of(pa)
            .and_then(|i| self.frames.get_mut(i))
            .and_then(Option::as_mut)
        {
            Some(slot) => {
                *slot = usage;
                true
            }
            None => false,
        }
    }

    /// The perms of the frame at `pa` if it is `owner`'s user page at `va`; O(1), the caller supplies the frame.
    pub fn user_at(
        &self,
        pa: PhysicalAddress,
        owner: ProcessId,
        va: VirtualAddress,
    ) -> Option<PagePerms> {
        let page = u32::try_from(va.raw() / PAGE_SIZE).ok()?;
        if !va.raw().is_multiple_of(PAGE_SIZE) {
            return None;
        }
        match self.get(pa)? {
            FrameUse::User {
                owner: o,
                page: p,
                perms,
            } if *o == owner && *p == page => Some(*perms),
            _ => None,
        }
    }

    pub fn find_user(
        &self,
        owner: ProcessId,
        va: VirtualAddress,
    ) -> Option<(PhysicalAddress, PagePerms)> {
        if !va.raw().is_multiple_of(PAGE_SIZE) {
            return None;
        }
        let page = u32::try_from(va.raw() / PAGE_SIZE).ok()?;
        self.live().iter().enumerate().find_map(|(idx, f)| match f {
            Some(FrameUse::User {
                owner: o,
                page: p,
                perms,
            }) if *o == owner && *p == page => Some((self.address_of(idx), *perms)),
            _ => None,
        })
    }

    /// Every allocated frame and its use.
    pub fn iter(&self) -> impl Iterator<Item = (PhysicalAddress, &FrameUse)> + '_ {
        self.live()
            .iter()
            .enumerate()
            .filter_map(|(idx, f)| f.as_ref().map(|u| (self.address_of(idx), u)))
    }

    /// User pages of `owner` as (frame, virtual address) pairs.
    pub fn user_frames_of(
        &self,
        owner: ProcessId,
    ) -> impl Iterator<Item = (PhysicalAddress, VirtualAddress)> + '_ {
        self.live()
            .iter()
            .enumerate()
            .filter_map(move |(idx, f)| match f {
                Some(FrameUse::User { owner: o, page, .. }) if *o == owner => Some((
                    self.address_of(idx),
                    VirtualAddress::from_raw(*page as usize * PAGE_SIZE),
                )),
                _ => None,
            })
    }

    /// Frees every user and page-table frame charged to `owner`; returns
    /// how many were freed.
    pub fn free_owned_by(&mut self, owner: ProcessId) -> usize {
        let mut freed = 0;
        for slot in self.live_mut().iter_mut() {
            let owned = matches!(
                slot,
                Some(FrameUse::User { owner: o, .. } | FrameUse::PageTable { owner: o }) if *o == owner
            );
            if owned {
                *slot = None;
                freed += 1;
            }
        }
        self.free += freed;
        freed
    }
}