//! Kernel's global heap allocator
use core::{
    alloc::{GlobalAlloc, Layout},
    ptr::NonNull,
};

use kernel_arch::KernelTableManagerImpl;
use libk_mm_interface::address::PhysicalAddress;
use libk_util::sync::IrqSafeSpinlock;
use libyalloc::{allocator::BucketAllocator, sys::PageProvider};

use crate::{address::Virtualize, phys, L3_PAGE_SIZE};

// TODO limits?

struct KernelPageProvider;

impl PageProvider for KernelPageProvider {
    fn unmap_pages(address: NonNull<u8>, count: usize) {
        log::trace!("Release {}K of heap", count * 4);
        let phys = PhysicalAddress::from_virtualized(address.as_ptr().addr());
        for i in 0..count {
            unsafe {
                phys::free_page(phys.add(i * L3_PAGE_SIZE));
            }
        }
    }

    fn map_pages(count: usize) -> Option<NonNull<u8>> {
        log::trace!("Grow heap by {}K", count * 4);
        let phys = phys::alloc_pages_contiguous(count).ok()?;
        let virt = phys.raw_virtualize::<KernelTableManagerImpl>();
        Some(unsafe { NonNull::new_unchecked(virt as *mut u8) })
    }
}

/// Kernel heap manager
pub struct KernelAllocator {
    inner: IrqSafeSpinlock<BucketAllocator<KernelPageProvider>>,
}

impl KernelAllocator {
    const fn empty() -> Self {
        Self {
            inner: IrqSafeSpinlock::new(BucketAllocator::new()),
        }
    }
}

unsafe impl GlobalAlloc for KernelAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        match self.inner.lock().allocate(layout) {
            Some(p) => p.as_ptr(),
            None => {
                panic!("Failed to allocate {:?}", layout);
            }
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let ptr = NonNull::new(ptr).unwrap();
        self.inner.lock().free(ptr, layout);
    }
}

unsafe impl Sync for KernelAllocator {}

/// Kernel's global allocator
#[global_allocator]
pub static GLOBAL_HEAP: KernelAllocator = KernelAllocator::empty();
