pub mod allocation;
pub mod mapping;
pub mod units;

use allocation::page_alloc::FoundryOSFrameAllocator;
use spin::{Mutex, Once};
use units::MemoryUnits::*;
use x86_64::{
    VirtAddr,
    registers::control::Cr3,
    structures::paging::{OffsetPageTable, PageTable},
};

pub const STACK_VIRTUAL_SPACE: usize = 0x5555_5555_0000; // start address of the memory space where we store allocated stacks
pub const HEAP_VIRTUAL_SPACE: usize = 0x4444_4444_0000; // start address of heap allocated memory
pub const HEAP_SIZE: usize = MiB(1).to_bytes();

pub static FRAME_ALLOCATOR: Once<Mutex<FoundryOSFrameAllocator>> = Once::new();
pub static OFFSET_PAGE_TABLE: Once<Mutex<OffsetPageTable>> = Once::new();
/// Returns a mutable reference to the current level 4 page table.
///
/// # Safety
///
/// The caller must ensure that the level 4 page table is not modified
/// simultaneously. The caller must also ensure that the physical memory offset
/// is correct, to ensure that the correct virtual address is constructed.
unsafe fn active_l4_table(
    physical_memory_offset: VirtAddr,
) -> &'static mut PageTable {
    let (level_4_frame, _) = Cr3::read();

    let phys_addr = level_4_frame.start_address();
    let virt = phys_addr.as_u64() + physical_memory_offset.as_u64();
    unsafe { &mut *(virt as *mut PageTable) }
}

/// Initializes the `OffsetPageTable` for the current CPU architecture.
///
/// # Safety
///
/// This function must be called only once and should be called before any
/// memory operations are performed that rely on virtual memory management.
/// The provided `physical_memory_offset` must be accurate to ensure correct
/// translation of physical addresses.
///
/// # Parameters
///
/// - `physical_memory_offset`: The offset to convert physical addresses to
///   virtual addresses in the higher-half direct map.
///
/// # Returns
///
/// Returns an `OffsetPageTable` that allows for manipulation of the page
/// tables for the current CPU architecture.
pub fn init_page_table(physical_memory_offset: VirtAddr) {
    unsafe {
        let l4_table = active_l4_table(physical_memory_offset);
        let offset_table =
            OffsetPageTable::new(l4_table, physical_memory_offset);
        OFFSET_PAGE_TABLE.call_once(|| Mutex::new(offset_table));
    }
}