#define pr_fmt(fmt) "[alloc]: " fmt
//#define DEBUG
#include "alloc.h"
#include "allocArea.h"
#include "assert.h"
#include "errno.h"
#include "irq.h"
#include "kernel.h"
#include "klibc.h"
#include "list.h"
#include "math.h"
#include "mem.h"
#include "paging.h"

#define IS_SELF_CONTAINED(desc) ((vaddr_t)((desc)->page) == (vaddr_t)(desc))
static struct slabDesc *slub;
static int allocInitialized = FALSE;

static int allocSlab(struct slabDesc **desc, size_t sizeEl, size_t sizeSlab,
                     int self_containing, int neverEmpty);
static int allocSlabEntry(struct slabEntry **desc, size_t sizeEl, size_t sizeSlab,
                          int selfContained);
static int formatPage(struct slabEntry *desc, size_t size, size_t sizeSlab, int selfContained);
static int freeFromSlab(void *ptr, struct slabEntry *slab);
static struct slabDesc *allocGetSlab(size_t size);

static struct {
    size_t elementSize;
    size_t slabSize;
    unsigned char isSelf;
    unsigned char neverEmpty;
} initSlab[] = {{4, PAGE_SIZE, 0, 0},
                {8, PAGE_SIZE, 0, 0},
                {16, PAGE_SIZE, 0, 0},
                {32, PAGE_SIZE, 0, 0},
                {64, PAGE_SIZE, 0, 0},
                {128, PAGE_SIZE, 0, 0},
                {256, 2 * PAGE_SIZE, 0, 0},
                {1024, 2 * PAGE_SIZE, 0, 0},
                {2048, 3 * PAGE_SIZE, 0, 0},
                {4096, 4 * PAGE_SIZE, 0, 0},
                {0, 0, 0, 0}};

int allocSetup(size_t sizeOfArea, vaddr_t *areaAddr, vaddr_t *descAddr, vaddr_t *entryAddr)
{
    list_init(slub);

    assert(allocBookSlab(sizeof(struct slabDesc), PAGE_SIZE, TRUE, FALSE) == 0);
    *descAddr = (vaddr_t)allocGetSlab(sizeof(struct slabDesc));
    assert(allocBookSlab(sizeof(struct slabEntry), PAGE_SIZE, TRUE, FALSE) == 0);
    *entryAddr = (vaddr_t)allocGetSlab(sizeof(struct slabEntry));
    assert(allocBookSlab(sizeOfArea, PAGE_SIZE, TRUE, TRUE) == 0);
    *areaAddr        = (vaddr_t)allocGetSlab(sizeOfArea);
    allocInitialized = TRUE;

    return 0;
}

int allocPopulate()
{
    for (uint i = 0; initSlab[i].elementSize != 0; i++) {
        int ret;

        if ((ret = allocBookSlab(initSlab[i].elementSize, initSlab[i].slabSize,
                                 initSlab[i].isSelf, initSlab[i].neverEmpty))) {
            if (ret == -EEXIST)
                continue;
            pr_err("Fail to allocBookSlab %d for %u \n", ret, (1U << i));

            return ret;
        }
    }

    return 0;
}

int allocBookSlab(size_t sizeEl, size_t sizeSlab, int selfContained, int neverEmpty)
{
    struct slabDesc *slab    = NULL;
    struct slabDesc *newSlab = NULL;
    int slabIdx;
    int ret;
    int flags;

    pr_devel("%s for element of size %lu is self %d\n", __func__, sizeEl, selfContained);

    disable_IRQs(flags);
    list_foreach(slub, slab, slabIdx)
    {
        if (slab->size == sizeEl) {
            restore_IRQs(flags);
            return -EEXIST;
        }
        if (slab->size > sizeEl) {
            break;
        }
    }

    if ((ret = allocSlab(&newSlab, sizeEl, sizeSlab, selfContained, neverEmpty))) {
        pr_devel("Failed to alloc Slab\n");
        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;
}

static int allocSlab(struct slabDesc **desc, size_t size, size_t sizeSlab, int selfContained,
                     int neverEmpty)
{
    uint nbPage, i;
    vaddr_t alloc;

    pr_devel("%s for size %d is self %d\n", __func__, size, selfContained);

    sizeSlab = MAX(sizeSlab, PAGE_SIZE);
    if (size > sizeSlab) {
        pr_devel("size of element %d are bigger than slab size %d\n", size, sizeSlab);
        return -ENOENT;
    }

    nbPage = DIV_ROUND_UP(sizeSlab, PAGE_SIZE);
    if (allocInitialized) {
        alloc = areaAlloc(nbPage, AREA_PHY_MAP);
        if (alloc == (paddr_t)NULL)
            return -ENOMEM;
    } else {
        alloc = (vaddr_t)allocPhyPage(nbPage);
        if (alloc == (paddr_t)NULL)
            return -ENOMEM;

        for (i = 0; i < nbPage; i++) {
            if (pageMap(alloc + i * PAGE_SIZE, alloc + i * PAGE_SIZE, PAGING_MEM_WRITE))
                goto free_page;
        }
    }

    if (selfContained) {
        *desc                  = (struct slabDesc *)alloc;
        ((*desc)->slab).freeEl = (char *)(*desc) + sizeof(struct slabDesc);
    } else {
        *desc = malloc(sizeof(struct slabDesc));
        if (*desc == NULL)
            return -ENOMEM;
        (*desc)->slab.freeEl = (void *)alloc;
    }

    struct slabEntry *slab = &(*desc)->slab;
    list_singleton(slab, slab);
    slab->page          = (vaddr_t)alloc;
    slab->full          = 0;
    slab->size          = sizeSlab;
    (*desc)->neverEmpty = neverEmpty;
    (*desc)->size       = size;

    return formatPage(&(*desc)->slab, size, sizeSlab, selfContained);

free_page:
    for (uint j = 0; j < i; j++) {
        pageUnmap((vaddr_t)alloc + i * PAGE_SIZE);
    }

    return -ENOMEM;
}

static int allocSlabEntry(struct slabEntry **desc, size_t size, size_t sizeSlab,
                          int selfContained)
{
    uint nbPage;

    pr_devel("%s for size %d is self %d\n", __func__, size, selfContained);

    sizeSlab = MAX(sizeSlab, PAGE_SIZE);
    if (size > sizeSlab) {
        pr_devel("size of element %d are bigger than slab size %d\n", size, sizeSlab);
        return -ENOENT;
    }

    nbPage        = DIV_ROUND_UP(sizeSlab, PAGE_SIZE);
    vaddr_t alloc = areaAlloc(nbPage, AREA_PHY_MAP);
    if (alloc == (paddr_t)NULL)
        return -ENOMEM;

    if (selfContained) {
        *desc           = (struct slabEntry *)alloc;
        (*desc)->freeEl = (char *)(*desc) + sizeof(struct slabEntry);
    } else {
        *desc = malloc(sizeof(struct slabEntry));
        if (*desc == NULL)
            return -ENOMEM;
        (*desc)->freeEl = (void *)alloc;
    }

    list_singleton(*desc, *desc);
    (*desc)->page = (vaddr_t)alloc;
    (*desc)->full = 0;
    (*desc)->size = sizeSlab;

    return formatPage((*desc), size, sizeSlab, selfContained);
}

static int formatPage(struct slabEntry *desc, size_t size, size_t sizeSlab, int selfContained)
{
    char *cur  = desc->freeEl;
    ulong nbEl = sizeSlab / size - 1;

    if (selfContained)
        nbEl = (sizeSlab - sizeof(struct slabDesc)) / size - 1;

    for (ulong i = 0; i < nbEl; i++) {
        *((vaddr_t *)cur) = (vaddr_t)cur + size;
        cur += size;
    }

    *((vaddr_t *)cur) = (vaddr_t)NULL;

    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;
    }
    slab->freeEl = (void *)(*next);

    return (void *)next;
}

