use crate::{bootinfo::BOOTINFO, println, virtual_memory::AsVirt};
use bootloader::boot_info::MemoryRegionKind;
use core::mem;
use spin::{Lazy, Mutex};
use tap::Tap;
use x86_64::{
    structures::paging::{FrameAllocator, FrameDeallocator, PhysFrame, Size4KiB},
    PhysAddr,
};

type FrameIterator = impl Iterator<Item = PhysFrame>;

#[derive(Default)]
struct FrameStack(Option<(&'static mut FrameStack, PhysFrame)>);
pub struct FrameManager {
    stack: FrameStack,
    iter: FrameIterator,
}

unsafe impl FrameAllocator<Size4KiB> for FrameManager {
    fn allocate_frame(&mut self) -> Option<PhysFrame> {
        self.iter.next().or_else(|| self.stack.pop())
    }
}

impl FrameStack {
    /// # Safety
    /// The frame must be unique
    unsafe fn push(&mut self, frame: PhysFrame) {
        self.0 = Some((
            unsafe { &mut *frame.as_virt_ptr::<Self>().tap_mut(|x| x.write(mem::take(self))) },
            frame,
        ));
    }

    fn pop(&mut self) -> Option<PhysFrame> {
        let (head, frame) = self.0.take()?;
        *self = mem::take(head);
        Some(frame)
    }
}

impl FrameDeallocator<Size4KiB> for FrameManager {
    unsafe fn deallocate_frame(&mut self, frame: PhysFrame) {
        unsafe { self.stack.push(frame) }
    }
}

pub static PHYSICAL_MEMORY: Lazy<Mutex<FrameManager>> = Lazy::new(|| {
    let region_iter =
        BOOTINFO.memory_regions.iter().filter(|region| region.kind == MemoryRegionKind::Usable);
    let frame_iter = region_iter
        .clone()
        .flat_map(|region| {
            ((region.start >> 12) + if region.start & 0xFFF > 0 { 1 } else { 0 })
                ..(region.end >> 12)
        })
        .map(|num| PhysFrame::from_start_address(PhysAddr::new(num << 12)).unwrap());

    #[allow(clippy::cast_precision_loss)]
    let mut mem_size = region_iter
        .map(|region| {
            4096 * ((region.end >> 12)
                - ((region.start >> 12) + if region.start & 0xFFF > 0 { 1 } else { 0 }))
        })
        .sum::<u64>() as f64;
    let mut prefix = 0;
    while mem_size >= 1024.0 {
        mem_size /= 1024.0;
        prefix += 1;
    }
    let prefix = ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei"][prefix];
    println!("[PMM] Initialized, found {:.2} {}B of memory", mem_size, prefix);
    Mutex::new(FrameManager { stack: FrameStack(None), iter: frame_iter })
});