#include <mm/slab.h>
#include <stddef.h>
#include <kprint.h>
#include <kmath.h>
#include <string.h>
struct ma_cache *kmalloc_caches[14] = {0};

// Create various sizes of caches to be used by kmalloc
void _kmalloc_init(void){
    kmalloc_caches[0] = ma_cache_create("kmalloc16",   16,     0, NULL, NULL);
    kmalloc_caches[1] = ma_cache_create("kmalloc32",   32,     0, NULL, NULL);
    kmalloc_caches[2] = ma_cache_create("kmalloc64",   64,     0, NULL, NULL);
    kmalloc_caches[3] = ma_cache_create("kmalloc128",  128,    0, NULL, NULL);
    kmalloc_caches[4] = ma_cache_create("kmalloc256",  256,    0, NULL, NULL);
    kmalloc_caches[5] = ma_cache_create("kmalloc512",  512,    0, NULL, NULL);
    kmalloc_caches[6] = ma_cache_create("kmalloc1K",   1024,   0, NULL, NULL);
    kmalloc_caches[7] = ma_cache_create("kmalloc4K",   4096,   0, NULL, NULL);
    kmalloc_caches[8] = ma_cache_create("kmalloc8K",   8192,   0, NULL, NULL);
    kmalloc_caches[9] = ma_cache_create("kmalloc32K",  32768,  0, NULL, NULL);
    kmalloc_caches[10] = ma_cache_create("kmalloc64K",  65536,  0, NULL, NULL);
    kmalloc_caches[11] = ma_cache_create("kmalloc131K", 131072, 0, NULL, NULL);
    kmalloc_caches[12] = ma_cache_create("kmalloc524K", 524288, 0, NULL, NULL);
    kmalloc_caches[13] = ma_cache_create("kmalloc1M", 1048576, 0, NULL, NULL);
}

size_t sizes[14] = {16, 32, 64, 128, 256, 512, 1024, 4096, 8192, 32768, 65536, 131072, 524288, 1048756};

void *kmalloc(size_t size){
    if(size > 1048576){
        klog(__func__, "Attempted to allocate more than max size (1M)");
        return NULL;
    }

    void *addr = NULL;
    for(int i = 0; i < 14; i++){
        if(sizes[i] >= size){
            addr = ma_cache_alloc(kmalloc_caches[i], 0);
            break;
        }
    }
    
    if(addr == NULL){
        klog(__func__, "Failed to allocate heap memory!");
    }

    return addr;

}

void *kzalloc(size_t size){
    void *addr = kmalloc(size);
    if(addr == NULL){
        return NULL;
    }
    memset(addr, 0, size);
    return addr;
}

kstatus kfree(void *addr){
    return ma_cache_dealloc(addr);
}