#include <stdlib.h>

/// -------- Basic functions

void *un_memory_alloc(u64 size, Allocator alloc) {
    assert(alloc.proc != NULL);
    return alloc.proc(NULL, size, UN_DEFAULT_ALIGN, UN_ALLOC_MSG_ALLOCATE, alloc.data);
}

void *un_memory_alloc_align(u64 size, u32 align, Allocator alloc) {
    assert(alloc.proc != NULL);
    return alloc.proc(NULL, size, align, UN_ALLOC_MSG_ALLOCATE, alloc.data);
}

void un_memory_free(void *ptr, Allocator alloc) {
    assert(alloc.proc != NULL);
    alloc.proc(ptr, 0, 0, UN_ALLOC_MSG_FREE, alloc.data);
}

/// -------- Std allocator

#if defined(OS_WINDOWS)
#   include "un_mem_win_x64.c"
#elif defined(OS_LINUX)
#   include "un_mem_linux_x64.c"
#else
#   include "un_mem_std.c"
#endif

Allocator un_alloc_std_get(void) {
    return CLITERAL(Allocator) { .proc = un_std_alloc_proc };
}

/// Arena allocator

typedef struct Arena {
    u8 *buffer;
    u64 size, curr_offset, prev_offset;
} Arena;

static Arena __temp_arena;

static Arena arena_create(u64 size) {
    Allocator alloc;
    Arena     arena;

    UN_CLEAR(arena);

    alloc = un_alloc_std_get();
    arena.buffer = un_memory_alloc(size, alloc);

    if (arena.buffer == NULL) {
        return arena;
    }

    arena.size = size;
    return arena;
}

// special thanks to ginger bill: https://www.gingerbill.org/article/2019/02/08/memory-allocation-strategies-002/
static uintptr_t align_forward(uintptr_t ptr, u32 align) {
    uintptr_t p, a, modulo;

    assert(UN_POW_OF2(align));

    p = ptr;
    a = (uintptr_t)align;
	modulo = p & (a-1);

	if (modulo != 0) {
		p += a - modulo;
	}

	return p;
}

static void* arena_allocate_align(Arena *arena, u32 align, u64 size) {
    uintptr_t curr_ptr, offset;
    u8 *ptr;

    assert(arena != NULL);

    curr_ptr = (uintptr_t)arena->buffer + arena->curr_offset;
    offset   = align_forward(curr_ptr, align);
    offset  -= (uintptr_t)arena->buffer;
    
    if ((offset + size) <= arena->size) {
        ptr = (u8*)(arena->buffer + offset);
        arena->prev_offset = offset;
        arena->curr_offset = offset + size;
        memset(ptr, 0, size);
    } else {
        ptr = NULL;
    }

    return ptr;
}

static void arena_free_memory(Arena *arena) {
    if (arena == &__temp_arena) return;

    assert(arena != NULL);

    if (arena->buffer != NULL) {
        un_memory_free(arena->buffer, un_alloc_std_get());
    }
}

static ALLOCATOR_PROC_SIGNATURE(un_arena_alloc_proc) {
    UNUSED(p);

    switch (message) {
        case UN_ALLOC_MSG_ALLOCATE:
            return arena_allocate_align((Arena*)data, align, size);
        case UN_ALLOC_MSG_FREE:
            break;
        case UN_ALLOC_MSG_SELF_DELETE:
            arena_free_memory((Arena*)data);
            break;
    }

    return NULL;
}

Allocator un_alloc_arena_create(u64 byte_size) {
    Allocator alloc, std;
    Arena arena;

    UN_CLEAR(alloc);
    std        = un_alloc_std_get();
    alloc.proc = un_arena_alloc_proc;
    arena      = arena_create(byte_size);

    if(arena.buffer == NULL) {
        UN_CLEAR(alloc);
    }

    alloc.data = un_memory_alloc(sizeof(arena), std); // yeah, it allocs 4kb for 32 bytes...

    if(alloc.data == NULL) {
        arena_free_memory(&arena);
        UN_CLEAR(alloc);
        return alloc;
    }

    *((Arena *)alloc.data) = arena;
    return alloc;
}

void un_alloc_arena_destroy(Allocator *alloc) {
    arena_free_memory((Arena*)alloc->data);

    un_memory_free(alloc->data, un_alloc_std_get());

    UN_CLEAR(*alloc);
}


void un_alloc_temp_init(u64 byte_size) {
    __temp_arena = arena_create(byte_size);
}

void un_alloc_temp_reset(void) {
    __temp_arena.curr_offset = 0;
    __temp_arena.prev_offset = 0;
}

Allocator un_alloc_temp_get(void) {
    assert(__temp_arena.buffer != NULL);

    return CLITERAL(Allocator) { un_arena_alloc_proc, &__temp_arena };
}