//! Paging: page tables, address translation and a physical frame allocator.
//!
//! x86_64 uses 4-level page tables (L4 -> L3 -> L2 -> L1 -> 4 KiB page).
//! CR3 holds the *physical* address of the active level-4 table. But the CPU
//! (and our code) only uses *virtual* addresses, so to read or edit a page
//! table we need it mapped somewhere. The bootloader's `map_physical_memory`
//! feature maps all of physical RAM at `physical_memory_offset`, so
//!     virtual address = physical_memory_offset + physical address.
//! `OffsetPageTable` from the x86_64 crate implements exactly that.

use bootloader::bootinfo::{MemoryMap, MemoryRegionType};
use x86_64::{
    PhysAddr, VirtAddr,
    structures::paging::{
        FrameAllocator, Mapper, OffsetPageTable, Page, PageTable, PageTableFlags, PhysFrame,
        Size4KiB,
    },
};

/// Create an `OffsetPageTable` for the currently active page tables.
///
/// Unsafe: the caller must guarantee that all physical memory really is
/// mapped at `physical_memory_offset`, and must call this only once (two
/// `&mut` to the same page table would be undefined behaviour).
pub unsafe fn init(physical_memory_offset: VirtAddr) -> OffsetPageTable<'static> {
    unsafe {
        let level_4_table = active_level_4_table(physical_memory_offset);
        OffsetPageTable::new(level_4_table, physical_memory_offset)
    }
}

/// Find the active level-4 table via CR3 and return a reference to it.
unsafe fn active_level_4_table(physical_memory_offset: VirtAddr) -> &'static mut PageTable {
    use x86_64::registers::control::Cr3;

    let (level_4_table_frame, _flags) = Cr3::read();
    let phys = level_4_table_frame.start_address();
    let virt = physical_memory_offset + phys.as_u64();
    let page_table_ptr: *mut PageTable = virt.as_mut_ptr();
    unsafe { &mut *page_table_ptr }
}

/// Demo: map the virtual page containing `page` to the VGA text buffer's
/// physical frame (0xb8000). Writing through that page then shows on screen.
pub fn create_example_mapping(
    page: Page,
    mapper: &mut OffsetPageTable,
    frame_allocator: &mut impl FrameAllocator<Size4KiB>,
) {
    let frame = PhysFrame::containing_address(PhysAddr::new(0xb8000));
    let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;

    // Unsafe because mapping the same frame twice could create aliasing
    // `&mut`s. Fine here: it's only the VGA buffer, for a demo.
    // `map_to` may need new page tables (L3/L2/L1) -> takes a frame allocator.
    let map_to_result = unsafe { mapper.map_to(page, frame, flags, frame_allocator) };
    // `flush` removes the stale entry from the TLB (the CPU's mapping cache).
    map_to_result.expect("map_to failed").flush();
}

/// Hands out unused physical 4 KiB frames, taken from the memory map the
/// bootloader got from the BIOS (E820). Very simple: it never frees, and
/// walks the map from the start on every call (fine for now).
pub struct BootInfoFrameAllocator {
    memory_map: &'static MemoryMap,
    next: usize, // how many frames we've handed out so far
}

impl BootInfoFrameAllocator {
    /// Unsafe: the caller must guarantee that every frame marked `Usable` in
    /// the memory map really is unused.
    pub unsafe fn init(memory_map: &'static MemoryMap) -> Self {
        BootInfoFrameAllocator { memory_map, next: 0 }
    }

    /// All usable frames, in order.
    fn usable_frames(&self) -> impl Iterator<Item = PhysFrame> {
        self.memory_map
            .iter()
            .filter(|r| r.region_type == MemoryRegionType::Usable) // only free RAM
            .map(|r| r.range.start_addr()..r.range.end_addr()) // -> address ranges
            .flat_map(|r| r.step_by(4096)) // -> each frame's start address
            .map(|addr| PhysFrame::containing_address(PhysAddr::new(addr)))
    }
}

unsafe impl FrameAllocator<Size4KiB> for BootInfoFrameAllocator {
    fn allocate_frame(&mut self) -> Option<PhysFrame> {
        let frame = self.usable_frames().nth(self.next);
        self.next += 1;
        frame
    }
}

/// Total usable RAM in bytes, according to the memory map (for printing).
pub fn usable_memory_bytes(memory_map: &MemoryMap) -> u64 {
    memory_map
        .iter()
        .filter(|r| r.region_type == MemoryRegionType::Usable)
        .map(|r| r.range.end_addr() - r.range.start_addr())
        .sum()
}