#define pr_fmt(fmt) "[alloc]: " fmt
//#define DEBUG
#include "alloc.h"
#include "errno.h"
#include "irq.h"
#include "klibc.h"
#include "list.h"
#include "math.h"
#include "mem.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)
static struct slabDesc *slub;

int allocSlab(struct slabDesc **desc, size_t size, int self_containing);
int allocSlabEntry(struct slabEntry **desc, size_t size, int selfContained);
static int formatPage(struct slabEntry *desc, size_t size, int selfContained);

int allocSetup(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))) {
            if (ret == -EEXIST)
                continue;
            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;
    int flags;

    disable_IRQs(flags);
    list_foreach(slub, slab, slabIdx)
    {
        if (slab->size == size) {
            restore_IRQs(flags);
            return -EEXIST;
        }
        if (slab->size > size) {
            break;
        }
    }
    struct slabDesc *newSlab;
    if ((ret = allocSlab(&newSlab, size, selfContained))) {
        restore_IRQs(flags);
        return ret;
    }
    if (list_foreach_early_break(slub, slab, slabIdx)) {
        list_insert_before(slub, slab, newSlab);
    } else {
        list_add_tail(slub, newSlab);
    }
    restore_IRQs(flags);
    return 0;
}

int allocSlab(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;
    (*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 allocSlabEntry(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;
    // 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 is now full\n", slab);
        slab->full = 1;
    } else {
        slab->freeEl = (void *)(*next);
    }
    return (void *)next;
}

void *malloc(size_t size)
{
    int flags;
    void *ret;

    if (size > (1U << SLUB_SIZE)) {
        printf("implement malloc for big size\n");
        return NULL;
    }
    struct slabDesc *slab;
    uint slubIdx;

    disable_IRQs(flags);
    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);
            ret = allocFromSlab(slabEntry);
            restore_IRQs(flags);
            return ret;
        }
    }

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

int freeFromSlab(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);
            if (slabEntry->full) {
                *((vaddr_t *)ptr) = (vaddr_t)NULL;
            } else {
                *((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;
    int flags;

    disable_IRQs(flags);
    list_foreach(slub, slab, slabIdx)
    {
        struct slabEntry *slabEntry;
        int entryIdx;
        list_foreach(&slab->slab, slabEntry, entryIdx)
        {
            if (freeFromSlab(ptr, slabEntry)) {
                restore_IRQs(flags);
                return;
            }
        }
    }
    restore_IRQs(flags);
    pr_devel("free: slab not found\n");
}