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

#include <zenbase/zenbase.h>
#include <zencore/memory/llm.h>
#include <zencore/memory/memorytrace.h>
#include <zencore/memory/tagtrace.h>

#if ZEN_PLATFORM_WINDOWS && ZEN_WITH_TRACE
#	define PLATFORM_SUPPORTS_TRACE_WIN32_MODULE_DIAGNOSTICS 1
#else
#	define PLATFORM_SUPPORTS_TRACE_WIN32_MODULE_DIAGNOSTICS 0
#endif

#include "moduletrace_events.h"

#if PLATFORM_SUPPORTS_TRACE_WIN32_MODULE_DIAGNOSTICS

#	include <zencore/windows.h>

ZEN_THIRD_PARTY_INCLUDES_START
#	include <winternl.h>
ZEN_THIRD_PARTY_INCLUDES_END

#	include <zencore/trace.h>

#	include <array>

namespace zen {

class FMalloc;

typedef uint32_t HeapId;

////////////////////////////////////////////////////////////////////////////////
struct FNtDllFunction
{
	FARPROC Addr;

	FNtDllFunction(const char* Name)
	{
		HMODULE NtDll = LoadLibraryW(L"ntdll.dll");
		ZEN_ASSERT(NtDll);
		Addr = GetProcAddress(NtDll, Name);
	}

	template<typename... ArgTypes>
	unsigned int operator()(ArgTypes... Args)
	{
		typedef unsigned int(NTAPI * Prototype)(ArgTypes...);
		return (Prototype((void*)Addr))(Args...);
	}
};

////////////////////////////////////////////////////////////////////////////////7777
class FModuleTrace
{
public:
	typedef void (*SubscribeFunc)(bool, void*, const char16_t*);

	FModuleTrace(FMalloc* InMalloc);
	~FModuleTrace();
	static FModuleTrace* Get();
	void				 Initialize();
	void				 Subscribe(SubscribeFunc Function);

private:
	void				 OnDllLoaded(const UNICODE_STRING& Name, uintptr_t Base);
	void				 OnDllUnloaded(uintptr_t Base);
	void				 OnDllNotification(unsigned int Reason, const void* DataPtr);
	static FModuleTrace* Instance;
	SubscribeFunc		 Subscribers[64];
	int					 SubscriberCount = 0;
	void*				 CallbackCookie	 = nullptr;
	HeapId				 ProgramHeapId	 = 0;
};

////////////////////////////////////////////////////////////////////////////////
FModuleTrace* FModuleTrace::Instance = nullptr;

////////////////////////////////////////////////////////////////////////////////
FModuleTrace::FModuleTrace(FMalloc* InMalloc)
{
	ZEN_UNUSED(InMalloc);
	Instance = this;
}

////////////////////////////////////////////////////////////////////////////////
FModuleTrace::~FModuleTrace()
{
	if (CallbackCookie)
	{
		FNtDllFunction UnregisterFunc("LdrUnregisterDllNotification");
		UnregisterFunc(CallbackCookie);
	}
}

////////////////////////////////////////////////////////////////////////////////
FModuleTrace*
FModuleTrace::Get()
{
	return Instance;
}

////////////////////////////////////////////////////////////////////////////////
void
FModuleTrace::Initialize()
{
	using namespace UE::Trace;

	ProgramHeapId = MemoryTrace_HeapSpec(SystemMemory, u"Module", EMemoryTraceHeapFlags::None);

	UE_TRACE_LOG(Diagnostics, ModuleInit, ModuleChannel, sizeof(char) * 3)
		<< ModuleInit.SymbolFormat("pdb", 3) << ModuleInit.ModuleBaseShift(uint8(0));

	// Register for DLL load/unload notifications.
	auto Thunk = [](ULONG Reason, const void* Data, void* Context) {
		auto* Self = (FModuleTrace*)Context;
		Self->OnDllNotification(Reason, Data);
	};

	typedef void(CALLBACK * ThunkType)(ULONG, const void*, void*);
	auto ThunkImpl = ThunkType(Thunk);

	FNtDllFunction RegisterFunc("LdrRegisterDllNotification");
	RegisterFunc(0, ThunkImpl, this, &CallbackCookie);

	// Enumerate already loaded modules.
	const TEB*		  ThreadEnvBlock  = NtCurrentTeb();
	const PEB*		  ProcessEnvBlock = ThreadEnvBlock->ProcessEnvironmentBlock;
	const LIST_ENTRY* ModuleIter	  = ProcessEnvBlock->Ldr->InMemoryOrderModuleList.Flink;
	const LIST_ENTRY* ModuleIterEnd	  = ModuleIter->Blink;
	do
	{
		const auto& ModuleData = *(LDR_DATA_TABLE_ENTRY*)(ModuleIter - 1);
		if (ModuleData.DllBase == 0)
		{
			break;
		}

		OnDllLoaded(ModuleData.FullDllName, UPTRINT(ModuleData.DllBase));
		ModuleIter = ModuleIter->Flink;
	} while (ModuleIter != ModuleIterEnd);
}

////////////////////////////////////////////////////////////////////////////////
void
FModuleTrace::Subscribe(SubscribeFunc Function)
{
	ZEN_ASSERT(SubscriberCount < ZEN_ARRAY_COUNT(Subscribers));
	Subscribers[SubscriberCount++] = Function;
}

////////////////////////////////////////////////////////////////////////////////
void
FModuleTrace::OnDllNotification(unsigned int Reason, const void* DataPtr)
{
	enum
	{
		LDR_DLL_NOTIFICATION_REASON_LOADED	 = 1,
		LDR_DLL_NOTIFICATION_REASON_UNLOADED = 2,
	};

	struct FNotificationData
	{
		uint32_t			  Flags;
		const UNICODE_STRING& FullPath;
		const UNICODE_STRING& BaseName;
		uintptr_t			  Base;
	};
	const auto& Data = *(FNotificationData*)DataPtr;

	switch (Reason)
	{
		case LDR_DLL_NOTIFICATION_REASON_LOADED:
			OnDllLoaded(Data.FullPath, Data.Base);
			break;
		case LDR_DLL_NOTIFICATION_REASON_UNLOADED:
			OnDllUnloaded(Data.Base);
			break;
	}
}

////////////////////////////////////////////////////////////////////////////////
void
FModuleTrace::OnDllLoaded(const UNICODE_STRING& Name, UPTRINT Base)
{
	const auto*					 DosHeader		= (IMAGE_DOS_HEADER*)Base;
	const auto*					 NtHeaders		= (IMAGE_NT_HEADERS*)(Base + DosHeader->e_lfanew);
	const IMAGE_OPTIONAL_HEADER& OptionalHeader = NtHeaders->OptionalHeader;
	uint8_t						 ImageId[20];

	// Find the guid and age of the binary, used to match debug files
	const IMAGE_DATA_DIRECTORY& DebugInfoEntry = OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG];
	const auto*					DebugEntries   = (IMAGE_DEBUG_DIRECTORY*)(Base + DebugInfoEntry.VirtualAddress);
	for (uint32_t i = 0, n = DebugInfoEntry.Size / sizeof(DebugEntries[0]); i < n; ++i)
	{
		const IMAGE_DEBUG_DIRECTORY& Entry = DebugEntries[i];
		if (Entry.Type == IMAGE_DEBUG_TYPE_CODEVIEW)
		{
			struct FCodeView7
			{
				uint32_t Signature;
				uint32_t Guid[4];
				uint32_t Age;
			};

			if (Entry.SizeOfData < sizeof(FCodeView7))
			{
				continue;
			}

			const auto* CodeView7 = (FCodeView7*)(Base + Entry.AddressOfRawData);
			if (CodeView7->Signature != 'SDSR')
			{
				continue;
			}

			memcpy(ImageId, (uint8_t*)&CodeView7->Guid, sizeof(uint32_t) * 4);
			memcpy(&ImageId[16], (uint8_t*)&CodeView7->Age, sizeof(uint32_t));
			break;
		}
	}

