#include "teensy_mem.h"

#include "teensy_common.h"
#include "teensy_list.h"

#define TEMP_ALLOC_ARENA_SIZE (1024*1024)

uint8_t temp_arena[TEMP_ALLOC_ARENA_SIZE];
uint8_t *next_temp_alloc = temp_arena;

#ifdef TEENSY_DEBUG

#include <execinfo.h>

#define BACKTRACE_SIZE 16

typedef struct {
    void *ptr;
    size_t size;
    int backtrace_len;
    char **backtrace;
} Alloc;

int allocations_cap;
int allocations_len;
Alloc *allocations;

#endif // TODO TEENSY_DEBUG

void ty_init_mem(void)
{
#ifdef TEENSY_DEBUG
    allocations_len = 0;
    allocations_cap = 8;
    allocations = malloc(sizeof(Alloc) * 8);
    if (!allocations)
        ty_log_fatal(TY_ERR_MEM, "(%s) (track) ran out of memory", __func__);
#endif // TEENSY_DEBUG
}

void ty_deinit_mem(void)
{
#ifdef TEENSY_DEBUG
    fprintf(stderr, "%d allocations leaked\n", allocations_len);
    for (int i = 0; i < allocations_len; i++) {
        Alloc *alloc = &allocations[i];
        fprintf(
            stderr,
            "leaked %zu bytes at %p\n",
            alloc->size,
            alloc->ptr
        );
        for (int j = 0; j < alloc->backtrace_len; j++) {
            fprintf(stderr, "\t%s\n", alloc->backtrace[j]);
        }
        putc('\n', stderr);

        free(alloc->backtrace);
    }
    free(allocations);

    if (allocations_len > 0)
        exit(TY_ERR_MEM_LEAK);
#endif // TEENSY_DEBUG
}

void *ty_talloc(size_t size)
{
    if (next_temp_alloc > temp_arena + TEMP_ALLOC_ARENA_SIZE)
        ty_log_fatal(TY_ERR_MEM, "bump up the temp alloc arena size");
        
    void *ptr = next_temp_alloc;
    next_temp_alloc += size;
    return ptr;
}

void ty_free_temp_allocs(void)
{
    next_temp_alloc = temp_arena;
}

void *ty_alloc(size_t size)
{
    void *ptr = malloc(size);
    if (!ptr)
        ty_log_fatal(TY_ERR_MEM, "(%s) ran out of memory", __func__);

#ifdef TEENSY_DEBUG
    // Log allocation
    Alloc alloc;
    void *bt[BACKTRACE_SIZE];
    int bt_len = backtrace(bt, BACKTRACE_SIZE);
    char **symbols = backtrace_symbols(bt, bt_len);

    alloc.ptr = ptr;
    alloc.size = size;
    alloc.backtrace_len = bt_len;
    alloc.backtrace = symbols;

    if (allocations_len + 1 > allocations_cap) {
        allocations_cap *= 2;
        allocations = realloc(
            allocations,
            sizeof(Alloc) * allocations_cap
        );
        if (!allocations)
            ty_log_fatal(
                TY_ERR_MEM,
                "(%s) (track) ran out of memory",
                __func__
            );
    }

    allocations[allocations_len++] = alloc;
#endif // TEENSY_DEBUG

    return ptr;
}

void *ty_realloc(void *ptr, size_t new_size)
{
    if (new_size == 0) {
        ty_free(ptr);
        return NULL;
    }

    void *new_ptr = realloc(ptr, new_size);
    if (!new_ptr)
        ty_log_fatal(TY_ERR_MEM, "(%s) ran out of memory", __func__);

#ifdef TEENSY_DEBUG
    for (int i = 0; i < allocations_len; i++) {
        Alloc *alloc = &allocations[i];
        if (alloc->ptr == ptr) {
            alloc->ptr = new_ptr;
            alloc->size = new_size;
            break;
        }
    }
#endif // TEENSY_DEBUG

    return new_ptr;
}

void ty_free(void *ptr)
{
#ifdef TEENSY_DEBUG
    bool double_free = true;

    for (int i = 0; i < allocations_len; i++) {
        Alloc *alloc = &allocations[i];
        if (alloc->ptr == ptr) {
            free(alloc->backtrace);
            *alloc = allocations[allocations_len - 1];
            allocations_len--;
            double_free = false;
            break;
        }
    }

    if (double_free) {
        ty_log_fatal(
            TY_ERR_MEM,
            "Double free or freed unowned memory (%p)",
            ptr
        );
    }
#endif // TEENSY_DEBUG
    free(ptr);
}

int ty_alloc_count(void)
{
#ifdef TEENSY_DEBUG
    return allocations_len;
#else
    return -1;
#endif
}