// vim:fileencoding=utf-8:foldmethod=marker

#include "mem_arena.h"
#include "../../mem/mem_os.h"
#include "../../../common/aliases/aliases.h"
#include "../../../common/assert/assert.h"
#include "../../../common/misc/misc_utils.h"
#include "../../../base/mem/utils/mem_utils.h"
#include <string.h>

#ifndef DEFAULT_ALIGNMENT
// Why 2 * sizeof(void *) instead of sizeof(void *)
// https://handmade.network/forums/t/6860-alignment_arena_allocator
#define DEFAULT_ALIGNMENT (2 * sizeof(void *))
#endif /* ifndef DEFAULT_ALIGNMENT */

#define ARENA_MINIMUM_CAPACITY KiB(16) // Allocate minimum of 4 pages

typedef enum {
	ARENA_STORAGE_TYPE_ALLOCATED,
	ARENA_STORAGE_TYPE_BUFFER,
} ArenaStorageType;

struct Arena {
	u8 *buf;
	u8 *offset;
	u8 *prev_offset;
	u64 capacity;
	ArenaStorageType type;
	b8 committed;

	wapp_misc_utils_reserve_padding(sizeof(u8 *) * 3 + sizeof(u64) + sizeof(ArenaStorageType) + sizeof(b8));
};

b8 wapp_mem_arena_init_buffer(Arena **arena, u8 *buffer, u64 buffer_size) {
	if (!arena || *arena || buffer_size < sizeof(Arena)) {
		return false;
	}

	*arena           = (Arena *)buffer;
	Arena *arena_ptr = *arena;

	arena_ptr->buf         = (u8 *)(arena_ptr + 1);
	arena_ptr->offset      = arena_ptr->buf;
	arena_ptr->prev_offset = NULL;
	arena_ptr->capacity    = buffer_size - sizeof(Arena);
	arena_ptr->type        = ARENA_STORAGE_TYPE_BUFFER;
	arena_ptr->committed   = true;

	return true;
}

b8 wapp_mem_arena_init_allocated_custom(Arena **arena, u64 base_capacity, MemAllocFlags flags, b8 zero_buffer) {
	if (!arena || *arena || base_capacity == 0) {
		return false;
	}

	u64 size       = sizeof(Arena) + (base_capacity >= ARENA_MINIMUM_CAPACITY ? base_capacity : ARENA_MINIMUM_CAPACITY);
	u64 alloc_size = wapp_misc_utils_u64_round_up_pow2(size);
	u8 *allocated  = (u8 *)wapp_os_mem_alloc(NULL, alloc_size, WAPP_MEM_ACCESS_READ_WRITE, flags,
	                                         zero_buffer ? WAPP_MEM_INIT_INITIALISED : WAPP_MEM_INIT_UNINITIALISED);
	if (!allocated) {
		return false;
	}

	b8 committed = (flags & WAPP_MEM_ALLOC_COMMIT) == WAPP_MEM_ALLOC_COMMIT;

#ifdef WAPP_PLATFORM_WINDOWS
	if (!committed) {
		wapp_os_mem_alloc(allocated, sizeof(Arena), WAPP_MEM_ACCESS_READ_WRITE, WAPP_MEM_ALLOC_COMMIT,
		                  WAPP_MEM_INIT_INITIALISED);
	}
#endif // ifdef WAPP_PLATFORM_WINDOWS

	if (!wapp_mem_arena_init_buffer(arena, allocated, alloc_size)) {
		wapp_mem_arena_destroy(arena);
		return false;
	}

	Arena *arena_ptr     = *arena;
	arena_ptr->type      = ARENA_STORAGE_TYPE_ALLOCATED;
	arena_ptr->committed = committed;

	return true;
}

void *wapp_mem_arena_alloc(Arena *arena, u64 size) {
	return wapp_mem_arena_alloc_aligned(arena, size, DEFAULT_ALIGNMENT);
}

void *wapp_mem_arena_alloc_aligned(Arena *arena, u64 size, u64 alignment) {
	wapp_debug_assert(arena != NULL, "`arena` should not be NULL");

	u8 *alloc_start = arena->offset;

	u8 *output = wapp_mem_util_align_forward((void *)alloc_start, alignment);
	if (output + size >= arena->buf + arena->capacity) {
		return NULL;
	}

	arena->offset = output + size;

#ifdef WAPP_PLATFORM_WINDOWS
	if (arena->type == ARENA_STORAGE_TYPE_ALLOCATED && !(arena->committed)) {
		wapp_os_mem_alloc(alloc_start, (uptr)(arena->offset) - (uptr)(alloc_start),
		                  WAPP_MEM_ACCESS_READ_WRITE, WAPP_MEM_ALLOC_COMMIT,
		                  WAPP_MEM_INIT_UNINITIALISED);
	}
#endif // ifdef WAPP_PLATFORM_WINDOWS

	memset(output, 0, size);

	return (void *)output;
}

void *wapp_mem_arena_realloc(Arena *arena, void *ptr, u64 old_size, u64 new_size) {
	wapp_debug_assert(arena != NULL, "`arena` should not be NULL");

	if ((u8*)ptr < arena->buf || (u8*)ptr > arena->offset ||
		  arena->offset + new_size >= arena->buf + arena->capacity) {
		return NULL;
	}

	void *new_ptr = wapp_mem_arena_alloc(arena, new_size);
	if (!new_ptr) {
		return NULL;
	}

	u64 copy_size = new_size <= old_size ? new_size : old_size;
	memcpy(new_ptr, ptr, copy_size);

	return new_ptr;
}

void *wapp_mem_arena_realloc_aligned(Arena *arena, void *ptr, u64 old_size, u64 new_size, u64 alignment) {
	wapp_debug_assert(arena != NULL, "`arena` should not be NULL");

	if ((u8*)ptr < arena->buf || (u8*)ptr > arena->offset ||
		  arena->offset + new_size >= arena->buf + arena->capacity) {
		return NULL;
	}

	void *new_ptr = wapp_mem_arena_alloc_aligned(arena, new_size, alignment);
	if (!new_ptr) {
		return NULL;
	}

	u64 copy_size = new_size <= old_size ? new_size : old_size;
	memcpy(new_ptr, ptr, copy_size);

	return new_ptr;
}

void wapp_mem_arena_temp_begin(Arena *arena) {
	wapp_debug_assert(arena != NULL, "`arena` should not be NULL");

	if (arena->prev_offset != NULL) {
		return;
	}

	arena->prev_offset = arena->offset;
}

void wapp_mem_arena_temp_end(Arena *arena) {
	wapp_debug_assert(arena != NULL, "`arena` should not be NULL");

	if (arena->prev_offset == NULL) {
		return;
	}

	arena->offset      = arena->prev_offset;
	arena->prev_offset = NULL;
}

void wapp_mem_arena_clear(Arena *arena) {
	wapp_debug_assert(arena != NULL, "`arena` should not be NULL");

	memset(arena->buf, 0, arena->offset - arena->buf);
	arena->offset = arena->buf;
}

void wapp_mem_arena_destroy(Arena **arena) {
	wapp_debug_assert(arena != NULL && (*arena) != NULL, "`arena` double pointer is not valid");

	Arena *arena_ptr = *arena;

	if (arena_ptr->type == ARENA_STORAGE_TYPE_ALLOCATED) {
		wapp_os_mem_free(*arena, sizeof(Arena) + arena_ptr->capacity);
	}

	*arena = NULL;
}