use core::{mem::MaybeUninit, ptr::{NonNull, slice_from_raw_parts_mut}};
use limine::{memmap::MEMMAP_USABLE, request::MemmapResponse};
use talc::{TalcLock, source::Claim};
use x86_64::PhysAddr;
use crate::translate_addr::OffsetMappedPhysAddr;

#[global_allocator]
static GLOBAL_ALLOCATOR: TalcLock<spin::Mutex<()>, Claim> = TalcLock::new(Claim::standby());

pub const GLOBAL_ALLOCATOR_SIZE: u64 = {
    // 4 MiB
    4 * 0x400 * 0x400
};

/// Finds unused physical memory for the global allocator and initializes the global allocator.
/// Returns the start address of the physical memory used for the global allocator.  
///
/// # Safety
/// This function must be called exactly once, and no page tables should be modified before calling this function.
pub unsafe fn init(memory_map: &'static MemmapResponse) -> PhysAddr {
    let global_allocator_physical_start = PhysAddr::new(
        memory_map
            .entries()
            .iter()
            .find(|entry| {
                entry.type_ == MEMMAP_USABLE && entry.length >= GLOBAL_ALLOCATOR_SIZE
            })
            .unwrap()
            .base
    );

    let global_allocator_mem = {
        let mut ptr = NonNull::new(slice_from_raw_parts_mut(
            global_allocator_physical_start
                .offset_mapped()
                .as_mut_ptr::<MaybeUninit<u8>>(),
            GLOBAL_ALLOCATOR_SIZE as usize
        ))
        .unwrap();
        // Safety: We've reserved the physical memory and it is already offset mapped
        unsafe { ptr.as_mut() }
    };

    let mut talc = GLOBAL_ALLOCATOR.lock();
    let ptr = global_allocator_mem.as_mut_ptr().cast::<u8>();
    let len = global_allocator_mem.len();
    // Safety: We got the span from valid memory
    unsafe { talc.claim(ptr, len) }.unwrap();

    global_allocator_physical_start
}