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

#include "zenutil/openprocesscache.h"
#include <zencore/fmtutils.h>
#include <zencore/logging.h>

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>
#else
#	include <sys/mman.h>
#endif

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

namespace zen {

OpenProcessCache::OpenProcessCache()
#if ZEN_PLATFORM_WINDOWS
: m_GcThread(&OpenProcessCache::GcWorker, this)
#endif	// ZEN_PLATFORM_WINDOWS
{
}

OpenProcessCache::~OpenProcessCache()
{
#if ZEN_PLATFORM_WINDOWS
	try
	{
		m_GcExitEvent.Set();
		if (m_GcThread.joinable())
		{
			m_GcThread.join();
		}
		RwLock::ExclusiveLockScope Lock(m_SessionsLock);
		for (const auto& It : m_Sessions)
		{
			if (It.second.ProcessHandle == nullptr)
			{
				continue;
			}
			CloseHandle(It.second.ProcessHandle);
		}
		m_Sessions.clear();
	}
	catch (std::exception& Ex)
	{
		ZEN_ERROR("OpenProcessCache destructor failed with reason: `{}`", Ex.what());
	}
#endif	// ZEN_PLATFORM_WINDOWS
}

void*
OpenProcessCache::GetProcessHandle(Oid SessionId, int ProcessPid)
{
	if (ProcessPid == 0)
	{
		return nullptr;
	}
#if ZEN_PLATFORM_WINDOWS
	RwLock::ExclusiveLockScope Lock(m_SessionsLock);

	void* DeadHandle = nullptr;
	if (auto It = m_Sessions.find(SessionId); It != m_Sessions.end())
	{
		if (ProcessPid == It->second.ProcessPid)
		{
			void* ProcessHandle = It->second.ProcessHandle;
			ZEN_ASSERT(ProcessHandle != nullptr);
			DWORD ExitCode = 0;
			GetExitCodeProcess(It->second.ProcessHandle, &ExitCode);
			if (ExitCode == STILL_ACTIVE)
			{
				return It->second.ProcessHandle;
			}
		}

		// Remove the existing stale handle
		ZEN_INFO("Closing stale target process pid {} for session: {}", It->second.ProcessPid, SessionId, It->second.ProcessHandle);
		DeadHandle = It->second.ProcessHandle;
		m_Sessions.erase(It);
	}

	// Cache new handle
	void* ProcessHandle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_DUP_HANDLE, FALSE, ProcessPid);
	ZEN_INFO("Opening target process pid {} for session {}: {}", ProcessPid, SessionId, ProcessHandle == nullptr ? "failed" : "success");
	m_Sessions.insert_or_assign(SessionId, Process{.ProcessPid = ProcessPid, .ProcessHandle = ProcessHandle});
	Lock.ReleaseNow();

	if (DeadHandle != nullptr)
	{
		CloseHandle(DeadHandle);
	}

	return ProcessHandle;
#else	// ZEN_PLATFORM_WINDOWS
	ZEN_UNUSED(SessionId);
	return nullptr;
#endif	// ZEN_PLATFORM_WINDOWS
}

#if ZEN_PLATFORM_WINDOWS
void
OpenProcessCache::GCHandles()
{
	std::vector<void*>		   DeadHandles;
	RwLock::ExclusiveLockScope Lock(m_SessionsLock);
	for (auto& It : m_Sessions)
	{
		if (It.second.ProcessHandle == nullptr)
		{
			continue;
		}
		ZEN_ASSERT(It.second.ProcessPid != 0);
		DWORD ExitCode = 0;
		GetExitCodeProcess(It.second.ProcessHandle, &ExitCode);
		if (ExitCode == STILL_ACTIVE)
		{
			continue;
		}
		ZEN_INFO("GC stale target process pid {} for session {}: {}", It.second.ProcessPid, It.first, It.second.ProcessHandle);
		DeadHandles.push_back(It.second.ProcessHandle);
		It.second.ProcessPid	= 0;
		It.second.ProcessHandle = nullptr;
	}
	Lock.ReleaseNow();

	for (auto Handle : DeadHandles)
	{
		CloseHandle(Handle);
	}
}

void
OpenProcessCache::GcWorker()
{
	while (!m_GcExitEvent.Wait(500))
	{
		try
		{
			GCHandles();
		}
		catch (std::exception& Ex)
		{
			ZEN_ERROR("gc of open process cache failed with reason: `{}`", Ex.what());
		}
	}
}
#endif	// ZEN_PLATFORM_WINDOWS

}  // namespace zen