// Copyright Epic Games, Inc. All Rights Reserved.

#include <zencore/memory/memoryarena.h>

#include <cstring>

namespace zen {

MemoryArena::~MemoryArena()
{
	for (auto Chunk : m_Chunks)
		delete[] Chunk;
}

void*
MemoryArena::AllocateAligned(size_t ByteCount, size_t align)
{
	if (ByteCount == 0)
	{
		return nullptr;
	}

	void* Ptr = nullptr;

	m_Lock.WithExclusiveLock([&] {
		size_t AlignedOffset = (m_CurrentOffset + (align - 1)) & ~(align - 1);

		if (m_CurrentChunk == nullptr || AlignedOffset + ByteCount > ChunkSize)
		{
			uint8_t* NewChunk = new uint8_t[ChunkSize];
			if (!NewChunk)
			{
				return;
			}

			m_Chunks.push_back(NewChunk);
			m_CurrentChunk = NewChunk;
			AlignedOffset  = 0;
		}

		Ptr				= m_CurrentChunk + AlignedOffset;
		m_CurrentOffset = AlignedOffset + ByteCount;
	});

	return Ptr;
}

void*
MemoryArena::AllocateAlignedWithOffset(size_t ByteCount, size_t align, size_t offset)
{
	if (ByteCount == 0)
	{
		return nullptr;
	}

	void* Ptr = nullptr;

	m_Lock.WithExclusiveLock([&] {
		size_t AlignedOffset = (m_CurrentOffset + (align - 1) + offset) & ~(align - 1);

		if (m_CurrentChunk == nullptr || AlignedOffset + ByteCount > ChunkSize)
		{
			uint8_t* NewChunk = new uint8_t[ChunkSize];
			if (!NewChunk)
			{
				return;
			}

			m_Chunks.push_back(NewChunk);
			m_CurrentChunk = NewChunk;
			AlignedOffset  = offset;
		}

		Ptr				= m_CurrentChunk + AlignedOffset;
		m_CurrentOffset = AlignedOffset + ByteCount;
	});

	return Ptr;
}

void*
MemoryArena::Allocate(size_t Size)
{
	if (Size == 0)
	{
		return nullptr;
	}

	void*			 Ptr	   = nullptr;
	constexpr size_t Alignment = alignof(std::max_align_t);

	m_Lock.WithExclusiveLock([&] {
		size_t AlignedOffset = (m_CurrentOffset + Alignment - 1) & ~(Alignment - 1);

		if (m_CurrentChunk == nullptr || AlignedOffset + Size > ChunkSize)
		{
			uint8_t* NewChunk = new uint8_t[ChunkSize];
			if (!NewChunk)
			{
				return;
			}

			m_Chunks.push_back(NewChunk);
			m_CurrentChunk = NewChunk;
			AlignedOffset  = 0;
		}

		Ptr				= m_CurrentChunk + AlignedOffset;
		m_CurrentOffset = AlignedOffset + Size;
	});

	return Ptr;
}

const char*
MemoryArena::DuplicateString(std::string_view Str)
{
	const size_t Len	= Str.size();
	char*		 NewStr = static_cast<char*>(Allocate(Len + 1));
	if (NewStr)
	{
		memcpy(NewStr, Str.data(), Len);
		NewStr[Len] = '\0';
	}
	return NewStr;
}

}  // namespace zen