#include "arena.h"
#include <string.h>

Arena *
arena_init(void *mem, ptrdiff_t size) {
    Arena *arena = mem;
    arena->mem = arena + 1;
    arena->top = arena->mem;
    arena->lim = (char *) mem + size;
    return arena;
}

static void *
align_forward(void *ptr, ptrdiff_t align) {
    return (void *)(((uintptr_t)ptr + align - 1) & ~(align - 1));
}

static void
pretty_print_size(char buf[64], ptrdiff_t size) {
    if (size < 1024) {
        sprintf(buf, "%ld bytes", size);
    }
    else if (size < 1024 * 1024) {
        sprintf(buf, "%ld KB", size / 1024);
    }
    else if (size < 1024 * 1024 * 1024) {
        sprintf(buf, "%ld MB", size / (1024 * 1024));
    }
    else {
        sprintf(buf, "%ld GB", size / (1024 * 1024 * 1024));
    }
}

static void
arena_oom(Arena *arena, ptrdiff_t size, ptrdiff_t align) {
    char used[64];
    pretty_print_size(used, (uintptr_t) arena->top - (uintptr_t) arena->mem);
    char available[64];
    pretty_print_size(available, (uintptr_t) arena->lim - (uintptr_t) arena->top);
    char requested[64];
    pretty_print_size(requested, size);
    loge("Out of memory in arena %p: %s used, %s available, %s requested, %ld alignment",
            arena, used, available, requested, align);
    BP();
}

void *
arena_alloc_aligned(Arena *arena, ptrdiff_t size, ptrdiff_t align) {
    uintptr_t start = (uintptr_t) align_forward(arena->top, align);
    if (start + size > (uintptr_t) arena->lim) {
        arena_oom(arena, size, align);
        return NULL;
    }
    memset((void *)start, 0, size);
    arena->top = (void *)(start + size);
    return (void *)start;
}

void *
arena_realloc_aligned(Arena *arena, void *ptr, ptrdiff_t old_size, ptrdiff_t new_size, ptrdiff_t align) {
    if (ptr) {
        // check that pointer was allocated in this arena
        ASSERT((uintptr_t) ptr >= (uintptr_t) arena->mem);
        ASSERT((uintptr_t) ptr < (uintptr_t) arena->lim);
        ASSERT((uintptr_t) ptr + old_size <= (uintptr_t) arena->top);
    }

    void *result = NULL;
    if (false && ptr && (uintptr_t) ptr + old_size == (uintptr_t) arena->top) {
        // The block is at the top of the arena, so we can just attempt to extend it.
        if ((uintptr_t) ptr + new_size <= (uintptr_t) arena->lim) {
            // clear new memory
            memset((char *) ptr + old_size, 0, new_size - old_size);
            arena->top = ((char *) ptr + new_size);
            result = ptr;
        }
        else {
            arena_oom(arena, new_size, align);
        }
    }
    else {
        result = arena_alloc_aligned(arena, new_size, align);
        if (result) {
            if (ptr) {
                memcpy(result, ptr, old_size);
            }
        }
        else {
            arena_oom(arena, new_size, align);
        }
    }
    return result;
}

void 
arena_clear(Arena *arena) {
    arena_clear_overwrite(arena);
    arena->top = arena->mem;
}

void arena_clear_overwrite_with(Arena *arena, u8 value) {
    memset(arena->mem, value, (uintptr_t) arena->top - (uintptr_t) arena->mem);
    arena->top = arena->mem;
}

void arena_clear_overwrite(Arena *arena) {
    arena_clear_overwrite_with(arena, 0xCC);
}