#include "mem.h"

#include <GLFW/glfw3.h>

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

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

#ifdef UE_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 // UE_DEBUG

static
void *glfw_alloc(size_t size, void *user)
{
        (void)user;
        return mem_alloc(size);
}

static
void *glfw_realloc(void *block, size_t size, void *user)
{
        (void)user;
        return mem_realloc(block, size);
}

static
void glfw_dealloc(void *block, void *user)
{
        (void)user;
        mem_free(block);
}

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

        GLFWallocator glfw_allocator;
        glfw_allocator.allocate = glfw_alloc;
        glfw_allocator.reallocate = glfw_realloc;
        glfw_allocator.deallocate = glfw_dealloc;
        glfwInitAllocator(&glfw_allocator);

        temp_allocations = da_create(void*, 0);
}

void deinit_mem(void)
{
        da_free(temp_allocations);
#ifdef UE_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(ERR_MEM_LEAK);
#endif // UE_DEBUG
}

void free_temp_allocs(void)
{
        for (int i = 0; i < da_len(temp_allocations); i++) {
                void *temp = temp_allocations[i];
                mem_free(temp);
        }
        da_clear(temp_allocations);
}

void *mem_alloc(size_t size)
{
        void *ptr = malloc(size);
        if (!ptr)
                log_fatal(ERR_MEM, "(%s) ran out of memory", __func__);

#ifdef UE_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)
                        log_fatal(
                                ERR_MEM,
                                "(%s) (track) ran out of memory",
                                __func__
                        );
        }

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

        return ptr;
}

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

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

#ifdef UE_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 // UE_DEBUG

        return new_ptr;
}

void *mem_talloc(size_t size)
{
        void *ptr = mem_alloc(size);
        da_append(void*, &temp_allocations, ptr);
        return ptr;
}

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

#ifdef UE_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 // UE_DEBUG
}

int alloc_count(void)
{
#ifdef UE_DEBUG
        return allocations_len;
#else
        return -1;
#endif
}