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

#include <zencore/commandline.h>
#include <zencore/memory/fmalloc.h>
#include <zencore/memory/mallocansi.h>
#include <zencore/memory/mallocmimalloc.h>
#include <zencore/memory/mallocrpmalloc.h>
#include <zencore/memory/mallocstomp.h>
#include <zencore/memory/memory.h>
#include <zencore/memory/memorytrace.h>
#include <zencore/string.h>

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>
ZEN_THIRD_PARTY_INCLUDES_START
#	include <shellapi.h>  // For command line parsing
ZEN_THIRD_PARTY_INCLUDES_END
#endif

#if ZEN_PLATFORM_LINUX
#	include <stdio.h>
#endif

namespace zen {

enum class MallocImpl
{
	None = 0,
	Ansi,
	Stomp,
	Mimalloc,
	Rpmalloc
};

static int
InitGMalloc()
{
	MallocImpl Malloc	  = MallocImpl::None;
	FMalloc*   InitMalloc = GMalloc;

	// Pick a default base allocator based on availability/platform
	// when using sanitizers, we should use the default/ansi allocator

#if ZEN_OVERRIDE_NEW_DELETE

#	if ZEN_MIMALLOC_ENABLED
	if (Malloc == MallocImpl::None)
	{
		Malloc = MallocImpl::Mimalloc;
	}
#	endif

#	if ZEN_RPMALLOC_ENABLED
	if (Malloc == MallocImpl::None)
	{
		Malloc = MallocImpl::Rpmalloc;
	}
#	endif
#endif

	// Process any command line overrides
	//
	// Note that calls can come into this function before we enter the regular main function
	// and we can therefore not rely on the regular command line parsing for the application

	using namespace std::literals;

	auto ProcessMallocArg = [&](const std::string_view& Arg) {
#if ZEN_RPMALLOC_ENABLED
		if (Arg == "rpmalloc"sv)
		{
			Malloc = MallocImpl::Rpmalloc;
		}
#endif

#if ZEN_MIMALLOC_ENABLED
		if (Arg == "mimalloc"sv)
		{
			Malloc = MallocImpl::Mimalloc;
		}
#endif

		if (Arg == "ansi"sv)
		{
			Malloc = MallocImpl::Ansi;
		}

		if (Arg == "stomp"sv)
		{
			Malloc = MallocImpl::Stomp;
		}
	};

	constexpr std::string_view MallocOption = "--malloc"sv;

	std::function<void(const std::string_view&)> ProcessArg = [&](const std::string_view& Arg) {
		if (Arg.starts_with(MallocOption))
		{
			std::string_view::value_type DelimChar = Arg[MallocOption.length()];

			if (DelimChar == ' ' || DelimChar == '=')
			{
				const std::string_view OptionArgs = Arg.substr(MallocOption.size() + 1);

				IterateCommaSeparatedValue(OptionArgs, ProcessMallocArg);
			}
		}
	};

	IterateCommandlineArgs(ProcessArg);

	switch (Malloc)
	{
#if ZEN_WITH_MALLOC_STOMP
		case MallocImpl::Stomp:
			GMalloc = new FMallocStomp();
			break;
#endif

#if ZEN_RPMALLOC_ENABLED
		case MallocImpl::Rpmalloc:
			GMalloc = new FMallocRpmalloc();
			break;
#endif

#if ZEN_MIMALLOC_ENABLED
		case MallocImpl::Mimalloc:
			GMalloc = new FMallocMimalloc();
			break;
#endif
		default:
			break;
	}

	if (GMalloc == InitMalloc)
	{
		GMalloc = new FMallocAnsi();
	}

	return 1;
}

void
Memory::GCreateMalloc()
{
	static int InitFlag = InitGMalloc();
}

void
Memory::Initialize()
{
	GCreateMalloc();
}

//////////////////////////////////////////////////////////////////////////

void*
Memory::SystemMalloc(size_t Size)
{
	void* Ptr = ::malloc(Size);
	MemoryTrace_Alloc(uint64_t(Ptr), Size, 0, EMemoryTraceRootHeap::SystemMemory);
	return Ptr;
}

void
Memory::SystemFree(void* Ptr)
{
	MemoryTrace_Free(uint64_t(Ptr), EMemoryTraceRootHeap::SystemMemory);
	::free(Ptr);
}

}  // namespace zen

