#include "slab.h"
#include "bump.h"
#include "memory.h"
#include "string.h"
#include "jove.h"
#include "print.h"
#include "klib/format.h"

static int
s_get_free_listw(size_t slabw, size_t objw)
{
    int freelistc = 1;
    while(freelistc < 256) {
        int maxobjc = (slabw - (freelistc * sizeof(uintptr_t))) / objw;
        if(maxobjc <= freelistc) return maxobjc;
        freelistc++;
    }
    return freelistc;
}

static slab_t
*s_slab_new(slab_cache_t *cache, slab_t *last)
{
    size_t slab_width = (cache->slab_pages << PAGE_SHIFT);
    uintptr_t descr_base = (uintptr_t)bump_alloc(slab_width);
    slab_t *descr = (slab_t*)descr_base;

    size_t free_listc = s_get_free_listw(
            slab_width - sizeof(slab_t),
            cache->obj_size);
    size_t descriptor_width = sizeof(slab_t)
        + (free_listc * sizeof(uintptr_t));
    uintptr_t obj_base = descr_base + descriptor_width;

    if(free_listc < 8) {
        free_listc = ((slab_width - sizeof(slab_t)) / cache->obj_size);
        descr = kmalloc(sizeof(slab_t) + (free_listc * sizeof(uintptr_t)));
        obj_base = descr_base;
    }
    cache->obj_capacity += free_listc;

    *descr = (slab_t) {
        .prev = last,
        .next = (last == NULL ? NULL : last->next),
        .slab_base = (void*)descr_base,
        .obj_base = (void*)obj_base,
        .free_count = free_listc,
        .free_index = free_listc - 1
    };
    for(size_t i = 0; i < free_listc; i++) {
        descr->free[i] = obj_base + (i * cache->obj_size);
    }

    return descr;
}

void
slabcache_new(slab_cache_t *cache, char *name, size_t objsize)
{
    if(objsize % 8 > 0) objsize += (8 - (objsize % 8));
    size_t pages = objsize > 512 ? (objsize >> 9) : 1;
    *cache = (slab_cache_t){
        .obj_size = objsize,
        .slab_pages = pages,
        .list_free = NULL,
        .list_partial = NULL,
        .list_full = NULL
    };
    size_t namelen = strlen(name);
    namelen = namelen > 32 ? 32 : namelen;
    memcpy(cache->name, name, namelen);

    //Allocate the first slab
    cache->list_free = s_slab_new(cache, NULL);
}

void*
slab_alloc(slab_cache_t *cache)
{
    // Get a free slab
    slab_t *slab = NULL;
    if(cache->list_partial != NULL) slab = cache->list_partial;
    if(slab == NULL && cache->list_free != NULL) {
        slab = cache->list_free;
        cache->list_free = slab->next;
    }
    if(slab == NULL) slab = s_slab_new(cache, cache->list_free);
    cache->list_partial = slab;

    // Take an object from the slab.
    uintptr_t objaddr = slab->free[slab->free_index];
    slab->free_index -= 1;

    if(slab->free_index < 0) {
        slab->next = cache->list_full;
        cache->list_full = slab;
    }
    return (void*)objaddr;
}

void
slab_free(slab_cache_t *cache, void *ptr)
{
    uintptr_t addr = (uintptr_t)ptr;
    //Look for the pointer in the bounds of every slab
    for(slab_t *slab = cache->list_full; 
            slab != NULL; slab = slab->next)
    {
        uintptr_t base = (uintptr_t)slab->obj_base;
        uintptr_t limit = ((uintptr_t)slab->slab_base) 
            + (cache->slab_pages << PAGE_SHIFT);
        if(addr > limit || addr < base) continue;
        if((addr - base) % cache->obj_size != 0) {
            klogf("Tried to free offset pointer %#016X in slab %s\n",
                    addr, cache->name);
            return;
        }
        slab->free_index++;
        slab->free[slab->free_index] = addr;

        cache->list_full = slab->next;
        slab->next = cache->list_partial;
        cache->list_partial = slab;
        return;
    }
    for(slab_t *slab = cache->list_partial; 
            slab != NULL; slab = slab->next)
    {
        uintptr_t base = (uintptr_t)slab->obj_base;
        uintptr_t limit = ((uintptr_t)slab->slab_base) 
            + (cache->slab_pages << PAGE_SHIFT);
        if(addr > limit || addr < base) continue;
        if((addr - base) % cache->obj_size != 0) {
            klogf("Tried to free offset pointer %#016X in slab %s\n",
                    addr, cache->name);
            return;
        }
        slab->free_index++;
        slab->free[slab->free_index] = addr;

        if(slab->free_index == ((int)slab->free_count - 1)) {
            cache->list_partial = slab->next;
            slab->next = cache->list_free;
            cache->list_free = slab;
        }
        return;
    }
}