#include <stdlib.h>

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

void un_memory_set(u8 *buffer, u8 value, u64 size) {
    if (size == 0) return;
    assert(buffer != NULL);

    while (size-- > 0) {
        *buffer++ = value;
    }
}

void un_memory_copy(u8 *dest, u8 *source, u64 size) {
    if (size == 0) return;

    assert(dest   != NULL);
    assert(source != NULL);

#ifdef DEBUG

    if (dest < source) {
        assert(((ptrdiff_t)source - (ptrdiff_t)dest) >= (ptrdiff_t)size);
    } else {
        assert(((ptrdiff_t)dest - (ptrdiff_t)source) >= (ptrdiff_t)size);
    }
#endif // DEBUG

    while (size-- > 0) {
        *dest++ = *source++;
    }
}

s32 un_memory_compare(u8 *left, u8 *right, u64 size) {
    while (size-- > 0) {
        if (*left++ == *right++)
            continue;

        return left[-1] > right[-1] ? 1 : -1;
    }

    return 0;
}

void un_memory_move(u8 *dest, u8 *src, u64 size) {
    Allocator talloc = un_allocator_get_temporary();

    void *temp = un_memory_alloc(size, talloc);

    un_memory_copy(temp, src, size);
    un_memory_copy(dest, temp, size);
}

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

void *un_memory_realloc(void *ptr, u64 size, Allocator alloc) {
    assert(alloc.proc != NULL);
    return alloc.proc(ptr, size, UN_ALLOC_MSG_REALLOCATE, alloc.data);
}

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

void un_memory_destroy(Allocator *alloc) {
    assert(alloc != NULL);
    assert(alloc->proc != NULL);

    alloc->proc(NULL, 0, UN_ALLOC_MSG_SELF_DELETE, alloc->data);

    UN_CLEAR(*alloc);
}

/// -------- Stdlib allocator

static ALLOCATOR_PROC_SIGNATURE(un_std_alloc_proc) {
    UNUSED(data);

    switch (message) {
        case UN_ALLOC_MSG_ALLOCATE:
            return calloc(1, size);
        case UN_ALLOC_MSG_REALLOCATE:
            return realloc(p, size);
        case UN_ALLOC_MSG_FREE:
            free(p);
            break;
        case UN_ALLOC_MSG_SELF_DELETE: break;
    }

    return NULL;
}

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

/// -------- Temp allocator

static struct {
    u64 index;
    u8 data[UN_TEMP_SIZE];
} temp; 

static void temp_reset(void) {
    temp.index = 0;
}

static void *temp_alloc(u64 size) {
    if ((temp.index + size) > UN_TEMP_SIZE) {
        // @todo: Decice what is the best behaviour for wrapping
        assert(false);
        temp_reset();
    }

    if ((temp.index + size) > UN_TEMP_SIZE) {
        return NULL;
    }

    void *ptr = (u8*)temp.data + temp.index;
    temp.index += size;
    un_memory_set(ptr, 0x00, size);

    return ptr;
}

static ALLOCATOR_PROC_SIGNATURE(un_temp_alloc_proc) {
    UNUSED(data);
    UNUSED(p);

    switch (message) {
        case UN_ALLOC_MSG_ALLOCATE:
            return temp_alloc(size);
        case UN_ALLOC_MSG_REALLOCATE:
        case UN_ALLOC_MSG_FREE:
            break;
        case UN_ALLOC_MSG_SELF_DELETE: break;
    }

    return NULL;
}

Allocator un_allocator_get_temporary(void) {
    return CLITERAL(Allocator) { .proc = un_temp_alloc_proc };
}

/// Arena allocator

typedef struct Arena {
    u64 size;
    u64 occupied;
    struct Arena *next;
    u8 data[1];
} Arena;

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

    alloc = un_allocator_get_standard();
    arena = CLITERAL(Arena*)un_memory_alloc(size, alloc);

    if (!arena) {
        return NULL;
    }

    UN_CLEAR(*arena);
    arena->size = size - sizeof(Arena);
    return arena;
}

static void arena_delete(Arena *arena) {
    if (arena->next != NULL) {
        arena_delete(arena->next);
    }

    Allocator alloc = un_allocator_get_standard();
    un_memory_free(arena, alloc);
}

static void *arena_allocate(u64 size, Arena *arena) {
    u8* pos;

    if (size <= (arena->size - arena->occupied)) {
        pos = arena->data + arena->occupied;
        arena->occupied += size;
        return (void*)pos;
    }

    if (!arena->next) {
        arena->next = arena_create((arena->size + sizeof(Arena)) * 2LL);
    } 

    return arena_allocate(size, arena->next);
}

static ALLOCATOR_PROC_SIGNATURE(un_arena_alloc_proc) {
    assert(data != NULL);
    UNUSED(p);

    switch (message) {
        case UN_ALLOC_MSG_ALLOCATE:
            return arena_allocate(size, (Arena*)data);
        case UN_ALLOC_MSG_REALLOCATE:
            break;
        case UN_ALLOC_MSG_FREE:
            break;
        case UN_ALLOC_MSG_SELF_DELETE:
            arena_delete((Arena*)data);
            break;
    }

    return NULL;
}

Allocator un_allocator_create_arena(u64 size) {
    Allocator alloc = { 0 };

    alloc.proc = un_arena_alloc_proc;
    alloc.data = arena_create(size);
    
    return alloc;
}