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

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

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
        // I don't know if I care to keep this or remove it. I'll see if it becomes a problem / is useful
    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", );
    if(bytes < 1)
        RETURNWERR(EINVAL, NULL);

    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
}

void arena_free(void *a) {
    if(!a)
        return;

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

        real->current = n;
    }
    free(real); 

    return;
}




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);
}

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