use bootloader::bootinfo::{MemoryMap, MemoryRegionType};
use spin::Mutex;
use x86_64::{
    registers::control::Cr3,
    structures::paging::{
        mapper::{MapToError, MapperFlush},
        *,
    },
    PhysAddr, VirtAddr,
};

static MAPPER: Mutex<Option<OffsetPageTable>> = Mutex::new(None);
static ALLOCATOR: Mutex<Option<Allocator>> = Mutex::new(None);

pub(crate) fn init(phys_offset: u64, memory_map: &'static MemoryMap) {
    let table4 = unsafe { &mut *(Cr3::read().0.start_address().as_u64() as *mut PageTable) };
    MAPPER
        .lock()
        .replace(unsafe { OffsetPageTable::new(table4, VirtAddr::new(phys_offset)) });
    ALLOCATOR
        .lock()
        .replace(unsafe { Allocator::new(memory_map) });
}

pub fn translate_addr(addr: VirtAddr) -> Option<PhysAddr> {
    MAPPER.lock().as_ref().unwrap().translate_addr(addr)
}

pub fn map(
    page: Page,
    frame: PhysFrame,
    flags: PageTableFlags,
) -> Result<MapperFlush<Size4KiB>, MapToError<Size4KiB>> {
    unsafe {
        MAPPER.lock().as_mut().unwrap().map_to(
            page,
            frame,
            flags,
            ALLOCATOR.lock().as_mut().unwrap(),
        )
    }
}

pub fn map_next(
    page: Page,
    flags: PageTableFlags,
) -> Result<MapperFlush<Size4KiB>, MapToError<Size4KiB>> {
    let next_page = ALLOCATOR
        .lock()
        .as_mut()
        .unwrap()
        .allocate_frame()
        .ok_or(MapToError::FrameAllocationFailed)?;
    map(page, next_page, flags)
}

struct Allocator {
    memory_map: &'static MemoryMap,
    next: usize,
}

impl Allocator {
    unsafe fn new(memory_map: &'static MemoryMap) -> Self {
        Self {
            memory_map,
            next: 0,
        }
    }
}

unsafe impl FrameAllocator<Size4KiB> for Allocator {
    fn allocate_frame(&mut self) -> Option<PhysFrame<Size4KiB>> {
        let next = self
            .memory_map
            .iter()
            .filter(|v| v.region_type == MemoryRegionType::Usable)
            .map(|v| v.range.start_addr()..v.range.end_addr())
            .flat_map(|v| v.step_by(4096))
            .nth(self.next);

        self.next += 1;

        Some(PhysFrame::from_start_address(PhysAddr::new(next?)).unwrap())
    }
}