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

#pragma once

#include <zenbase/zenbase.h>

#if ZEN_PLATFORM_WINDOWS && !defined(PLATFORM_SUPPORTS_TRACE_WIN32_VIRTUAL_MEMORY_HOOKS)
#	define PLATFORM_SUPPORTS_TRACE_WIN32_VIRTUAL_MEMORY_HOOKS 1
#endif

#ifndef PLATFORM_SUPPORTS_TRACE_WIN32_VIRTUAL_MEMORY_HOOKS
#	define PLATFORM_SUPPORTS_TRACE_WIN32_VIRTUAL_MEMORY_HOOKS 0
#endif

#if PLATFORM_SUPPORTS_TRACE_WIN32_VIRTUAL_MEMORY_HOOKS
#	include <zencore/windows.h>

namespace zen {

class FVirtualWinApiHooks
{
public:
	static void Initialize(bool bInLight);

private:
	FVirtualWinApiHooks();
	static bool			 bLight;
	static LPVOID WINAPI VmAlloc(LPVOID Address, SIZE_T Size, DWORD Type, DWORD Protect);
	static LPVOID WINAPI VmAllocEx(HANDLE Process, LPVOID Address, SIZE_T Size, DWORD Type, DWORD Protect);
#	if (NTDDI_VERSION >= NTDDI_WIN10_RS4)
	static PVOID WINAPI VmAlloc2(HANDLE					 Process,
								 PVOID					 BaseAddress,
								 SIZE_T					 Size,
								 ULONG					 AllocationType,
								 ULONG					 PageProtection,
								 MEM_EXTENDED_PARAMETER* ExtendedParameters,
								 ULONG					 ParameterCount);
	static PVOID(WINAPI* VmAlloc2Orig)(HANDLE, PVOID, SIZE_T, ULONG, ULONG, MEM_EXTENDED_PARAMETER*, ULONG);
	typedef PVOID(__stdcall* FnVirtualAlloc2)(HANDLE, PVOID, SIZE_T, ULONG, ULONG, MEM_EXTENDED_PARAMETER*, ULONG);
#	else
	static LPVOID WINAPI VmAlloc2(HANDLE Process,
								  LPVOID BaseAddress,
								  SIZE_T Size,
								  ULONG	 AllocationType,
								  ULONG	 PageProtection,
								  void*	 ExtendedParameters,
								  ULONG	 ParameterCount);
	static LPVOID(WINAPI* VmAlloc2Orig)(HANDLE, LPVOID, SIZE_T, ULONG, ULONG, /*MEM_EXTENDED_PARAMETER* */ void*, ULONG);
	typedef LPVOID(__stdcall* FnVirtualAlloc2)(HANDLE, LPVOID, SIZE_T, ULONG, ULONG, /* MEM_EXTENDED_PARAMETER* */ void*, ULONG);
#	endif
	static BOOL WINAPI VmFree(LPVOID Address, SIZE_T Size, DWORD Type);
	static BOOL WINAPI VmFreeEx(HANDLE Process, LPVOID Address, SIZE_T Size, DWORD Type);
	static LPVOID(WINAPI* VmAllocOrig)(LPVOID, SIZE_T, DWORD, DWORD);
	static LPVOID(WINAPI* VmAllocExOrig)(HANDLE, LPVOID, SIZE_T, DWORD, DWORD);
	static BOOL(WINAPI* VmFreeOrig)(LPVOID, SIZE_T, DWORD);
	static BOOL(WINAPI* VmFreeExOrig)(HANDLE, LPVOID, SIZE_T, DWORD);
};

}  // namespace zen

#endif