// 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 WpArena {
	u8 *buf;
	u8 *offset;
	u8 *prev_offset;
	u64 capacity;
	ArenaStorageType type;
	b8 committed;

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

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

	*arena           = (WpArena *)buffer;
	WpArena *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(WpArena);
	arena_ptr->type        = ARENA_STORAGE_TYPE_BUFFER;
	arena_ptr->committed   = true;

	return true;
}

b8 wpMemArenaInitAllocatedCustom(WpArena **arena, u64 base_capacity, WpMemAllocFlags flags, b8 zero_buffer) {
	if (!arena || *arena || base_capacity == 0) {
		return false;
	}

	u64 size       = sizeof(WpArena) + (base_capacity >= ARENA_MINIMUM_CAPACITY ? base_capacity : ARENA_MINIMUM_CAPACITY);
	u64 alloc_size = wpMiscUtilsU64RoundUpPow2(size);
	u8 *allocated  = (u8 *)wpOsMemAlloc(NULL, alloc_size, WP_MEM_ACCESS_READ_WRITE, flags,
	                                         zero_buffer ? WP_MEM_INIT_INITIALISED : WP_MEM_INIT_UNINITIALISED);
	if (!allocated) {
		return false;
	}

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

#ifdef WP_PLATFORM_WINDOWS
	if (!committed) {
		wpOsMemAlloc(allocated, sizeof(WpArena), WP_MEM_ACCESS_READ_WRITE, WP_MEM_ALLOC_COMMIT,
		                  WP_MEM_INIT_INITIALISED);
	}
#endif // ifdef WP_PLATFORM_WINDOWS

	if (!wpMemArenaInitBuffer(arena, allocated, alloc_size)) {
		wpMemArenaDestroy(arena);
		return false;
	}

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

	return true;
}

void *wpMemArenaAlloc(WpArena *arena, u64 size) {
	return wpMemArenaAllocAligned(arena, size, DEFAULT_ALIGNMENT);
}

void *wpMemArenaAllocAligned(WpArena *arena, u64 size, u64 alignment) {
	wpDebugAssert(arena != NULL, "`arena` should not be NULL");

	u8 *alloc_start = arena->offset;

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

	arena->offset = output + size;

#ifdef WP_PLATFORM_WINDOWS
	if (arena->type == ARENA_STORAGE_TYPE_ALLOCATED && !(arena->committed)) {
		wpOsMemAlloc(alloc_start, (uptr)(arena->offset) - (uptr)(alloc_start),
		                  WP_MEM_ACCESS_READ_WRITE, WP_MEM_ALLOC_COMMIT,
		                  WP_MEM_INIT_UNINITIALISED);
	}
#endif // ifdef WP_PLATFORM_WINDOWS

	memset(output, 0, size);

	return (void *)output;
}

void *wpMemArenaRealloc(WpArena *arena, void *ptr, u64 old_size, u64 new_size) {
	wpDebugAssert(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 = wpMemArenaAlloc(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 *wpMemArenaReallocAligned(WpArena *arena, void *ptr, u64 old_size, u64 new_size, u64 alignment) {
	wpDebugAssert(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 = wpMemArenaAllocAligned(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 wpMemArenaTempBegin(WpArena *arena) {
	wpDebugAssert(arena != NULL, "`arena` should not be NULL");

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

	arena->prev_offset = arena->offset;
}

void wpMemArenaTempEnd(WpArena *arena) {
	wpDebugAssert(arena != NULL, "`arena` should not be NULL");

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

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

void wpMemArenaClear(WpArena *arena) {
	wpDebugAssert(arena != NULL, "`arena` should not be NULL");

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

void wpMemArenaDestroy(WpArena **arena) {
	wpDebugAssert(arena != NULL && (*arena) != NULL, "`arena` double pointer is not valid");

	WpArena *arena_ptr = *arena;

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

	*arena = NULL;
}