use crate::memory::{FaultInfo, PhysAddr, VirtAddr, tlb::TLB};

pub struct MMU {
    buffer: TLB,
    enable_paging: bool,

    // Memory Management Register (MMR) Tells MMU where to find page tables
    mmr_value: u32,

    page_table: Box<[Option<PageTable>; 1024]>,
}

impl MMU {
    pub fn new() -> Self {
        MMU {
            buffer: TLB::new(),
            enable_paging: false,
            mmr_value: 0,
            page_table: Box::new([const { None }; 1024]),
        }
    }

    pub fn paging_enabled(&mut self, enable: bool) {
        self.enable_paging = enable;
    }

    pub fn tlb_insert(&mut self, virtual_address: VirtAddr, physical_address: PhysAddr) {
        self.buffer.insert(virtual_address, physical_address);
    }

    pub fn create_mapping(&mut self, virtual_address: VirtAddr, physical_address: PhysAddr) {
        let table = (virtual_address >> 22) as usize;
        let page = ((virtual_address >> 12) & 0x3FF) as usize;

        self.page_table[table].get_or_insert_default().entries[page] =
            Some(PageTableEntry::new(physical_address, true, true, false));
    }

    #[inline(always)]
    pub fn lookup(&mut self, virtual_address: VirtAddr) -> Result<PhysAddr, FaultInfo> {
        if !self.enable_paging {
            return Ok(virtual_address);
        }

        if let Some(addr) = self.buffer.lookup(virtual_address) {
            return Ok(addr);
        }

        self.walk(virtual_address)
    }

    #[cold]
    #[inline(never)]
    fn walk(&mut self, virtual_address: VirtAddr) -> Result<PhysAddr, FaultInfo> {
        let phys_addr = self.walk_page_table(virtual_address)?;
        self.buffer.insert(virtual_address, phys_addr);
        Ok(phys_addr)
    }

    fn walk_page_table(&mut self, virtual_address: VirtAddr) -> Result<PhysAddr, FaultInfo> {
        let table = (virtual_address >> 22) as usize;
        let page = ((virtual_address >> 12) & 0x3FF) as usize;

        if let Some(table) = &self.page_table[table] {
            return table.entries[page]
                .map(|entry| entry.physical_addr)
                .ok_or(FaultInfo::PageFault);
        } else {
            Err(FaultInfo::PageFault)
        }
    }
}

#[derive(Debug, Clone)]
struct PageDir {
    entries: Box<[PageTable; 1024]>,
}

#[derive(Debug, Clone)]
struct PageTable {
    entries: Box<[Option<PageTableEntry>; 1024]>,
}

impl Default for PageTable {
    fn default() -> Self {
        Self {
            entries: Box::new([None; 1024]),
        }
    }
}

#[derive(Debug, Clone, Copy)]
struct PageTableEntry {
    physical_addr: PhysAddr,
    writable: bool,
    present: bool,
    user: bool,
}

impl PageTableEntry {
    fn new(physical_addr: PhysAddr, writable: bool, present: bool, user: bool) -> Self {
        Self {
            physical_addr,
            writable,
            present,
            user,
        }
    }
}
