// 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 (const 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
	ZEN_ASSERT(ProcessPid != 0);
	{
		RwLock::SharedLockScope Lock(m_SessionsLock);

		if (auto It = m_Sessions.find(SessionId); It != m_Sessions.end())
		{
			if (ProcessPid == It->second.ProcessPid)
			{
				void* ProcessHandle = It->second.ProcessHandle;
				if (ProcessHandle == nullptr)
				{
					// Session is stale, just ignore it
					return nullptr;
				}
				DWORD ExitCode = 0;
				GetExitCodeProcess(ProcessHandle, &ExitCode);
				if (ExitCode == STILL_ACTIVE)
				{
					return ProcessHandle;
				}
			}
		}
	}
	void* ProcessHandle = nullptr;
	void* DeadHandle	= nullptr;
	{
		// Add/replace the session

		RwLock::ExclusiveLockScope Lock(m_SessionsLock);
		if (auto It = m_Sessions.find(SessionId); It != m_Sessions.end())
		{
			// Check to see if someone beat us to it...
			if ((It->second.ProcessHandle != nullptr) && (ProcessPid == It->second.ProcessPid))
			{
				DWORD ExitCode = 0;
				GetExitCodeProcess(It->second.ProcessHandle, &ExitCode);
				if (ExitCode == STILL_ACTIVE)
				{
					return It->second.ProcessHandle;
				}
			}

			ZEN_INFO("Purging stale target process pid {} for session: {}", It->second.ProcessPid, SessionId, It->second.ProcessHandle);
			DeadHandle = It->second.ProcessHandle;
			m_Sessions.erase(It);
		}

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

	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;
	{
		std::vector<Oid>		   DeadSessions;
		RwLock::ExclusiveLockScope Lock(m_SessionsLock);
		for (auto& It : m_Sessions)
		{
			if (It.second.ProcessHandle == nullptr)
			{
				// Stale session, remove it on second pass of GC
				DeadSessions.push_back(It.first);
				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);

			// We just mark the session as "dead" for one GC pass to avoid re-opening a stale process
			It.second.ProcessHandle = nullptr;
		}
		for (auto SessionId : DeadSessions)
		{
			ZEN_INFO("Purging stale target process cache for session: {}", SessionId);
			m_Sessions.erase(SessionId);
		}
	}

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

void
OpenProcessCache::GcWorker()
{
	SetCurrentThreadName("ProcessCache_GC");

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

}  // namespace zen