// Testing allocation tracking stuff. (Very scuffed, not thread-safe!)

#include "../include/al/lib.h"

#if AL_USE_STDLIB
static void *(*_al_malloc)(size_t) = malloc;
static void *(*_al_calloc)(size_t, size_t) = calloc;
static void *(*_al_realloc)(void *, size_t) = realloc;
static void (*_al_free)(void *) = free;
#ifdef HAVE_POSIX_MEMALIGN
s32 (*_al_posix_memalign)(void **, size_t, size_t) = posix_memalign;
#endif

#define MEMORY_TRACKING 0

#if MEMORY_TRACKING
struct alloc_t {
    void *ptr;
    size_t size;
};

// Use stdlib allocation for the allocation tracking array.
// This avoids a circular dependency and also keeps allocations
// from the memory tracking code out of the stats.
#include "../include/al/types.h"
#define AL_ARRAY_MALLOC malloc
#define AL_ARRAY_REALLOC realloc
#define AL_ARRAY_FREE free
#include "array.c"

static array(struct alloc_t) allocations;
static u32 currently_allocated = 0;
static u32 peak_allocated = 0;
static u32 total_allocated = 0;
#endif

void al_malloc_init(void)
{
#if MEMORY_TRACKING
    al_array_init(allocations);
#endif
}

void al_malloc_close(void)
{
#if MEMORY_TRACKING
    al_array_free(allocations);
#endif
}

void *al_malloc(size_t n)
{
    void *ptr = _al_malloc(n);
#if MEMORY_TRACKING
    printf("**al_malloc(%zu)\n", n);
    struct alloc_t a = {
        .ptr = ptr,
        .size = n
    };
    total_allocated += a.size;
    currently_allocated += a.size;
    if (currently_allocated > peak_allocated) {
        peak_allocated = currently_allocated;
    }
    al_array_push(allocations, a);
#endif
    return ptr;
}

void *al_calloc(size_t n, size_t size)
{
    void *ptr = _al_calloc(n, size);
#if MEMORY_TRACKING
    printf("**al_calloc(%zu, %zu)\n", n, size);
    struct alloc_t a = {
        .ptr = ptr,
        .size = n * size
    };
    total_allocated += a.size;
    currently_allocated += a.size;
    if (currently_allocated > peak_allocated) {
        peak_allocated = currently_allocated;
    }
    al_array_push(allocations, a);
#endif
    return ptr;
}

void *al_realloc(void *ptr, size_t n)
{
    void *nptr = _al_realloc(ptr, n);
#if MEMORY_TRACKING
    printf("**al_realloc(%p, %zu)\n", ptr, n);
    for (u32 i = 0; i < allocations.size; i++) {
        struct alloc_t *a = &al_array_at(allocations, i);
        if (a->ptr == ptr) {
            currently_allocated += (ssize_t)n - (ssize_t)a->size;
            if (currently_allocated > peak_allocated) {
                peak_allocated = currently_allocated;
            }
            a->ptr = nptr;
            a->size = n;
        }
    }
    if (!ptr) {
        struct alloc_t a = {
            .ptr = nptr,
            .size = n
        };
        currently_allocated += a.size;
        if (currently_allocated > peak_allocated) {
            peak_allocated = currently_allocated;
        }
        al_array_push(allocations, a);
    }
    total_allocated += n;
#endif
    return nptr;
}

void al_free(void *ptr)
{
#if MEMORY_TRACKING
    printf("**al_free(%p)\n", ptr);
    for (u32 i = 0; i < allocations.size; i++) {
        struct alloc_t *a = &al_array_at(allocations, i);
        if (a->ptr == ptr) {
            currently_allocated -= a->size;
            al_array_remove_at(allocations, i);
            break;
        }
    }
#endif
    _al_free(ptr);
}

#ifdef HAVE_POSIX_MEMALIGN
s32 al_posix_memalign(void **ptr, size_t alignment, size_t n)
{
    s32 res = _al_posix_memalign(ptr, alignment, n);
#if MEMORY_TRACKING
    size_t actual_size = (n + alignment) - ((n + alignment) % alignment);
    printf("**al_posix_memalign(%zu, %zu(%zu))\n", alignment, n, actual_size);
    struct alloc_t a = {
        .ptr = *ptr,
        .size = actual_size
    };
    total_allocated += actual_size;
    currently_allocated += actual_size;
    if (currently_allocated > peak_allocated) {
        peak_allocated = currently_allocated;
    }
    al_array_push(allocations, a);
#endif
    return res;
}
#endif

void al_malloc_stats(u32 *current, u32 *peak, u32 *total)
{
#if MEMORY_TRACKING
    *current = currently_allocated;
    *peak = peak_allocated;
    *total = total_allocated;
#else
    *current = 0;
    *peak = 0;
    *total = 0;
#endif
}

void al_set_alloc(void *(*malloc_func)(size_t), void *(*calloc_func)(size_t, size_t),
    void *(*realloc_func)(void *, size_t), void (*free_func)(void *))
{
    _al_malloc = malloc_func;
    _al_calloc = calloc_func;
    _al_realloc = realloc_func;
    _al_free = free_func;
}
#endif

// Default to 4KB.
size_t al_page_size = 0x1000;
size_t al_grow_limit = 0xfff;

void al_set_page_size(size_t size)
{
    al_assert(size > 0);
    al_page_size = size;
    al_grow_limit = al_page_size - 1;
}