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

#include "zentrack/zentrack.h"

#if ZEN_PLATFORM_WINDOWS
#	include <zenbase/zenbase.h>
#	include <zencore/windows.h>

#	include <detours/detours.h>

#	define ZEN_EXITCODE_INIT_FAILED	20
#	define ZEN_EXITCODE_HOOKS_FAILED	21
#	define ZEN_EXITCODE_PAYLOAD_FAILED 22

__declspec(dllexport) void dummy()	// we need at least one export in order for detours to be happy
{
}

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

#	define DETOURED_FUNCTIONS_KERNELBASE \
		DETOURED_FUNCTION(VirtualAlloc)   \
		DETOURED_FUNCTION(VirtualFree)    \
		DETOURED_FUNCTION(Sleep)

#	define DETOURED_FUNCTIONS DETOURED_FUNCTIONS_KERNELBASE

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

static const _GUID DetoursPayloadGuid = __uuidof(zen::DetoursPayload);

[[noreturn]] void
TerminateCurrentProcess(uint32_t ExitCode)
{
	TerminateProcess(GetCurrentProcess(), ExitCode);
}

#	define DETOURED_FUNCTION(func)         \
		using Decl_##func = decltype(func); \
		Decl_##func* Real_##func;
DETOURED_FUNCTIONS
#	undef DETOURED_FUNCTION

BOOL
Hooked_VirtualFree(LPVOID lpAddress, SIZE_T dwSize, DWORD dwFreeType)
{
	BOOL Result = Real_VirtualFree(lpAddress, dwSize, dwFreeType);

	return Result;
}

LPVOID
Hooked_VirtualAlloc(LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWORD flProtect)
{
	void* Result = Real_VirtualAlloc(lpAddress, dwSize, flAllocationType, flProtect);

	return Result;
}

void
Hooked_Sleep(DWORD dwMilliseconds)
{
	Real_Sleep(dwMilliseconds);
}

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

void
HookFunction(void** RealFunc, void* HookFunc, const char* FunctionName)
{
	if (!*RealFunc)
	{
		return;
	}

	LONG Result = DetourAttach(RealFunc, HookFunc);
	if (Result == NO_ERROR)
	{
		return;
	}

	// TODO: notify any listener about problem
	ZEN_UNUSED(FunctionName);

	ExitProcess(Result);
}

void
HookFunctions()
{
	// First grab the original pre-hook function pointers

#	define DETOURED_FUNCTION(FuncName) Real_##FuncName = (decltype(Real_##FuncName))GetProcAddress(hModule, #	 FuncName);

	if (HMODULE hModule = GetModuleHandleW(L"kernelbase.dll"))
	{
		DETOURED_FUNCTIONS_KERNELBASE
	}

#	undef DETOURED_FUNCTION

	// Then install our hooks

#	define DETOURED_FUNCTION(Func) HookFunction((PVOID*)&Real_##Func, Hooked_##Func, #	  Func);

	DETOURED_FUNCTIONS
}

void
InstallHooks(zen::DetoursPayload& Payload)
{
	ZEN_UNUSED(Payload);

	DetourTransactionBegin();
	DetourUpdateThread(GetCurrentThread());
	HookFunctions();

	if (LONG Result = DetourTransactionCommit(); Result != NO_ERROR)
	{
		ExitProcess(ZEN_EXITCODE_HOOKS_FAILED);
	}
}

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

BOOL WINAPI
DllMain(HINSTANCE, DWORD dwReason, LPVOID)
{
	if (DetourIsHelperProcess())
		return TRUE;

	if (dwReason == DLL_PROCESS_ATTACH)
	{
		// Restore the contents in memory import table after a process was started
		// with DetourCreateProcessWithDllEx or DetourCreateProcessWithDlls
		if (!DetourRestoreAfterWith())
		{
			TerminateCurrentProcess(ZEN_EXITCODE_INIT_FAILED);
		}

		// Grab state
		DWORD				 cbData	 = 0;
		zen::DetoursPayload* Payload = reinterpret_cast<zen::DetoursPayload*>(DetourFindPayloadEx(DetoursPayloadGuid, &cbData));
		if (!Payload || cbData != sizeof(zen::DetoursPayload))
		{
			TerminateCurrentProcess(ZEN_EXITCODE_PAYLOAD_FAILED);
		}

		InstallHooks(*Payload);
	}
	else if (dwReason == DLL_PROCESS_DETACH)
	{
	}

	return TRUE;
}

void
foo()
{
	DetourTransactionBegin();

	return;
}
#endif