#define pr_fmt(fmt) "[alloc]: " fmt
//#define DEBUG
#include "alloc.h"
#include "errno.h"
#include "list.h"
#include "math.h"
#include "mem.h"
#include "vga.h"

#define IS_SELF_CONTAINED(desc) ((vaddr_t)((desc)->page) == (vaddr_t)(desc))
// Slab will contains object from sizeof(void *) to PAGE_SIZE/2 by pow2
#define SLUB_SIZE (PAGE_SHIFT - 1)
struct slabDesc *slub;

int addSlab(struct slabDesc **desc, size_t size, int self_containing);
static int formatPage(struct slabDesc *desc, size_t size, int selfContained);

int allocInit(void)
{
        uint start = log2(sizeof(void *));
        list_init(slub);
        int ret;
        if ((ret = allocBookSlab(sizeof(struct slabDesc), 1))) {
                pr_devel("Fail to allocBookSlab %d for slabDesc :( \n");
                return ret;
        }
        for (uint i = start; i < SLUB_SIZE; i++) {
                if ((ret = allocBookSlab(1U << i, 0))) {
                        pr_devel("Fail to allocBookSlab %d for %d \n", ret, (1U << i));
                        return ret;
                }
        }
        return 0;
}

int allocBookSlab(size_t size, int selfContained)
{
        // pr_devel("%s for size %d is self %d\n", __func__, size, selfContained );
        struct slabDesc *slabEntry;
        struct slabDesc *slab;
        int slabIdx;
        int ret;
        if ((ret = addSlab(&slabEntry, size, selfContained)))
                return ret;
        list_foreach(slub, slab, slabIdx)
        {
                if (slab->size == size) {
                        pr_devel("Slab already exist");
                        return -EEXIST;
                }
                if (slab->size > size) {
                        list_insert_before(slub, slab, slabEntry);
                        return 0;
                }
        }
        list_add_tail(slub, slabEntry);
        return 0;
}

int addSlab(struct slabDesc **desc, size_t size, int selfContained)
{
        pr_devel("%s for size %d is self %d\n", __func__, size, selfContained);
        if (size > PAGE_SIZE)
                return -ENOENT;
        paddr_t alloc = allocPhyPage();
        if (alloc == (paddr_t)NULL)
                return -ENOMEM;
        if (pageMap((vaddr_t)alloc, alloc, PAGING_MEM_WRITE))
                return -ENOMEM;

        if (selfContained) {
                *desc           = (struct slabDesc *)alloc;
                (*desc)->freeEl = (char *)(*desc) + sizeof(struct slabDesc);
        } else {
                *desc           = malloc(sizeof(struct slabDesc));
                (*desc)->freeEl = (void *)alloc;
        }
        list_singleton(*desc, *desc);
        (*desc)->page = (vaddr_t)alloc;
        (*desc)->full = 0;
        (*desc)->size = size;
        pr_devel("got page %d for size %d first %d", alloc, size, (*desc)->freeEl);
        return formatPage((*desc), size, selfContained);
}

static int formatPage(struct slabDesc *desc, size_t size, int selfContained)
{
        char *cur  = desc->freeEl;
        ulong nbEl = PAGE_SIZE / size - 1;
        if (selfContained)
                nbEl = (PAGE_SIZE - sizeof(struct slabDesc)) / size - 1;
        ulong i;
        for (i = 0; i < nbEl; i++) {
                *((vaddr_t *)cur) = (vaddr_t)cur + size;
                cur += size;
        }
        *((vaddr_t *)cur) = (vaddr_t)NULL;
        pr_devel("last at %d allocated %d\n", cur, i + 1);
        return 0;
}

static void *allocFromSlab(struct slabDesc *slab)
{
        vaddr_t *next = slab->freeEl;
        if (*next == (vaddr_t)NULL) {
                pr_devel("Slab @%d for %d full\n", slab, slab->size);
                slab->full = 1;
        } else {
                slab->freeEl = (void *)(*next);
        }
        return (void *)next;
}

void *malloc(size_t size)
{
        if (size > (1U << SLUB_SIZE)) {
                printf("implement malloc for big size\n");
                return NULL;
        }
        struct slabDesc *slubEntry;
        uint slubIdx;
        list_foreach(slub, slubEntry, slubIdx)
        {
                if (size <= slubEntry->size)
                        break;
        }
        struct slabDesc *slab;
        int slabIdx;
        list_foreach(slubEntry, slab, slabIdx)
        {
                if (!slab->full) {
                        pr_devel("found place in slub %d at idx %d for size %d\n", slubIdx,
                                 slabIdx, size);
                        return allocFromSlab(slab);
                }
        }

        // No room found
        struct slabDesc *newSlab;
        int ret;
        if ((ret = addSlab(&newSlab, slab->size, IS_SELF_CONTAINED(slubEntry)))) {
                pr_devel("Fail to addSlab %d\n", ret);
                return NULL;
        }
        list_add_head(slubEntry, newSlab);
        return allocFromSlab(newSlab);
}

int slabFree(void *ptr, struct slabDesc *slab)
{
        struct slabDesc *slabEntry;
        int slabIdx;
        list_foreach(slab, slabEntry, slabIdx)
        {
                if ((slabEntry->page <= (vaddr_t)ptr) &&
                    ((vaddr_t)ptr < (slabEntry->page + PAGE_SIZE))) {
                        pr_devel("free place! was %d is now %d\n", slabEntry->freeEl, ptr);
                        *((vaddr_t *)ptr) = (vaddr_t)slabEntry->freeEl;
                        slabEntry->freeEl = ptr;
                        slabEntry->full   = 0;
                        return 1;
                }
        }
        return 0;
}
void free(void *ptr)
{
        if (!ptr)
                return;

        struct slabDesc *slab;
        int slabIdx;
        list_foreach(slub, slab, slabIdx)
        {
                if (slabFree(ptr, slab))
                        return;
        }
        pr_devel("free: slab not found\n");
}