#include "object_pool.h"

#include "../memory/memory.h"

typedef union BzObjectIdx {
    struct {
        u16 idx;
        u16 page;
    };
    u32 value;
} BzObjectIdx;

#define INVALID_OBJECT_IDX (UINT32_MAX)

static_assert(sizeof(BzObjectIdx) == 4, "");

typedef struct BzObjectPool {
    void **pages;
    size_t pageCapacity;
    size_t pageSize;
    size_t objectsPerPage;
    size_t stride;

    size_t      numFree;
    BzObjectIdx firstFree;

    BzObjectPoolFunc constructor;
    BzObjectPoolFunc destructor;
} BzObjectPool;

static void initializePage(BzObjectPool *pool, u16 pageIdx) {
    BZ_ASSERT(pageIdx < UINT8_MAX);
    // Link free list
    for (size_t i = 0; i < pool->objectsPerPage - 1; i++) {
        BzObjectIdx idx  = {.page = pageIdx, .idx = i};
        BzObjectIdx *obj = bzObjectPoolGetObject(pool, idx.value);

        BzObjectIdx nextIdx = {.page = pageIdx, .idx = i + 1};
        obj->value = nextIdx.value;
    }
    BzObjectIdx  lastIdx = {.page = pageIdx, .idx = pool->objectsPerPage - 1};
    BzObjectIdx *lastObj = bzObjectPoolGetObject(pool, lastIdx.value);
    lastObj->value       = pool->firstFree.value;
    pool->firstFree      = (BzObjectIdx) {.page = pageIdx, .idx = 0};

    pool->numFree += pool->objectsPerPage;
}
static void addNewPage(BzObjectPool *pool) {
    if (pool->pageSize >= pool->pageCapacity) {
        size_t newCapacity = pool->pageCapacity << 1;
        void **newPages    = bzResize(pool->pages, sizeof(*newPages) * newCapacity);
        pool->pages        = newPages;
        pool->pageCapacity = newCapacity;
    }
    void *newPage = bzAlloc(pool->stride * pool->objectsPerPage);
    BZ_ASSERT(newPage);
    size_t pageIdx = pool->pageSize;
    pool->pages[pageIdx] = newPage;
    pool->pageSize++;
    initializePage(pool, pageIdx);
}

BzObjectPool *bzObjectPoolCreate(const BzObjectPoolDesc *desc) {
    BZ_ASSERT(desc->objectSize > 0);

    // NOTE: Since object bits are used as a free list
    // when not in use, we must ensure they can hold i32.
    size_t stride = desc->objectSize;
    if (stride < sizeof(i32)) {
        stride = sizeof(i32);
    }
    BZ_ASSERT(desc->objectsPerPage >= 0);
    BZ_ASSERT(desc->objectsPerPage < (2 << 23));

    size_t objectsPerPage = desc->objectsPerPage;
    if (objectsPerPage == 0) objectsPerPage = 512;

    size_t pageCapacity = 8;
    void **pages        = bzAlloc(sizeof(*pages) * pageCapacity);

    BzObjectPool *pool = bzAlloc(sizeof(*pool));
    *pool = (BzObjectPool) {
        .pages = pages,
        .pageSize = 0,
        .pageCapacity = pageCapacity,
        .objectsPerPage = objectsPerPage,
        .stride = stride,
    };
    pool->firstFree.value = INVALID_OBJECT_IDX;

    addNewPage(pool);
    return pool;
}
void bzObjectPoolDestroy(BzObjectPool *pool) {
    for (size_t i = 0; i < pool->pageSize; i++) {
        bzFree(pool->pages[i]);
        pool->pages[i] = NULL;
    }
    bzFree(pool->pages);
    pool->pages = NULL;
    bzFree(pool);
}

size_t bzObjectPoolGetObjectSize(BzObjectPool *pool) {
    return pool->stride;
}

size_t bzObjectPoolGetNumFree(BzObjectPool *pool) {
    return pool->numFree;
}
void *bzObjectPool(BzObjectPool *pool) {
    if (pool->firstFree.value == INVALID_OBJECT_IDX) {
        addNewPage(pool);
        BZ_ASSERT(pool->firstFree.value !=  INVALID_OBJECT_IDX);
    }
    BzObjectIdx *object = (BzObjectIdx *) bzObjectPoolGetObject(pool, pool->firstFree.value);
    pool->firstFree     = *object;

    BZ_ASSERT(pool->numFree > 0);
    pool->numFree--;
    return object;
}
void *bzObjectPoolGetObject(BzObjectPool *pool, u32 idx) {
    BzObjectIdx objectIdx = {.value = idx};
    BZ_ASSERT(objectIdx.page >= 0 && objectIdx.page < pool->pageSize);
    BZ_ASSERT(objectIdx.idx >= 0 && objectIdx.idx < pool->objectsPerPage);

    void *page = pool->pages[objectIdx.page];
    return (void *) ((u8 *) page + objectIdx.idx * pool->stride);
}
u32 bzObjectPoolGetIdx(BzObjectPool *pool, void *object) {
    u16 pageIdx = pool->pageSize;
    u16 idx     = 0;

    for (size_t i = 0; i < pool->pageSize; i++) {
        void *start = pool->pages[i];
        void *end   = (u8 *) start + pool->objectsPerPage * pool->stride;

        if (object >= start && object < end) {
            pageIdx = i;
            idx = (size_t) object - (size_t) start;
            idx = idx / pool->stride;
            break;
        }
    }
    BZ_ASSERT(pageIdx < pool->pageSize);
    BzObjectIdx objectIdx = {.page = pageIdx, .idx = idx};
    return objectIdx.value;
}
void bzObjectPoolRelease(BzObjectPool *pool, void *object) {
    BzObjectIdx objectIdx = {.value = bzObjectPoolGetIdx(pool, object)};

    void *page = pool->pages[objectIdx.page];
    *(BzObjectIdx *) ((u8 *) page + objectIdx.idx * pool->stride) = pool->firstFree;
    pool->firstFree = objectIdx;
    pool->numFree++;
}