//////////////////////////////////////////////////////////////////////////

static ZEN_NOINLINE bool
InvokeNewHandler(bool NoThrow)
{
	std::new_handler h = std::get_new_handler();

	if (!h)
	{
#if defined(_CPPUNWIND) || defined(__cpp_exceptions)
		if (NoThrow == false)
			throw std::bad_alloc();
#else
		ZEN_UNUSED(NoThrow);
#endif
		return false;
	}
	else
	{
		h();
		return true;
	}
}

//////////////////////////////////////////////////////////////////////////

ZEN_NOINLINE void*
RetryNew(size_t Size, bool NoThrow)
{
	void* Ptr = nullptr;
	while (!Ptr && InvokeNewHandler(NoThrow))
	{
		Ptr = zen::Memory::Malloc(Size, zen::DEFAULT_ALIGNMENT);
	}
	return Ptr;
}

void*
zen_new(size_t Size)
{
	void* Ptr = zen::Memory::Malloc(Size, zen::DEFAULT_ALIGNMENT);

	if (!Ptr) [[unlikely]]
	{
		const bool NoThrow = false;
		return RetryNew(Size, NoThrow);
	}

	return Ptr;
}

void*
zen_new_nothrow(size_t Size) noexcept
{
	void* Ptr = zen::Memory::Malloc(Size, zen::DEFAULT_ALIGNMENT);

	if (!Ptr) [[unlikely]]
	{
		const bool NoThrow = true;
		return RetryNew(Size, NoThrow);
	}

	return Ptr;
}

void*
zen_new_aligned(size_t Size, size_t Alignment)
{
	void* Ptr;

	do
	{
		Ptr = zen::Memory::Malloc(Size, uint32_t(Alignment));
	} while (!Ptr && InvokeNewHandler(/* NoThrow */ false));

	return Ptr;
}

void*
zen_new_aligned_nothrow(size_t Size, size_t Alignment) noexcept
{
	void* Ptr;

	do
	{
		Ptr = zen::Memory::Malloc(Size, uint32_t(Alignment));
	} while (!Ptr && InvokeNewHandler(/* NoThrow */ true));

	return Ptr;
}

void
zen_free(void* Ptr) noexcept
{
	zen::Memory::Free(Ptr);
}

void
zen_free_size(void* Ptr, size_t Size) noexcept
{
	ZEN_UNUSED(Size);
	zen::Memory::Free(Ptr);
}

void
zen_free_size_aligned(void* Ptr, size_t Size, size_t Alignment) noexcept
{
	ZEN_UNUSED(Size, Alignment);
	zen::Memory::Free(Ptr);
}

void
zen_free_aligned(void* Ptr, size_t Alignment) noexcept
{
	ZEN_UNUSED(Alignment);
	zen::Memory::Free(Ptr);
}

// EASTL operator new

#if ZEN_OVERRIDE_NEW_DELETE

void*
operator new[](size_t size, const char* pName, int flags, unsigned debugFlags, const char* file, int line)
{
	ZEN_UNUSED(pName, flags, debugFlags, file, line);
	return zen_new(size);
}

void*
operator new[](size_t	   size,
			   size_t	   alignment,
			   size_t	   alignmentOffset,
			   const char* pName,
			   int		   flags,
			   unsigned	   debugFlags,
			   const char* file,
			   int		   line)
{
	ZEN_UNUSED(alignmentOffset, pName, flags, debugFlags, file, line);

	ZEN_ASSERT_SLOW(alignmentOffset == 0);	// currently not supported

	return zen_new_aligned(size, alignment);
}

#else

void*
operator new[](size_t size, const char* pName, int flags, unsigned debugFlags, const char* file, int line)
{
	ZEN_UNUSED(pName, flags, debugFlags, file, line);
	return ::operator new[](size);
}

void*
operator new[](size_t size,
			   size_t alignment,
			   size_t alignmentOffset,
			   const char* pName,
			   int flags,
			   unsigned debugFlags,
			   const char* file,
			   int line)
{
	ZEN_UNUSED(alignmentOffset, pName, flags, debugFlags, file, line);
	ZEN_ASSERT_SLOW(alignmentOffset == 0);	// currently not supported
	return ::operator new[](size, std::align_val_t(alignment));
}

#endif