use crate::print;
use core::sync::atomic::{AtomicUsize, Ordering};
use atomalloc::noalloc::AtomAlloc;
use limine::memory_map::{Entry, EntryType};
use spin::{Mutex, Once};
use x86_64::PhysAddr;
use x86_64::structures::paging::{FrameAllocator, FrameDeallocator, PhysFrame, Size4KiB};
use crate::println;

static FRAME_ALLOCATOR: Once<Mutex<PFrameAllocator>> = Once::new();

pub const MAX_ATOMS: usize = 1024;
static ATOMS: Mutex<AtomAlloc<MAX_ATOMS, 4096>> = Mutex::new(AtomAlloc::new());


static TOTAL_MEMORY: AtomicUsize = AtomicUsize::new(0);

#[repr(C)]
#[derive(PartialEq, Eq, Clone, Copy)]
pub struct Region {
    /// The base of the memory region, in *physical space*.
    pub base: u64,
    /// The end address of the memory region, in *physical space*
    pub end: u64,
    /// The length of the memory region, in bytes.
    pub length: u64,
    /// The type of the memory region. See [`EntryType`] for specific values.
    pub entry_type: EntryType,
}

pub struct PFrameAllocator {}

impl PFrameAllocator {
    pub unsafe fn init(memory_map: &[&Entry]) {
        let mut total_memory: usize = 0;

        // lock atoms so we can add to it
        let mut atoms = ATOMS.lock();

        for e in memory_map {
            if e.entry_type == EntryType::USABLE {
                total_memory += e.length as usize;

                atoms.add_memory(e.base as usize, e.length.div_ceil(4096) as usize);
            }
        }

        // we dont need it anymore, drop
        drop(atoms);

        TOTAL_MEMORY.store(total_memory, Ordering::Relaxed);

        FRAME_ALLOCATOR.call_once(|| {
            Mutex::new(PFrameAllocator {})
        });
    }

    pub fn allocate_single_frame(&mut self) -> Option<PhysFrame<Size4KiB>> {
        let mut atoms = ATOMS.lock();

        let frame = if let Some(frame) = atoms.allocate(1) {
            frame
        } else {
            // out of memory? lets defrag first
            atoms.defragment();

            if let Some(frame) = atoms.allocate(1) {
                frame
            } else {
                println!("Out of memory!");

                return None;
            }
        };

        drop(atoms);

        PhysFrame::from_start_address(PhysAddr::new(frame as u64)).ok()
    }

    pub fn deallocate_single_frame(&mut self, frame: PhysFrame<Size4KiB>) {
        ATOMS.lock().deallocate(frame.start_address().as_u64() as usize, 1)
    }

    pub fn total_memory() -> usize {
        TOTAL_MEMORY.load(Ordering::Relaxed)
    }
}

unsafe impl FrameAllocator<Size4KiB> for PFrameAllocator {
    fn allocate_frame(&mut self) -> Option<PhysFrame<Size4KiB>> {
        self.allocate_single_frame()
    }
}

impl FrameDeallocator<Size4KiB> for PFrameAllocator {
    unsafe fn deallocate_frame(&mut self, frame: PhysFrame<Size4KiB>) {
        self.deallocate_single_frame(frame);
    }
}

pub fn get_frame_allocator() -> &'static Mutex<PFrameAllocator> {
    FRAME_ALLOCATOR.get().expect("Cannot get frame allocator?")
}