#include "arena_allocator.h"

#include <assert.h>
#include <stddef.h>
#include <stdbool.h>

SandArena sand_create_arena(SandAllocator *a) {
	return (SandArena) {
		.parent = a,
		.current_region = NULL,
	};
}

#define MAX(a, b) ({ \
		__auto_type _a = (a);\
		__auto_type _b = (b);\
		_a > _b ? _a : _b;\
		})

#define REGION_DEFUALT_SIZE ((size_t)2048)

static SandRegion *allocate_new_region(SandAllocator *a, size_t size, size_t alignment) {
	size_t padded_size = size + (alignment - 1);
	size_t capacity = MAX(padded_size, REGION_DEFUALT_SIZE) ;
	size_t region_size = offsetof(SandRegion, data) + capacity;
	SandRegion *result = sand_allocate_aligned(a, region_size, _Alignof(SandRegion));
	if (result == NULL) {
		return NULL;
	}
	result->capacity = capacity;
	result->used = 0;
	return result;
}

static void *allocate(size_t size, size_t alignment, void *user_data) {
	SandArena *arena = user_data;

	// This is the first allocation made within the arena. Let's allocate the initial region.
	if (arena->current_region == NULL) {
		SandRegion *first_region = allocate_new_region(arena->parent, size, alignment);
		if (first_region == NULL) {
			return NULL;
		}
		first_region->prev = NULL;
		arena->current_region = first_region;
	}

	SandRegion *iter = arena->current_region;
	while (true) { // Loop looking for a region.
		void *aligned_ptr = sand_align_pointer_forward(iter->data + iter->used, alignment);
		void *end_ptr = iter->data + iter->capacity;
		if (aligned_ptr + size > end_ptr) { // This pointer region doesn't have room; allocate a new one and try again.
			SandRegion *new_region = allocate_new_region(arena->parent, size, alignment);
			if (new_region == NULL) {
				return NULL;
			}

			new_region->prev = arena->current_region;
			arena->current_region = iter;

			iter = new_region; // Next time is guaranteed to succeed.
			continue;
		} else { // Yay, we found a spot!
			iter->used += aligned_ptr - (void *)iter->data;
			return aligned_ptr;
		}
	}

	assert(false); // unreachable
}

static void deallocate(void *old_ptr, size_t old_size, void *user_data) {
	(void)old_ptr;
	(void)old_size;
	(void)user_data;
	return; // No-op, as all allocations are freed when the arena is freed.
}

// Returns a `SandAllocator` which is valid as long as the arena is.
// The arena must not be moved afterwards, as the allocator retains a pointer to the arena.
SandAllocator sand_get_allocator_for_arena(SandArena *arena) {
	return (SandAllocator) {
		.allocate = allocate,
		.reallocate = NULL, // TODO: Special case for if it is the last allocation of the current region.
		.deallocate = deallocate,
		.user_data = arena,
	};
}

void sand_destroy_arena(SandArena *arena) {
	SandRegion *iter = arena->current_region;
	while (iter != NULL) {
		SandRegion *next = iter->prev;
		size_t region_size = offsetof(SandRegion, data) + iter->capacity;
		sand_deallocate(arena->parent, iter, region_size);
		iter = next;
	}
}