#include "memory_arena.h"
#include "o1heap.h"

#define CC_NO_SHORT_NAMES
#include "cc.h"

#include "sc_mutex.h"
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <assert.h>

/**
 * This memory arena is engine-wide
 * Thus, this can take on a singleton design.
 *
 * o1heap does not store the actual buffer size.
 * To reimplement realloc, we need to keep track of the size of memory
 * allocation for a pointer, so that we can copy over. Thus, we use a
 * unordered map to do so.
 *
 * So, there are two allocation from system heap:
 * 1. Memory arena buffer
 * 2. The memory size map
 */

static void* mem_buffer = NULL;
static O1HeapInstance* heap_handle = NULL;
static struct sc_mutex lock = {0};
static cc_map(uintptr_t, size_t) mmap;
static size_t total_mallocd = 0;

bool mem_arena_init(uint16_t size_mb)
{
    if (mem_buffer != NULL) return true;

    const size_t sz = size_mb * 1024U * 1024U;
    const size_t raw_sz = sz + O1HEAP_ALIGNMENT - 1U;
    mem_buffer = calloc(1, raw_sz);

    if (mem_buffer == NULL) return false;

    int res = sc_mutex_init(&lock);
    if (res != 0) {
        free(mem_buffer);
        mem_buffer = NULL;
        return false;
    }

    cc_init(&mmap);

    heap_handle = o1heapInit(mem_buffer + ((uintptr_t)mem_buffer % O1HEAP_ALIGNMENT), sz);
    total_mallocd = 0;
    return true;
}

bool mem_arena_deinit(void)
{
    if (mem_buffer == NULL) return true;

    free(mem_buffer);
    mem_buffer = NULL;
    heap_handle = NULL;

    sc_mutex_term(&lock);
    cc_cleanup(&mmap);
    total_mallocd = 0;
    return true;
}

void mem_arena_print()
{
    O1HeapDiagnostics diag = o1heapGetDiagnostics(heap_handle);
    printf("O1heap Memory Arena Info\n");
    printf("--------------------\n");
    printf("Capacity: %.3f MB\n", diag.capacity * 1.0 / 1024 / 1024);
    printf("Allocated: %.3f MB (%lu)\n", diag.allocated * 1.0 /1024/1024, diag.allocated);
    printf("Peak allocated: %.3f MB\n", diag.peak_allocated * 1.0 /1024/1024);
    printf("Peak request: %lu\n", diag.peak_request_size);
    printf("OOM count: %lu\n\n", diag.oom_count);
}

size_t mem_arena_get_allocated(void)
{
    return total_mallocd;
}

void* mem_arena_malloc(size_t size)
{
    sc_mutex_lock(&lock);
    void* buf = o1heapAllocate(heap_handle, size);
    assert(cc_insert(&mmap, (uintptr_t)buf, size) != NULL);
    total_mallocd += size;
    sc_mutex_unlock(&lock);
    return buf;
}

void mem_arena_free(void* ptr)
{
    // Because we keep track of the memory alloc'd, we can check if the memory
    // being freed has been allocated by the heap.
    // The free will be ignored if it is not from the heap.
    // This is useful in case you want to mix static and dynamic memory in a component field.
    // And you want to universally free the pointer without care.
    sc_mutex_lock(&lock);
    size_t* sz = cc_get(&mmap, (uintptr_t)ptr);

    //assert(sz != NULL);
    if (sz == NULL) {
        return;
    }
    o1heapFree(heap_handle, ptr);

    total_mallocd -= *sz;
    assert(cc_erase(&mmap, (uintptr_t)ptr));
    assert(o1heapDoInvariantsHold(heap_handle));
    sc_mutex_unlock(&lock);
}

void* mem_arena_calloc(size_t nmemb, size_t sz)
{
    const size_t total_sz = nmemb * sz;
    void* buf = mem_arena_malloc(total_sz);
    if (buf == NULL) return NULL;

    memset(buf, 0, total_sz);
    return buf;
}

void* mem_arena_realloc(void* buf, size_t new_size)
{
    /**
     * Simple implementation: Malloc a new memory and copy over.
     * This is not friendly towards the arena, but it works.
     *
     */
    if (buf == NULL) {
        return mem_arena_malloc(new_size);
    }

    size_t* old_sz = cc_get(&mmap, (uintptr_t)buf);
    if (old_sz == NULL) return NULL;

    void* new_buf = mem_arena_malloc(new_size);
    if (new_buf == NULL) return NULL;

    size_t cpy_sz = (*old_sz > new_size) ? new_size : *old_sz;
    memcpy(new_buf, buf, cpy_sz);
    mem_arena_free(buf);

    return new_buf;
}