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

/// 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();
    &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 unsafe fn init(physical_memory_offset: VirtAddr) -> OffsetPageTable<'static> {
    let l4_table = active_l4_table(physical_memory_offset);
    OffsetPageTable::new(l4_table, physical_memory_offset)
}

pub(crate) 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) -> FoundryOSFrameAllocator {
        FoundryOSFrameAllocator {
            memory_map,
            next: 0,
        }
    }

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

// pub unsafe fn translate_addr(addr: VirtAddr, physical_memory_offset: VirtAddr) -> Option<PhysAddr> {
//     translate_addr_inner(addr, physical_memory_offset)
// }

// fn translate_addr_inner(addr: VirtAddr, physical_memory_offset: VirtAddr) -> Option<PhysAddr> {
//     let (l4_table_frame, _) = Cr3::read();

//     let table_indexes = [
//         addr.p4_index(),
//         addr.p3_index(),
//         addr.p2_index(),
//         addr.p1_index(),
//     ];
//     let mut frame = l4_table_frame;

//     for &i in &table_indexes {
//         let virt = physical_memory_offset + frame.start_address().as_u64();
//         let table_ptr: *const PageTable = virt.as_ptr();
//         let table = unsafe { &*table_ptr };

//         let entry = &table[i];

//         frame = match entry.frame() {
//             Ok(frame) => frame,
//             Err(FrameError::FrameNotPresent) => return None,
//             Err(FrameError::HugeFrame) => panic!("huge frames are not supported!"),
//         };
//     }

//     Some(frame.start_address() + u64::from(addr.page_offset()))
// }