#include "mem_arena.h"
#include "config.h"

#include <stdlib.h>

RT_CVAR_I(rt_TemporaryArenaSize, "Size of temporary arenas in bytes. Default: 32 MB", RT_MB(32));

#define ALIGNMENT 0xf
#define ALIGN(n)  (((n) + ALIGNMENT) & ~ALIGNMENT)

extern void *memset(void *, int, size_t);

RT_DLLEXPORT rt_create_arena_result rtCreateArena(void *memory, size_t size) {
    rt_arena arena = {.needs_free = 0};
    if (!memory) {
        size   = ALIGN(size);
        memory = malloc(size);
        if (!memory) {
            return (rt_create_arena_result){
                .ok = 0,
            };
        }
        arena.needs_free = 1;
    }
    arena.base = memory;
    arena.size = size;
    arena.at   = 0;
    return (rt_create_arena_result){
        .arena = arena,
        .ok    = 1,
    };
}

RT_DLLEXPORT void *rtArenaPush(rt_arena *arena, size_t n) {
    if (n == 0)
        return NULL;
    n = ALIGN(n);
    if (arena->at + n > arena->size)
        return NULL;
    void *p = (char *)arena->base + arena->at;
    arena->at += n;
    return p;
}

RT_DLLEXPORT void *rtArenaPushZero(rt_arena *arena, size_t n) {
    void *p = rtArenaPush(arena, n);
    if (p)
        memset(p, 0, ALIGN(n));
    return p;
}

RT_DLLEXPORT void rtArenaPop(rt_arena *arena, size_t n) {
    n = ALIGN(n);
    if (arena->at < n)
        return;
    arena->at -= n;
}

RT_DLLEXPORT void rtReleaseArena(rt_arena *arena) {
    if (arena->needs_free)
        free(arena->base);
    arena->base       = NULL;
    arena->at         = 0;
    arena->size       = 0;
    arena->needs_free = 0;
}

/* Temporary arena pool */

typedef uint32_t rt_thread_id;
extern RT_DLLEXPORT rt_thread_id rtGetCurrentThreadId(void);

#define NUM_TEMP_ARENAS_PER_THREAD 2
typedef struct {
    rt_arena arenas[NUM_TEMP_ARENAS_PER_THREAD];
} rt_thread_temp_arenas;

static RT_THREAD_LOCAL rt_thread_temp_arenas t_arenas;

RT_DLLEXPORT rt_temp_arena rtGetTemporaryArena(const rt_arena **permanent_arenas, int count) {
    if (!t_arenas.arenas[0].base) {
        /* Initialize */
        for (int i = 0; i < NUM_TEMP_ARENAS_PER_THREAD; ++i) {
            rt_create_arena_result res = rtCreateArena(NULL, (size_t)rt_TemporaryArenaSize.i);
            if (!res.ok) {
                rtLog("CORE",
                      "Failed to initialize thread-local temporary arenas for thread: %u",
                      rtGetCurrentThreadId());
                return (rt_temp_arena){.arena = NULL, .at = 0};
            }
            t_arenas.arenas[i] = res.arena;
        }
    }
    for (int i = 0; i < NUM_TEMP_ARENAS_PER_THREAD; ++i) {
        int conflict_found = 0;
        for (int j = 0; j < count; ++j) {
            if (permanent_arenas[j] == &t_arenas.arenas[i])
                conflict_found = 1;
        }
        if (!conflict_found) {
            return rtBeginTempArena(&t_arenas.arenas[i]);
        }
    }
    rtLog("CORE",
          "Failed to find a usable thread-local temporary arena for thread: %u",
          rtGetCurrentThreadId());
    return (rt_temp_arena){.arena = NULL, .at = 0};
}