#include <assert.h>
#include "ucore/spool.h"
#include "ucore/utils.h"
#include <string.h>

void uc_pool_init(UCPool *pool, void *buffer, size_t buffer_size, size_t block_size, size_t align)
{
    assert(align > 0);
    assert((align & (align - 1)) == 0); // power of 2


    // enforce minimum for storing free list pointers

    if (align < _Alignof(UCPoolBlock)) {
        align = _Alignof(UCPoolBlock);
    }

    if (block_size < sizeof(UCPoolBlock)) {
        block_size = sizeof(UCPoolBlock);
    }
    // ensure initial alignment
    size_t stride = UC_ALIGN(block_size, align);

    // align first usable block
    uintptr_t start_addr = (uintptr_t)buffer;
    uintptr_t first_addr = UC_ALIGN(start_addr, align);
    size_t padding = first_addr - start_addr;

    if (padding + stride > buffer_size) {
        pool->capacity = 0;
        pool->free_list = NULL;
    } else {
        size_t usable = buffer_size - padding;
        pool->capacity = usable / stride;

        // build intrusive free list
        UCPoolBlock *p = (UCPoolBlock*)first_addr;
        for (size_t i = 0; i < pool->capacity - 1; ++i) {
            void *punt = __builtin_assume_aligned(&p->data[0] + pool->block_size, _Alignof(UCPoolBlock));
            UCPoolBlock *next = punt;
            p->next = next;
            p = next;
        }
        p->next = NULL; // last element
        pool->free_list = (UCPoolBlock *)first_addr;
    }

    pool->block_size = stride;
    pool->align      = align;
    pool->start      = buffer;
    pool->end        = buffer + buffer_size;
}


void *uc_pool_alloc(UCPool *pool)
{
    if (!pool->free_list) {
        return NULL;
    }
    UCPoolBlock *block = pool->free_list;
    pool->free_list = block->next;

#ifdef DEBUG
    memset(block, 0xCB, pool->block_size);
#endif
    return block->data;
}
void uc_pool_free(UCPool *pool, void *p)
{
    if (p == NULL) {
        return;
    }
    UCPoolBlock *block = p;
#ifdef DEBUG
    assert(block >= pool->start && block < pool->end);
    assert(((uintptr_t)p - (uintptr_t)pool->start) % pool->block_size == 0);

    // double-free detection
    for (UCPoolBlock *f = pool->free_list; f; f = f->next) {
        if (f == p) {
            assert(0 && "Double free detected in pool!");
            return;
        }
    }

    memset(p, 0xDD, pool->block_size);
#endif
     // When not in use, store pointer to next free element inline in the block
    block->next = pool->free_list;
    pool->free_list = p;
}

void uc_pool_reset(UCPool *pool)
{
   if (pool->capacity == 0)
        return;

    uintptr_t start_addr = (uintptr_t)pool->start;
    uintptr_t first_addr = UC_ALIGN(start_addr, pool->align);

    // build intrusive free list
    UCPoolBlock *p = (UCPoolBlock*)first_addr;
    for (size_t i = 0; i < pool->capacity - 1; ++i) {
        void *punt = __builtin_assume_aligned(&p->data[0] + pool->block_size, _Alignof(UCPoolBlock));
        UCPoolBlock *next = punt;
        p->next = next;
        p = next;

    }

    p->next = NULL;  // last element

    pool->free_list = (UCPoolBlock *)first_addr;
}