static struct slabDesc *allocGetSlab(size_t size)
{
    struct slabDesc *slab = NULL;
    uint slubIdx;

    list_foreach(slub, slab, slubIdx)
    {
        if (size <= slab->size)
            return slab;
    }

    return NULL;
}

void *malloc(size_t size)
{
    struct slabEntry *slabEntry;
    struct slabDesc *slab = NULL;
    void *ret;
    int flags;
    int slabIdx;

    disable_IRQs(flags);

    if (size >= PAGE_SIZE) {
        vaddr_t area = areaAlloc(DIV_ROUND_UP(size, PAGE_SIZE), AREA_PHY_MAP);

        return (void *)area;
    }
    if ((slab = allocGetSlab(size)) == NULL) {
        pr_devel("No slab found for %d\n", size);
        return NULL;
    }

    struct slabEntry *slabList = &slab->slab;
    list_foreach(&slab->slab, slabEntry, slabIdx)
    {
        if (!slabEntry->full) {
            // pr_devel("found place in slab idx %d for size %d\n", slabIdx, slab->size);
            ret = allocFromSlab(slabEntry);
            if (slabEntry->full && slab->neverEmpty) {

                pr_devel("Prealloc place for slab for object size %d\n", slab->size);
                assert(IS_SELF_CONTAINED(slabEntry));
                size_t slabSize         = slabEntry->size;
                int nbPages             = DIV_ROUND_UP(slabSize, PAGE_SIZE);
                struct slabEntry *entry = (struct slabEntry *)areaBook(nbPages, AREA_PHY_MAP);
                if (entry == NULL) {
                    restore_IRQs(flags);

                    return NULL;
                }

                (entry)->freeEl = (char *)(entry) + sizeof(struct slabEntry);
                list_singleton(entry, entry);
                entry->page = (vaddr_t)entry;
                entry->full = 0;
                entry->size = slabSize;

                formatPage(entry, slab->size, slabSize, 1);
                list_add_tail(slabList, entry);
                areaAdd((vaddr_t)entry, (vaddr_t)entry + nbPages * PAGE_SIZE, FALSE);
            }

            restore_IRQs(flags);
            return ret;
        }
    }

    // No room found
    struct slabEntry *newSlabEntry;
    size_t slabSize = MAX(PAGE_SIZE, size);
    int retSlab;

    if ((retSlab = allocSlabEntry(&newSlabEntry, slab->size, slabSize,
                                  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;
}

void *zalloc(size_t size)
{
    void *alloc = malloc(size);

    if (alloc != NULL)
        memset(alloc, 0, size);

    return alloc;
}

static 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 + slabEntry->size))) {

            *((vaddr_t *)ptr) = (vaddr_t)slabEntry->freeEl;
            slabEntry->freeEl = ptr;
            slabEntry->full   = 0;
            return 1;
        }
    }

    return 0;
}

int freeSlabAllocated(void *ptr)
{
    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 0;
            }
        }
    }
    restore_IRQs(flags);

    return -1;
}

void free(void *ptr)
{
    if (!ptr)
        return;
    if (!freeSlabAllocated(ptr))
        return;
    if (areaFree((vaddr_t)ptr))
        pr_err("free: cannot found origin\n");
}