#include <assert.h>
#include <stddef.h>
#include <stdlib.h>
#include <limits.h>
#include <string.h>

struct ucl_arena {
    size_t object_size;
    size_t capacity;
    void *objects;
    int *used_map;
};

#define DIV_ROUND_UP(numerator, denominator) \
    (((numerator) + (denominator) - 1) / (denominator))

#define INT_BITS (CHAR_BIT * sizeof(int))

struct ucl_arena *ucl_arena_create(size_t object_size, size_t capacity) {
    struct ucl_arena *arena = malloc(sizeof(*arena));
    size_t used_map_size = DIV_ROUND_UP(capacity, INT_BITS) * sizeof(int);

    arena->object_size = object_size;
    arena->capacity = capacity;
    arena->objects = malloc(capacity * object_size);
    arena->used_map = malloc(used_map_size);

    memset(arena->used_map, 0, used_map_size);

    assert(arena->objects != NULL);
    assert(arena->used_map != NULL);

    return arena;
}

void ucl_arena_map(struct ucl_arena *arena, void (*map_function)(struct ucl_arena * arena, void *object)) {
    size_t used_map_ints = DIV_ROUND_UP(arena->capacity, INT_BITS);

    for (unsigned int i = 0; i < used_map_ints; i++ ) {
        // TODO: Allow for 'put' in map
        int map = arena->used_map[i];
        while (map) {
            int bit_index = __builtin_ffs(map) - 1;
            int index = bit_index + INT_BITS * i;
            void *obj = (char *) arena->objects + (index * arena->object_size);
            map_function(arena, obj);
            map &= ~(1 << bit_index);
        }
    }
}

int total_arena_gets = 0;

void *ucl_arena_get(struct ucl_arena *arena) {
    int used_map_ints = DIV_ROUND_UP(arena->capacity, INT_BITS);
    for (int i = 0; i < used_map_ints; i++ ) {
        // TODO: Maybe keep a cache of available used_map_ints
        int map = arena->used_map[i];
        if (!(~map)) {
            continue;
        }

        unsigned int bit_index = __builtin_ffs(~map) - 1;
        unsigned int index = bit_index + INT_BITS * i;
        if (index >= arena->capacity) {
            // This might happen in the last used_map_int when (capacity % int_bits != 0)
            return NULL;
        }
        arena->used_map[i] |= 1 << bit_index;
        total_arena_gets++;
        return (char *) arena->objects + (index * arena->object_size);
    }

    return NULL;
}

int total_arena_puts = 0;

void ucl_arena_put(struct ucl_arena *arena, void *object) {

    if (object == NULL) {
        return;
    }

    ptrdiff_t offset = (char *) object - (char *) arena->objects;
    unsigned int index = offset / arena->object_size;
    unsigned int int_index = index / INT_BITS;
    unsigned int bit_index = index % INT_BITS;

    assert(offset % arena->object_size == 0);
    assert(index < arena->capacity);

    assert(arena->used_map[int_index] & (1 << bit_index));
    //if (arena->used_map[int_index] & (1 << bit_index)) {
        total_arena_puts++;
    //}


    arena->used_map[int_index] &= ~(1 << bit_index);
}