	// Note: UNICODE_STRING.Length is the size in bytes of the string buffer.
	UE_TRACE_LOG(Diagnostics, ModuleLoad, ModuleChannel, uint32_t(Name.Length + sizeof(ImageId)))
		<< ModuleLoad.Name((const char16_t*)Name.Buffer, Name.Length / 2) << ModuleLoad.Base(uint64_t(Base))
		<< ModuleLoad.Size(OptionalHeader.SizeOfImage) << ModuleLoad.ImageId(ImageId, uint32_t(sizeof(ImageId)));

#	if UE_MEMORY_TRACE_ENABLED
	{
		ZEN_MEMSCOPE(ELLMTag::ProgramSize);
		MemoryTrace_Alloc(Base, OptionalHeader.SizeOfImage, 4 * 1024, EMemoryTraceRootHeap::SystemMemory);
		MemoryTrace_MarkAllocAsHeap(Base, ProgramHeapId);
		MemoryTrace_Alloc(Base, OptionalHeader.SizeOfImage, 4 * 1024, EMemoryTraceRootHeap::SystemMemory);
	}
#	endif	// UE_MEMORY_TRACE_ENABLED

	for (int i = 0; i < SubscriberCount; ++i)
	{
		Subscribers[i](true, (void*)Base, (const char16_t*)Name.Buffer);
	}
}

////////////////////////////////////////////////////////////////////////////////
void
FModuleTrace::OnDllUnloaded(UPTRINT Base)
{
#	if UE_MEMORY_TRACE_ENABLED
	MemoryTrace_Free(Base, EMemoryTraceRootHeap::SystemMemory);
	MemoryTrace_UnmarkAllocAsHeap(Base, ProgramHeapId);
	MemoryTrace_Free(Base, EMemoryTraceRootHeap::SystemMemory);
#	endif	// UE_MEMORY_TRACE_ENABLED

	UE_TRACE_LOG(Diagnostics, ModuleUnload, ModuleChannel) << ModuleUnload.Base(uint64(Base));

	for (int i = 0; i < SubscriberCount; ++i)
	{
		Subscribers[i](false, (void*)Base, nullptr);
	}
}

////////////////////////////////////////////////////////////////////////////////
void
Modules_Create(FMalloc* Malloc)
{
	if (FModuleTrace::Get() != nullptr)
	{
		return;
	}

	static FModuleTrace Instance(Malloc);
}

////////////////////////////////////////////////////////////////////////////////
void
Modules_Initialize()
{
	if (FModuleTrace* Instance = FModuleTrace::Get())
	{
		Instance->Initialize();
	}
}

////////////////////////////////////////////////////////////////////////////////
void
Modules_Subscribe(void (*Function)(bool, void*, const char16_t*))
{
	if (FModuleTrace* Instance = FModuleTrace::Get())
	{
		Instance->Subscribe(Function);
	}
}

}  // namespace zen

#endif	// PLATFORM_SUPPORTS_WIN32_MEMORY_TRACE