#include "mem.h"

#include "common.h"
#include "dyn_arr.h"

// Array to track temp allocations.
void **temp_allocations;

#ifdef TEENSY_DEBUG

#include <execinfo.h>

#define BACKTRACE_SIZE 16

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

int allocations_cap;
int allocations_len;
struct alloc *allocations;

#endif // TODO TEENSY_DEBUG

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

    temp_allocations = ty_list_create(void*, 0);
}

void ty_deinit_mem(void)
{
    ty_list_free(temp_allocations);
#ifdef TEENSY_DEBUG
    fprintf(stderr, "%d allocations leaked\n", allocations_len);
    for (int i = 0; i < allocations_len; i++) {
        struct 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_free_temp_allocs(void)
{
    for (int i = 0; i < ty_list_len(temp_allocations); i++) {
        void *temp = temp_allocations[i];
        ty_free(temp);
    }
    ty_list_clear(temp_allocations);
}

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
    struct 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(struct 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++) {
        struct 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_talloc(size_t size)
{
    void *ptr = ty_alloc(size);
    ty_list_append(void*, &temp_allocations, ptr);
    return ptr;
}

void ty_free(void *ptr)
{
    free(ptr);

#ifdef TEENSY_DEBUG
    for (int i = 0; i < allocations_len; i++) {
        struct alloc *alloc = &allocations[i];
        if (alloc->ptr == ptr) {
            free(alloc->backtrace);
            *alloc = allocations[allocations_len - 1];
            allocations_len--;
            break;
        }
    }
#endif // TEENSY_DEBUG
}

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