#ifndef __AMR_INCLUDE_AMR_MEMORY_H
#define __AMR_INCLUDE_AMR_MEMORY_H

#define AMRM_KB(x) ((x)*1024)
#define AMRM_MB(x) ((AMR_KB(x))*1024)
#define AMRM_GB(x) ((AMR_MB(x))*1024)

#ifndef AMRM_DEFAULT_ALIGNMENT
#define AMRM_DEFAULT_ALIGNMENT (2*sizeof(void *))
#endif

////////////////////////////////////////////////////
//
// Arena
//


// This struct is defined publicly but it really is supposed to be opaque
// and is used by all the functions here
struct amrm_Arena {
    unsigned char *arena;
    size_t prev_offset;
    size_t curr_offset;
    size_t capacity;
};

typedef struct amrm_Arena amrm_Arena;

int8_t amrm__is_power_of_two(uintptr_t x); 
// Uses some bit manipulation to check if x is a power of two.
// This is needed for alignment

uintptr_t amrm__fast_modulo(uintptr_t p, uintpt_t a) {
// Uses a bitwise operator to calculate modulus

uintptr_t amrm__align_forward(uintptr_t ptr, size_t alignment);
// Ensures a ptr is aligned by 'alignment'
// It does this by fast_modulus(ptr, alignment) and moves ptr ahead
// by the necessary amount to have it be word aligned

void  amrm_arena_init(amrm_Arena *a, unsigned char *raw_buffer, size_t capacity);
// Initialise memory arena struct
// Pass in raw buffer, acquired through malloc 
// (but in future, test for os syscalls to acquire virtual pages)
// and the capacity of buffer in bytes

void* amrm_arena_alloc(amrm_Arena *a, size_t size);
// Allocate memory in arena, which roughly translates to
// perform bound checks, request memory from the arena and update offsets

void* amrm_arena_alloc_aligned(amrm_Arena *a, size_t size, size_t alignment);
// This is the function called by amrm_arena_alloc. The main item is that it 
// expects an alignment

void* amrm_arena_resize(amrm_Arena *a, void *old,
                        size_t old_size, size_t new_size);
// Resize memory in arena
// Pass the current pointer to old, along with its size and the new requested size

void* amrm_arena_resize_aligned(amrm_Arena *a, void *old, 
                                size_t old_size, size_t new_size, size_t alignment);
// Same as amrm_arena_resize, this is what it calls. It just expects an alignment

void  amrm_arena_clear(amrm_Arena *a);
// "clears" an arena, though what this actually does it clears the offsets
// so the arena considers itself empty.

////////////////////////////////////////////////////
//
// IMPLEMENTATION
//

#ifdef AMR_ARRAY_IMPLEMENTATION

int8_t amrm__is_power_of_two(uintptr_t x) {
    return (x & (x-1)) == 0;
}

uintptr_t amrm__fast_modulo(uintptr_t p, uintpt_t a) {
    return (p & (a-1));
}

uintptr_t amrm__align_forward(uintptr_t ptr, size_t alignment) {
    uintptr_t p, a, mod;
    p = ptr;
    a = (uintptr_t)alignment;
    mod = amrm__fast_modulo(p, a);
    if (mod != 0) {
        p += (a-mod);
    }

    return p;
}

void amrm_arena_init(amrm_Arena *a, unsigned char *raw_buffer, size_t capacity) {
    a->arena = raw_buffer;
    a->prev_offset = 0;
    a->curr_offset = 0;
    a->capacity = capacity;
}

void* amrm_arena_alloc(amrm_Arena *a, size_t size) {
    return amrm_arena_alloc_aligned(a, size, AMRM_DEFAULT_ALIGNMENT);
}

void* amrm_arena_alloc_aligned(amrm_Arena *a, size_t size, size_t alignment) {
    void *ptr = NULL;

    assert(is_power_of_two(alignment));

    uintptr_t curr_ptr = (uintptr_t)a->buffer + a->curr_offset;
    uintptr_t curr_ptr_aligned = amrm__align_forward(curr_ptr, alignment);
    // if I really want to do a bounds check here, I think that is on the 
    // compiler to catch, mainly because if someone is typing that they need 
    // an insane amount of memory, then catch that
    if (curr_ptr_aligned + size <= a->buffer + a->capacity) {
        ptr = curr_ptr_aligned;
        a->prev_offset = a->curr_offset;
        a->curr_offset = (ptr - (uintptr_t)a->buffer) + size;
        memset(ptr, 0, size);
    }

    return ptr;
}

void* amrm_arena_resize(amrm_Arena *a, void *old, size_t old_size, size_t new_size) {
    return amrm_arena_resize_aligned(a, old, old_size, new_size, AMRM_DEFAULT_ALIGNMENT);
}

void* amrm_arena_resize_aligned(amrm_Arena *a, void *old, 
                                size_t old_size, size_t new_size, size_t alignment) {
    unsigned char *mem_old = (unsigned char*)old;
    void *ptr = NULL;

    assert(amrm__is_power_of_two(alignment));
    assert(old >= a->buffer && old + old_size < a->buffer + a->capacity); 

    ptr = old;
    if (a->buffer + a->curr_offset == old) {
        // if last element, just extend and update state
        size_t curr_offset_new = a->prev_offset + new_size;
        if (curr_offset_new <= a->capacity) {
            a->curr_offset = curr_offset_new;
        }
    } else {
        uintptr_t curr_ptr = a->buffer + a->curr_offset;
        uintptr_t curr_ptr_aligned = amrm__align_forward(curr_ptr, alignment);
        if (curr_ptr_aligned + new_size <= (uintptr_t)a->buffer + a->capacity) {
            ptr = curr_ptr_aligned;
            memset(ptr, 0, new_size);
            memmove(ptr, old, new_size);
        }
    }
    return ptr;
}

void amrm_arena_clear(amrm_Arena *a) {
    a->prev_offset = 0;
    a->curr_offset = 0;
}

#endif // AMR_ARRAY_IMPLEMENTATION

#endif // __AMR_INCLUDE_AMR_MEMORY_H