use crate::arch::x86_64::memory::{FRAME_ALLOCATOR, OFFSET_PAGE_TABLE};
use limine::memory_map::EntryType;
use limine::response::MemoryMapResponse;
use spin::Mutex;
use x86_64::structures::paging::mapper::{MapToError, TranslateResult};
use x86_64::structures::paging::{
    FrameAllocator, Mapper, Page, PageTableFlags, PhysFrame, Size4KiB,
    Translate,
};
use x86_64::{PhysAddr, VirtAddr};

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 fn init(memory_map: &'static MemoryMapResponse) {
        FRAME_ALLOCATOR.call_once(|| {
            Mutex::new(Self {
                memory_map,
                next: 0,
            })
        });
    }

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

    pub fn available_memory(&self) -> u64 {
        self.memory_map
            .entries()
            .iter()
            .map(|region| region.base..region.base + region.length)
            .flat_map(|r| r.step_by(4096))
            .count() as u64
            * 4096
    }

    /// 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 page map.
    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()
        })
    }

    pub(crate) fn is_mapped(virt_addr: VirtAddr) -> bool {
        let mapper = OFFSET_PAGE_TABLE.get().unwrap().lock();
        matches!(mapper.translate(virt_addr), TranslateResult::Mapped { .. })
    }

    pub(crate) fn allocate_page(
        &mut self,
        start_addr: VirtAddr,
        flags: PageTableFlags,
    ) -> Result<(), MapToError<Size4KiB>> {
        let page = Page::containing_address(start_addr);
        let frame = self
            .allocate_frame()
            .ok_or(MapToError::<Size4KiB>::FrameAllocationFailed)?;
        let mut mapper = OFFSET_PAGE_TABLE.get().unwrap().lock();

        unsafe { mapper.map_to(page, frame, flags, self)? }.flush();

        Ok(())
    }
}

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