extern crate alloc;
use alloc::alloc::{GlobalAlloc, Layout};

use lazy_static::lazy_static;

use super::{frame, paging, regions::KERNEL_HEAP_REGION, Allocator, FreeListAllocator};
use crate::prelude::*;

pub const STARTING_SIZE: usize = 64 * MIB;

#[global_allocator]
static GLOBAL_ALLOCATOR: GlobalAllocator = GlobalAllocator;

lazy_static! {
    // Lazily initialized the heap allocator; this will happen on the first call to `kmalloc()`
    static ref HEAP_ALLOCATOR: Mutex<FreeListAllocator> = unsafe {
        info!("Kernel heap located at {}", KERNEL_HEAP_REGION);

        let start = KERNEL_HEAP_REGION.start as *mut u8;

        // Map the heap
        for i in (0..STARTING_SIZE).step_by(paging::L0_PAGE_SIZE) {
            let frame = frame::alloc().expect("failed to allocate frame");

            // We don't need to log every one of these
            quiet! {
                paging::map(
                    start.add(i) as usize,
                    frame as usize,
                    paging::Permissions::ReadWrite,
                    true
                )
                .expect("Failed to map portions of the heap");
            }
        }
        // Initialize the allocator.
        Mutex::new(FreeListAllocator::new(start, STARTING_SIZE))
    };
}

/// Allocate `size` bytes in the kernel heap
pub fn kmalloc(size: usize) -> Option<*mut u8> {
    trace!("Allocate {} bytes", size);
    match HEAP_ALLOCATOR.lock().alloc_bytes(size) {
        Ok(ptr) => Some(ptr),
        Err(why) => {
            error!("Heap allocation error: {}", why);
            None
        }
    }
}

/// Free a heap allocation
///
/// # Safety
///
/// `ptr` must be a pointer returned by `kmalloc()`
pub unsafe fn kfree(ptr: *mut u8) {
    trace!("Free {:p}", ptr);
    if let Err(why) = HEAP_ALLOCATOR.lock().free(ptr) {
        error!("Heap free error: {}", why)
    }
}

struct GlobalAllocator;

unsafe impl GlobalAlloc for GlobalAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        match kmalloc(layout.size()) {
            Some(ptr) => ptr,
            None => ptr::null_mut(),
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, _: Layout) {
        kfree(ptr)
    }
}

#[alloc_error_handler]
fn alloc_error_handler(layout: Layout) -> ! {
    panic!("Global allocator error: size={}", layout.size());
}
