#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 slabEntry *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", ret);
                return ret;
        }
        if ((ret = allocBookSlab(sizeof(struct slabEntry), 1))) {
                pr_devel("Fail to allocBookSlab %d for slabEntry :( \n", ret);
                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 *slab = NULL;
        int slabIdx;
        int ret;
        list_foreach(slub, slab, slabIdx)
        {
                if (slab->size == size) {
                        return -EEXIST;
                }
                if (slab->size > size) {
                        break;
                }
        }
        struct slabDesc *newSlab;
        if ((ret = addSlab(&newSlab, size, selfContained)))
                return ret;
        if (list_foreach_early_break(slub, slab, slabIdx)) {
                list_insert_before(slub, slab, newSlab);
        } else {
                list_add_tail(slub, newSlab);
        }
        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)->slab).freeEl = (char *)(*desc) + sizeof(struct slabDesc);
        } else {
                *desc                = malloc(sizeof(struct slabDesc));
                (*desc)->slab.freeEl = (void *)alloc;
        }
        struct slabEntry *slab = &(*desc)->slab;
        list_singleton(slab, slab);
        slab->page    = (vaddr_t)alloc;
        slab->full    = 0;
        slab->size    = size;
        (*desc)->size = size;
        // pr_devel("got page %d for size %d first %d", alloc, size, (*desc)->slab.freeEl);
        return formatPage(&(*desc)->slab, size, selfContained);
}

int addSlabEntry(struct slabEntry **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 slabEntry *)alloc;
                (*desc)->freeEl = (char *)(*desc) + sizeof(struct slabEntry);
        } else {
                *desc           = malloc(sizeof(struct slabEntry));
                (*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 slabEntry *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 slabEntry *slab)
{
        vaddr_t *next = slab->freeEl;
        if (*next == (vaddr_t)NULL) {
                pr_devel("Slab @%d for %d is now 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 *slab;
        uint slubIdx;
        list_foreach(slub, slab, slubIdx)
        {
                if (size <= slab->size)
                        break;
        }
        struct slabEntry *slabEntry;
        int slabIdx;
        list_foreach(&slab->slab, slabEntry, slabIdx)
        {
                if (!slabEntry->full) {
                        // pr_devel("found place in slub %d at idx %d for size %d\n", slubIdx,
                        //         slabIdx, size);
                        return allocFromSlab(slabEntry);
                }
        }

        // No room found
        struct slabEntry *newSlabEntry;
        struct slabEntry *slabList = &slab->slab;
        int ret;
        if ((ret = addSlabEntry(&newSlabEntry, slab->size, IS_SELF_CONTAINED(&slab->slab)))) {
                pr_devel("Fail to addSlabEntry %d\n", ret);
                return NULL;
        }
        pr_devel("Allocate new slab for object of size %d\n", slab->size);
        list_add_tail(slabList, newSlabEntry);
        return allocFromSlab(newSlabEntry);
}

int slabFree(void *ptr, struct slabEntry *slab)
{
        struct slabEntry *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)
        {
                struct slabEntry *slabEntry;
                int entryIdx;
                list_foreach(&slab->slab, slabEntry, entryIdx)
                {
                        if (slabFree(ptr, slabEntry))
                                return;
                }
        }
        pr_devel("free: slab not found\n");
}