#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;

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

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

    an->membase = mem;
    an->memcur = mem;

    return an;
}

arena * arena_init(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");

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

    arenanode *base = arenanode_init(bytes);
    if(!base) {
        free(a);
        return NULL;
    }

    a->start = base;
    a->current = base;

    return a;
}

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 = arenanode_init(arena->node_memspace);
        if(!new)
            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;
    (*arena) = arena_init(bytes);

    return (!(*arena)) ? -1 : 0;
}



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

typedef arena simplearena;

simplearena * simplearena_init(size_t bytes) {
    return arena_init(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);
}