constexpr u64 DEFAULT_ALIGNMENT = sizeof(void*);

struct ArenaPool
{
	ArenaPool *next;
	u8        *buffer;
	u64        used;
	u64        length;
};

struct Arena
{
	ArenaPool *pool;
	u64        default_size;
	u64        active_pool;
};

struct TempArena
{
	Arena     *arena;
	ArenaPool *start_pool;
	u64        start_pos;
};

struct Scratch
{
	Arena *arena[2];
	u64    index;
	bool   init;
};

Scratch g_scratch;

void *
Malloc(usize size)
{
	void *ptr = malloc(size);
	memset(ptr, 0, size);
	return ptr;
}

void
AddArenaPool(Arena *arena, usize size)
{
	usize buffer_size = size+sizeof(ArenaPool);
	ArenaPool *pool = (ArenaPool *)Malloc(buffer_size);

	pool->buffer = (u8 *)MovePtrForward(pool, sizeof(ArenaPool));
	pool->length = buffer_size - PtrSub(pool->buffer, pool);

	pool->next  = arena->pool;
	arena->pool = pool;
}

Arena*
CreateArena(usize size)
{
	usize buffer_size = size+sizeof(Arena)+sizeof(ArenaPool);
	void *buffer = Malloc(buffer_size);

	memset(buffer, 0, buffer_size);

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

	arena->default_size = size;
	arena->active_pool  = 0;

	buffer = MovePtrForward(buffer, sizeof(Arena));
	arena->pool = (ArenaPool *)buffer;

	buffer = MovePtrForward(buffer, sizeof(ArenaPool));
	arena->pool->buffer   = (u8 *)buffer;

	arena->pool->length = buffer_size - PtrSub(arena->pool->buffer, arena);

	return arena;
}

TempArena
Begin(Arena* arena)
{
	TempArena temp_arena = { .arena = arena };

	ArenaPool *pool = arena->pool;
	u64 pool_index = 0;
	for(;;)
	{
		assert(pool);

		if(arena->active_pool == pool_index++)
		{
			break;
		}

		pool = pool->next;
	}

	temp_arena.start_pool = pool;
	temp_arena.start_pos  = pool->used;

	return temp_arena;
}

void
End(TempArena *temp_arena)
{
	if(temp_arena->arena)
	{
		temp_arena->start_pool->used = temp_arena->start_pos;
		ArenaPool *pool = temp_arena->start_pool->next;
		while(pool)
		{
			pool->used = 0;
			pool       = pool->next;
		}
	}
}

template<typename T> T*
AllocAlign(Arena* arena, usize size, usize alignment)
{
	T *ptr = NULL;

	uintptr buffer_pos = 0, current = 0, offset = 0;
	ArenaPool *pool = arena->pool;
	u64 selected_pool = 0;
	for(;;)
	{
		if(!pool)
		{
			if(size > arena->default_size)
			{
				size += arena->default_size;
			}

			AddArenaPool(arena, Max(size, arena->default_size));
			pool = arena->pool;
		}

		buffer_pos = (uintptr)pool->buffer;
		current    = buffer_pos + pool->used;
		offset     = AlignPow2(current, alignment) - buffer_pos;

		if(selected_pool >= arena->active_pool && offset+size <= pool->length)
		{
			break;
		}

		pool           = pool->next;
		selected_pool += 1;
	}

	if(selected_pool > arena->active_pool)
	{
		arena->active_pool = selected_pool;
	}

	ptr        = (T *)((uintptr)(pool->buffer) + offset);
	pool->used = offset+size;

	return ptr;	
}

template<typename T> T*
Alloc(Arena* arena)
{
	return AllocAlign<T>(arena, sizeof(T), DEFAULT_ALIGNMENT);
}

template<typename T> Array<T>
Alloc(Arena* arena, u64 length)
{
	Array<T> array;

	array.ptr    = AllocAlign<T>(arena, sizeof(T)*length, DEFAULT_ALIGNMENT);
	array.length = length;

	return array;
}

template<typename T> T *
AllocPtrArray(Arena *arena, u64 length)
{
	return AllocAlign<T>(arena, sizeof(T)*length, DEFAULT_ALIGNMENT);
}

template<typename T> Array<T>
Alloc(TempArena temp_arena, u64 length)
{
	return Alloc<T>(temp_arena.arena, length);
}

template<typename T> T *
Alloc(TempArena temp_arena)
{
	return Alloc<T>(temp_arena.arena);
}

template<typename T> T *
AllocPtrArray(TempArena temp_arena, u64 length)
{
	return AllocPtrArray<T>(temp_arena, length);
}

void
Reset(Arena* arena)
{
	arena->active_pool = 0;
	ArenaPool *pool   = arena->pool;
	while(pool)
	{
		pool->used = 0;
		memset(pool->buffer, 0, pool->length);
		pool       = pool->next;
	}
}

void
BeginScratch(u64 size = 0)
{
	if(!g_scratch.init)
	{
		assert(size);
		g_scratch.arena[0] = CreateArena(size);
		g_scratch.arena[1] = CreateArena(size);
		g_scratch.init     = true;
	}
	else
	{
		g_scratch.index = g_scratch.index%2;
		Reset(g_scratch.arena[g_scratch.index]);
	}
}

template<typename T> T *
ScratchAlloc()
{
	return Alloc<T>(g_scratch.arena[g_scratch.index]);
}

template<typename T> Array<T>
ScratchAlloc(u64 length)
{
	return Alloc<T>(g_scratch.arena[g_scratch.index], length);
}

template<typename T> T *
ScratchAllocPtrArray(u64 length)
{
	return AllocPtrArray<T>(g_scratch.arena[g_scratch.index], length);
}

template<typename T> T*
Alloc()
{
	return (T *)Malloc(sizeof(T));
}

template<typename T> T *
AllocPtrArray(u64 length)
{
	return (T *)Malloc(sizeof(T)*length);
}

template<typename T> Array<T>
Alloc(u64 length)
{
	Array<T> array;

	array.ptr    = (T *)Malloc(sizeof(T)*length);
	array.length = length;

	return array;
}

template<typename T> void
Free(Array<T> *array)
{
	free(array->ptr);

	array->ptr    = NULL;
	array->length = 0;
}

template<typename T> void
Free(T **ptr)
{
	free(*ptr);

	*ptr = NULL;
}