extern crate alloc;
use core::alloc::{GlobalAlloc, Layout};
use alloc::collections::BTreeMap;
use crate::ALLOCATOR;

pub unsafe fn init(start: usize, size: usize) {
    ALLOCATOR.init(start, size);
}

/*
#[cfg(not(alloc_trace))]
pub unsafe fn alloc(layout: Layout) -> *mut u8 {
    ALLOCATOR.alloc(layout)
}
*/

pub unsafe fn alloc(layout: Layout) -> *mut u8 {
    let ptr = ALLOCATOR.alloc(layout);
    #[cfg(alloc_trace)]
    {
        crate::console::get_global_console().write("+");
        crate::util::print_hex_padded(ptr as u32, 8);
        crate::console::get_global_console().write(" ");
    }
    ptr
}

pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) {
    #[cfg(alloc_trace)]
    {
        crate::console::get_global_console().write("-");
        crate::util::print_hex_padded(ptr as u32, 8);
        crate::console::get_global_console().write(" ");
    }
    ALLOCATOR.dealloc(ptr, layout);
}

pub unsafe fn realloc(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
    #[cfg(alloc_trace)]
    {
        crate::console::get_global_console().write("=");
        crate::util::print_hex_padded(ptr as u32, 8);
    }
    let ptr = ALLOCATOR.realloc(ptr, layout, new_size);
    #[cfg(alloc_trace)]
    {
        crate::console::get_global_console().write("/");
        crate::util::print_hex_padded(ptr as u32, 8);
        crate::console::get_global_console().write(" ");
    }
    ptr
}

pub struct RustAllocOperations {
    pub init: unsafe fn(start: usize, size: usize),
    pub alloc: unsafe fn(layout: Layout) -> *mut u8,
    pub dealloc: unsafe fn(ptr: *mut u8, layout: Layout),
    pub realloc: unsafe fn(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8,
}

pub const RUST_ALLOC_OPERATIONS: RustAllocOperations = RustAllocOperations {
    init,
    alloc,
    dealloc,
    realloc,
};

pub unsafe extern "C" fn c_init(start: usize, size: usize) {
    init(start, size)
}

pub unsafe extern "C" fn c_alloc(size: usize, align: usize) -> *mut u8 {
    let layout = match Layout::from_size_align(size, align) {
        Ok(l) => l,
        Err(_) => return 0 as *mut u8,
    };

    alloc(layout)
}

pub unsafe extern "C" fn c_dealloc(ptr: *mut u8, size: usize, align: usize) {
    let layout = match Layout::from_size_align(size, align) {
        Ok(l) => l,
        Err(_) => return,
    };

    dealloc(ptr, layout)
}

#[link_section = ".kgbss"]
static mut MALLOC_MAP: Option<BTreeMap<*mut u8, usize>> = None;

pub unsafe extern "C" fn malloc(size: usize) -> *mut u8 {
    let layout = match Layout::from_size_align(size, 4) {
        Ok(l) => l,
        Err(_) => return 0 as *mut u8,
    };

    if MALLOC_MAP.is_none() {
        MALLOC_MAP = Some(BTreeMap::new());
    }

    let ptr = alloc(layout);
    if ptr == (0 as *mut u8) {
        return 0 as *mut u8;
    }

    let map = MALLOC_MAP.as_mut().unwrap();
    map.insert(ptr, size);

    ptr
}

pub unsafe extern "C" fn free(ptr: *mut u8) {
    if ptr.is_null() {
        return;
    }

    if MALLOC_MAP.is_none() {
        return;
    }

    let map = MALLOC_MAP.as_mut().unwrap();
    if map.contains_key(&ptr) {
        let size = *map.get(&ptr).unwrap();
        let layout = match Layout::from_size_align(size, 4) {
            Ok(l) => l,
            Err(_) => return,
        };
        dealloc(ptr, layout);
        map.remove(&ptr);
    }
}

pub unsafe extern "C" fn c_realloc(ptr: *mut u8, nbytes: usize) -> *mut u8 {
    if MALLOC_MAP.is_none() {
        MALLOC_MAP = Some(BTreeMap::new());
    }

    if ptr as u32 == 0 {
        return malloc(nbytes);
    } else if nbytes == 0 {
        free(ptr);
        return 0 as *mut u8;
    }

    let map = MALLOC_MAP.as_mut().unwrap();
    if map.contains_key(&ptr) {
        let size = *map.get(&ptr).unwrap();
        let layout = match Layout::from_size_align(size, 4) {
            Ok(l) => l,
            Err(_) => return 0 as *mut u8,
        };
        let new_ptr = realloc(ptr, layout, nbytes);
        map.remove(&ptr);
        map.insert(new_ptr, nbytes);
        new_ptr
    } else {
        0 as *mut u8
    }
}

#[repr(C)]
pub struct CAllocOperations {
    pub init: unsafe extern "C" fn(start: usize, size: usize),
    pub alloc: unsafe extern "C" fn(size: usize, align: usize) -> *mut u8,
    pub dealloc: unsafe extern "C" fn(ptr: *mut u8, size: usize, align: usize),
    pub malloc: unsafe extern "C" fn(size: usize) -> *mut u8,
    pub free: unsafe extern "C" fn(ptr: *mut u8),
    pub realloc: unsafe extern "C" fn(ptr: *mut u8, nbytes: usize) -> *mut u8,
}

pub const C_ALLOC_OPERATIONS: CAllocOperations = CAllocOperations {
    init: c_init,
    alloc: c_alloc,
    dealloc: c_dealloc,
    malloc,
    free,
    realloc: c_realloc,
};

#[alloc_error_handler]
pub fn handle_alloc_error(_layout: Layout) -> ! {
    crate::console::get_global_console().write("OOM!\r\n");
    loop{}
}