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

AL_ASSERT_TYPE_SIZE(int, 4);
AL_ASSERT_TYPE_SIZE(unsigned, 4);
#if defined AL_WE_64BIT
AL_ASSERT_TYPE_SIZE(intptr_t, 8);
#elif defined AL_WE_32BIT
AL_ASSERT_TYPE_SIZE(long, 4);
AL_ASSERT_TYPE_SIZE(intptr_t, 4);
#endif

#ifdef AL_MEMORY_TRACKING
static void *(*_al_malloc)(size_t) = NULL;
static void *(*_al_calloc)(size_t, size_t) = NULL;
static void *(*_al_realloc)(void *, size_t) = NULL;
static void (*_al_free)(void *) = NULL;
#ifdef AL_HAVE_POSIX_MEMALIGN
s32 (*_al_posix_memalign)(void **, size_t, size_t) = NULL;
#endif

static void (*_al_malloc_lock)(void) = NULL;
static void (*_al_malloc_unlock)(void) = NULL;

// Testing allocation tracking stuff (Very scuffed).
struct al_alloc_t {
    void *ptr;
    size_t size;
};

// Use stdlib malloc 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 al_alloc_t) allocations;
static size_t current_alloc = 0;
static size_t peak_alloc = 0;
static size_t total_alloc = 0;
static size_t alloc_operations = 0;

#define total_changed_by(n) \
    total_alloc += (n); \

#define current_changed_by(n) \
    current_alloc += (n); \
    if (current_alloc > peak_alloc) { \
        peak_alloc = current_alloc; \
    } \
    alloc_operations++;

void al_malloc_init(void)
{
    al_array_init(allocations);
}

void al_malloc_close(void)
{
    al_array_free(allocations);
}

void *al_malloc(size_t n)
{
    al_printf("**al_malloc(%zu)\n", n);
    void *ptr = _al_malloc(n);
    _al_malloc_lock();
    total_changed_by(n);
    current_changed_by(n);
    al_array_push(allocations, ((struct al_alloc_t){ ptr, n }));
    _al_malloc_unlock();
    return ptr;
}

void *al_calloc(size_t n, size_t size)
{
    al_printf("**al_calloc(%zu, %zu)\n", n, size);
    void *ptr = _al_calloc(n, size);
    size_t total_size = n * size;
    _al_malloc_lock();
    total_changed_by(total_size);
    current_changed_by(total_size);
    al_array_push(allocations, ((struct al_alloc_t){ ptr, total_size }));
    _al_malloc_unlock();
    return ptr;
}

void *al_realloc(void *ptr, size_t n)
{
    al_printf("**al_realloc(%p, %zu)\n", ptr, n);
#ifdef AL_PARANOID_REALLOC
    void *new_ptr = ptr ? _al_realloc(ptr, n) : _al_malloc(n);
#else
    void *new_ptr = _al_realloc(ptr, n);
#endif
    _al_malloc_lock();
    total_changed_by(n);
    if (!ptr) {
        current_changed_by(n);
        al_array_push(allocations, ((struct al_alloc_t){ new_ptr, n }));
    } else {
        struct al_alloc_t *alloc;
        al_array_foreach_ptr(allocations, i, alloc) {
            al_assert(alloc->ptr);
            if (alloc->ptr == ptr) {
                ptrdiff_t diff = (ptrdiff_t)n - (ptrdiff_t)alloc->size;
                current_changed_by(diff);
                alloc->ptr = new_ptr;
                alloc->size = n;
                break;
            }
        }
    }
    _al_malloc_unlock();
    return new_ptr;
}

void al_free(void *ptr)
{
    al_printf("**al_free(%p)\n", ptr);
    _al_malloc_lock();
    bool found = false;
    struct al_alloc_t *alloc;
    al_array_foreach_ptr(allocations, i, alloc) {
        if (alloc->ptr == ptr) {
            al_assert(!found);
            current_changed_by(-alloc->size);
            al_array_remove_at(allocations, i);
            found = true;
        }
    }
    al_assert(found);
    _al_malloc_unlock();
    _al_free(ptr);
}

#ifdef AL_HAVE_POSIX_MEMALIGN
s32 al_posix_memalign(void **ptr, size_t alignment, size_t n)
{
    size_t real_size = (n + alignment) - ((n + alignment) % alignment);
    al_printf("**al_posix_memalign(%zu, %zu(%zu))\n", alignment, n, real_size);
    s32 result = _al_posix_memalign(ptr, alignment, n);
    _al_malloc_lock();
    total_changed_by(real_size);
    current_changed_by(real_size);
    al_array_push(allocations, ((struct al_alloc_t){ *ptr, real_size }));
    _al_malloc_unlock();
    return result;
}
#endif

void al_malloc_stats(size_t *current, size_t *peak, size_t *total, size_t *ops)
{
    _al_malloc_lock();
    *current = current_alloc;
    *peak = peak_alloc;
    *total = total_alloc;
    *ops = alloc_operations;
    _al_malloc_unlock();
}

void al_set_alloc(void *(*malloc_func)(size_t), void *(*calloc_func)(size_t, size_t),
    void *(*realloc_func)(void *, size_t),
#ifdef AL_HAVE_POSIX_MEMALIGN
    s32 (*posix_memalign_func)(void **, size_t, size_t),
#endif
    void (*free_func)(void *), void (*lock_func)(void), void (*unlock_func)(void))
{
    _al_malloc = malloc_func;
    _al_calloc = calloc_func;
    _al_realloc = realloc_func;
#ifdef AL_HAVE_POSIX_MEMALIGN
    _al_posix_memalign = posix_memalign_func;
#endif
    _al_free = free_func;
    _al_malloc_lock = lock_func;
    _al_malloc_unlock = unlock_func;
}
#endif

u32 al_grow_step = 0xfff; // 0b111111111111, 4095
//u32 al_grow_step = 0x3fff; // 0b11111111111111, 16383
u32 al_growing_allocation(void **ptr, u32 prev_size, u32 size, void *(*realloc_func)(void *, size_t))
{
    if (size <= prev_size) return prev_size;

    al_assert(size > 0);
    if (size <= al_grow_step) {
        // next_power_of_two(4096) = 4096.
        size = al_next_power_of_two(size);
    } else {
        size = al_u32_add_wrap(size, al_grow_step, MAX_ALLOC_32) & ~al_grow_step;
    }
    // size can only be 0 after wrapping around MAX_ALLOC_32 (zero all possible bits).
    if (size == 0) size = MAX_ALLOC_32;

    *ptr = realloc_func(*ptr, size);
    al_assert(*ptr);

    return size;
}

size_t al_page_size = 0x1000; // 4096
void al_set_page_size(size_t size)
{
    al_assert(size > 0);
    al_page_size = size;
}