pub mod allocation;
pub mod memory_map;

// use lib_alloc::allocator::FoundryAllocator;
use limine::{memory_map::EntryType, response::MemoryMapResponse};
use spin::{Mutex, Once};
use x86_64::{
    PhysAddr, VirtAddr,
    registers::control::Cr3,
    structures::paging::{FrameAllocator, OffsetPageTable, PageTable, PhysFrame, Size4KiB},
};

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));
    }
}

pub struct FoundryOSFrameAllocator {
    memory_map: &'static MemoryMapResponse,
    next: usize,
}

impl FoundryOSFrameAllocator {
    /// Creates a new `FoundryOSFrameAllocator` from a memory map.
    ///
    /// This function takes a reference to a `MemoryMapResponse` and initializes a
    /// `FoundryOSFrameAllocator` with it. The `next` field is set to 0, indicating that
    /// the first frame to be allocated is the first frame in the memory map.
    pub unsafe fn init(memory_map: &'static MemoryMapResponse) { unsafe {
        FRAME_ALLOCATOR.call_once(|| Mutex::new(Self {
            memory_map,
            next: 0,
        }));
    }}

    pub fn count_usable_frames(&self) -> u32 {
        self.usable_frames().count() as u32
    }

    /// An iterator over all usable frames in the memory map.
    ///
    /// Yields one `PhysFrame` for each available 4KiB frame in the memory map.
    ///
    /// This function is used to allocate frames for the pagemap.
    fn usable_frames(&self) -> impl Iterator<Item = PhysFrame> + use<> {
        let regions = self.memory_map.entries().iter();
        let usable_regions = regions.filter(|region| region.entry_type == EntryType::USABLE);
        let addr_ranges = usable_regions.map(|region| region.base..region.base + region.length);
        let frame_addresses = addr_ranges.flat_map(|r| r.step_by(4096));
        frame_addresses.map(|addr| PhysFrame::from_start_address(PhysAddr::new(addr)).unwrap())
    }
}

unsafe impl FrameAllocator<Size4KiB> for FoundryOSFrameAllocator {
    /// Allocates a frame from the list of usable frames.
    ///
    /// This function returns the next available `PhysFrame` from the memory map,
    /// if one exists. Once a frame is allocated, the internal counter is incremented
    /// to point to the next frame for future allocations.
    ///
    /// # Returns
    ///
    /// - `Some(PhysFrame)`: If a usable frame is available.
    /// - `None`: If there are no more usable frames to allocate.
    fn allocate_frame(&mut self) -> Option<PhysFrame> {
        let frame = self.usable_frames().nth(self.next);
        self.next += 1;
        frame
    }
}