#include "arena.h"
#include "shared.h"

#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <errno.h>
#include <error.h>

#define VALLOC(nmemb, size) ((___VXGG___USE_XALLOC_FOR_ARENAS___ > 0) ? xcalloc((nmemb), (size)) : malloc((nmemb) * (size)))

typedef struct an {
    void *membase;
    void *memcur;
    size_t allocated;
    size_t used;

    struct an *next;
} arenanode;

typedef struct arena {
    arenanode *start;
    arenanode *current;

    // For keeping track of the largest possible thing that can be allocated to one node
    size_t node_memspace;
} arena;

int arenanode_init(arenanode **an, size_t bytes) {
    (*an) = VALLOC(1, sizeof(**an));
    if(!(*an))
        return -1;

    (*an)->allocated = bytes;
    (*an)->used = 0;
    (*an)->next = NULL;
    
    void *mem = VALLOC(bytes, 1);
    if(!mem) {
        free((*an));
        (*an) = NULL;
        return -1;
    }

    (*an)->membase = mem;
    (*an)->memcur = mem;

    return 0;
}

int arena_init(arena **a, size_t bytes) {
    if(!ISPOWOF2(MEM_ALIGN_BYTES))
        XALLOC_EXIT("<arena_init> \"MEM_ALIGN_BYTES\" is not a power of 2. Refusing to create a new arena");

    (*a) = VALLOC(1, sizeof(**a));
    if(!(*a))
        return -1;
    (*a)->start = NULL;
    (*a)->current = NULL;
    (*a)->node_memspace = bytes;

    arenanode *base;
    if(arenanode_init(&base, bytes) < 0) {
        free(*a);
        (*a) = NULL;
        return -1;
    }

    (*a)->start = base;
    (*a)->current = base;

    return 0;
}

void* arena_alloc(arena * const arena, size_t bytes) {
    if(!ISPOWOF2(MEM_ALIGN_BYTES))
        XALLOC_EXIT("<arena_alloc> \"MEM_ALIGN_BYTES\" is not a power of 2. Refusing to allocate any memory");
    if(!arena)
        RETURNWERR(EINVAL, NULL);
    if(bytes > arena->node_memspace)
        RETURNWERR(ENOMEM, NULL);

    if(bytes > ((arena->current)->allocated - (arena->current)->used)) {
        arenanode *new;
        if(arenanode_init(&new, arena->node_memspace) < 0)
            return NULL;

        arena->current->next = new;
        arena->current = new;
    }

    size_t alignment = MEM_ALIGN(bytes);
    size_t offset = bytes + alignment;

    void *mem = arena->current->memcur;
    arena->current->memcur = (void*)(((uint8_t*)arena->current->memcur) + offset); 

    arena->current->used += offset;

    return mem;

        // Note: This implementation assumes that malloc provides already-aligned memory. If it
        // doesn't, that sucks and blows everything up
}

int arena_free(arena **arena) {
    if(!arena)
        RETURNWERR(EINVAL, -1);
    if(!(*arena))
        RETURNWERR(EINVAL, -1);

    (*arena)->current = (*arena)->start;
    for(arenanode *n; (*arena)->current != NULL;) {
        n = (*arena)->current->next;
        
        free((*arena)->current->membase);
        free((*arena)->current);

        (*arena)->current = n;
    }
    free((*arena));
    (*arena) = NULL; 

    return 0;
}

int arena_clear(arena **arena) {
    if(!arena)
        RETURNWERR(EINVAL, -1);
    if(!(*arena))
        RETURNWERR(EINVAL, -1);

    size_t bytes = (*arena)->node_memspace;
    
    int ret = 0;
    if((ret = arena_free(arena)) < 0)
        return ret;
    return arena_init(arena, bytes);
}



// Simple Arena is an arena that can't expand whenever allocating memory, meaning what you originally allocated is what you get

typedef arena simplearena;

int simplearena_init(simplearena **a, size_t bytes) {
    return arena_init(a, bytes);
}

void* simplearena_alloc(simplearena * const a, size_t bytes) {
    // The criteria to allocate new memory in arena_alloc is 'bytes > ((a->current)->allocated - (a->current)->used)', so if this
    // is true, just return NULL & set errno

    if(!a)
        RETURNWERR(EINVAL, NULL);
    if(bytes > ((a->current)->allocated - (a->current)->used))
        RETURNWERR(ENOMEM, NULL);

    return arena_alloc(a, bytes);
}

int simplearena_free(simplearena **a) {
    return arena_free(a);
}

int simplearena_clear(simplearena **a) {
    return arena_clear(a);
}