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

#include <zencore/memory/memoryarena.h>

#include <algorithm>
#include <cstring>

namespace zen {

MemoryArena::MemoryArena(size_t ChunkSize) : m_ChunkSize(ChunkSize == 0 ? kDefaultChunkSize : ChunkSize)
{
}

MemoryArena::~MemoryArena()
{
	for (const Chunk& C : m_Chunks)
	{
		delete[] C.Base;
	}
}

// An allocation whose worst-case footprint won't fit inside a freshly
// allocated standard-sized chunk is "oversize" — we route it to its own
// dedicated chunk so it doesn't waste the rest of the regular chunk's
// space, and so a subsequent small allocation still has a near-empty
// regular chunk to bump-pointer into.
//
// Note: WorstCaseBytes is the upper bound on bytes the allocation could
// consume from a freshly aligned chunk start. The actual returned
// pointer + size fits within this footprint by construction.
bool
MemoryArena::IsOversizeRequest(size_t WorstCaseBytes) const
{
	return WorstCaseBytes > m_ChunkSize;
}

uint8_t*
MemoryArena::AllocateDedicatedChunk(size_t Capacity)
{
	uint8_t* NewChunk = new (std::nothrow) uint8_t[Capacity];
	if (!NewChunk)
	{
		return nullptr;
	}
	m_Chunks.push_back({NewChunk, Capacity});
	return NewChunk;
}

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

	void* Ptr = nullptr;

	m_Lock.WithExclusiveLock([&] {
		// Oversize fast path — give the request its own chunk, leave the
		// regular bump pointer untouched so subsequent small allocs still
		// reuse it.
		const size_t WorstCase = ByteCount + Align - 1;
		if (IsOversizeRequest(WorstCase))
		{
			if (uint8_t* Dedicated = AllocateDedicatedChunk(WorstCase))
			{
				const size_t Aligned = (reinterpret_cast<uintptr_t>(Dedicated) + (Align - 1)) & ~(uintptr_t(Align - 1));
				Ptr					 = reinterpret_cast<void*>(Aligned);
			}
			return;
		}

		size_t AlignedOffset = (m_CurrentOffset + (Align - 1)) & ~(Align - 1);
		if (m_CurrentChunk == nullptr || AlignedOffset + ByteCount > m_CurrentChunkCapacity)
		{
			uint8_t* NewChunk = AllocateDedicatedChunk(m_ChunkSize);
			if (!NewChunk)
			{
				return;
			}
			m_CurrentChunk		   = NewChunk;
			m_CurrentChunkCapacity = m_ChunkSize;
			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([&] {
		const size_t WorstCase = ByteCount + Align - 1 + Offset;
		if (IsOversizeRequest(WorstCase))
		{
			if (uint8_t* Dedicated = AllocateDedicatedChunk(WorstCase))
			{
				const uintptr_t Base	= reinterpret_cast<uintptr_t>(Dedicated);
				const uintptr_t Aligned = (Base + (Align - 1) + Offset) & ~(uintptr_t(Align - 1));
				Ptr						= reinterpret_cast<void*>(Aligned);
			}
			return;
		}

		size_t AlignedOffset = (m_CurrentOffset + (Align - 1) + Offset) & ~(Align - 1);
		if (m_CurrentChunk == nullptr || AlignedOffset + ByteCount > m_CurrentChunkCapacity)
		{
			uint8_t* NewChunk = AllocateDedicatedChunk(m_ChunkSize);
			if (!NewChunk)
			{
				return;
			}
			m_CurrentChunk		   = NewChunk;
			m_CurrentChunkCapacity = m_ChunkSize;
			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([&] {
		const size_t WorstCase = Size + Alignment - 1;
		if (IsOversizeRequest(WorstCase))
		{
			if (uint8_t* Dedicated = AllocateDedicatedChunk(WorstCase))
			{
				const uintptr_t Aligned = (reinterpret_cast<uintptr_t>(Dedicated) + Alignment - 1) & ~(uintptr_t(Alignment - 1));
				Ptr						= reinterpret_cast<void*>(Aligned);
			}
			return;
		}

		size_t AlignedOffset = (m_CurrentOffset + Alignment - 1) & ~(Alignment - 1);
		if (m_CurrentChunk == nullptr || AlignedOffset + Size > m_CurrentChunkCapacity)
		{
			uint8_t* NewChunk = AllocateDedicatedChunk(m_ChunkSize);
			if (!NewChunk)
			{
				return;
			}
			m_CurrentChunk		   = NewChunk;
			m_CurrentChunkCapacity = m_ChunkSize;
			AlignedOffset		   = 0;
		}

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

	return Ptr;
}

size_t
MemoryArena::GetReservedBytes() const
{
	size_t Total = 0;
	m_Lock.WithSharedLock([&] {
		for (const Chunk& Item : m_Chunks)
		{
			Total += Item.Capacity;
		}
	});
	return Total;
}

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