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

#include <zencore/process.h>

#include <zencore/except.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/scopeguard.h>
#include <zencore/string.h>
#include <zencore/testing.h>
#include <zencore/timer.h>

#include <thread>

ZEN_THIRD_PARTY_INCLUDES_START

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

#	include <Psapi.h>
#	include <shellapi.h>
#	include <Shlobj.h>
#	include <TlHelp32.h>
#else
#	include <fcntl.h>
#	include <pthread.h>
#	include <signal.h>
#	include <sys/file.h>
#	include <sys/sem.h>
#	include <sys/stat.h>
#	include <sys/syscall.h>
#	include <sys/wait.h>
#	include <time.h>
#	include <unistd.h>
#endif

#if ZEN_PLATFORM_MAC
#	include <libproc.h>
#	include <sys/types.h>
#	include <sys/sysctl.h>
#endif

#include <fmt/format.h>

ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

#if ZEN_PLATFORM_LINUX
void
IgnoreChildSignals()
{
	// When a child process exits it is put into a zombie state until the parent
	// collects its result. This doesn't fit the Windows-like model that Zen uses
	// where there is a less strict familial model and no zombification. Ignoring
	// SIGCHLD signals removes the need to call wait() on zombies. Another option
	// would be for the child to call setsid() but that would detatch the child
	// from the terminal.
	struct sigaction Action = {};
	sigemptyset(&Action.sa_mask);
	Action.sa_handler = SIG_IGN;
	sigaction(SIGCHLD, &Action, nullptr);
}

static char
GetPidStatus(int Pid, std::error_code& OutEc)
{
	std::filesystem::path EntryPath = std::filesystem::path("/proc") / fmt::format("{}", Pid);
	std::filesystem::path StatPath	= EntryPath / "stat";
	if (IsFile(StatPath))
	{
		FILE* StatFile = fopen(StatPath.c_str(), "r");
		if (StatFile)
		{
			char Buffer[5120];
			int	 Size = fread(Buffer, 1, 5120 - 1, StatFile);
			fclose(StatFile);
			if (Size > 0)
			{
				Buffer[Size]  = 0;
				char* ScanPtr = strrchr(Buffer, ')');
				if (ScanPtr && ScanPtr[1] != '\0')
				{
					ScanPtr += 2;
					char State = *ScanPtr;
					return State;
				}
			}
		}
		else
		{
			OutEc = MakeErrorCodeFromLastError();
		}
	}
	return 0;
}

bool
IsZombieProcess(int pid, std::error_code& OutEc)
{
	char Status = GetPidStatus(pid, OutEc);
	if (OutEc)
	{
		return false;
	}
	if (Status == 'Z' || Status == 0)
	{
		return true;
	}
	return false;
}

#endif	// ZEN_PLATFORM_LINUX

#if ZEN_PLATFORM_MAC
bool
IsZombieProcess(int pid, std::error_code& OutEc)
{
	struct kinfo_proc Info;
	int				  Mib[4]   = {CTL_KERN, KERN_PROC, KERN_PROC_PID, pid};
	size_t			  InfoSize = sizeof Info;

	int Res = sysctl(Mib, 4, &Info, &InfoSize, NULL, 0);
	if (Res != 0)
	{
		OutEc = MakeErrorCodeFromLastError();
		return false;
	}
	if (Info.kp_proc.p_stat == SZOMB)
	{
		// Zombie process
		return true;
	}
	return false;
}
#endif	// ZEN_PLATFORM_MAC

ProcessHandle::ProcessHandle() = default;

#if ZEN_PLATFORM_WINDOWS
void
ProcessHandle::Initialize(void* ProcessHandle)
{
	ZEN_ASSERT(m_ProcessHandle == nullptr);

	if (ProcessHandle == INVALID_HANDLE_VALUE)
	{
		ProcessHandle = nullptr;
	}

	// TODO: perform some debug verification here to verify it's a valid handle?
	m_ProcessHandle = ProcessHandle;
	m_Pid			= GetProcessId(m_ProcessHandle);
}
#endif	// ZEN_PLATFORM_WINDOWS

ProcessHandle::~ProcessHandle()
{
	Reset();
}

void
ProcessHandle::Initialize(int Pid, std::error_code& OutEc)
{
	OutEc.clear();
	ZEN_ASSERT(m_ProcessHandle == nullptr);

#if ZEN_PLATFORM_WINDOWS
	m_ProcessHandle = OpenProcess(PROCESS_QUERY_INFORMATION | SYNCHRONIZE, FALSE, Pid);
#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	if (Pid > 0)
	{
		m_ProcessHandle = (void*)(intptr_t(Pid));
	}
#endif

	if (!m_ProcessHandle)
	{
		OutEc = MakeErrorCodeFromLastError();
	}

	m_Pid = Pid;
}

void
ProcessHandle::Initialize(int Pid)
{
	std::error_code Ec;
	Initialize(Pid, Ec);

	if (Ec)
	{
		throw std::system_error(Ec, fmt::format("ProcessHandle::Initialize(pid: {}) failed", Pid));
	}

	m_Pid = Pid;
}

bool
ProcessHandle::IsRunning() const
{
	bool bActive = false;

#if ZEN_PLATFORM_WINDOWS
	DWORD ExitCode = 0;
	GetExitCodeProcess(m_ProcessHandle, &ExitCode);
	bActive = (ExitCode == STILL_ACTIVE);
#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	std::error_code _;
	bActive = IsProcessRunning(m_Pid, _);
#endif

	return bActive;
}

bool
ProcessHandle::IsValid() const
{
	return (m_ProcessHandle != nullptr);
}

bool
ProcessHandle::Kill()
{
	if (!IsRunning())
	{
		return true;
	}

#if ZEN_PLATFORM_WINDOWS
	SetConsoleCtrlHandler(nullptr, TRUE);  // Prevent this process from terminating itself
	auto _ = MakeGuard([] { SetConsoleCtrlHandler(nullptr, FALSE); });

	// Try graceful shutdown first
	if (GenerateConsoleCtrlEvent(CTRL_BREAK_EVENT, m_Pid))
	{
		// Wait briefly for graceful shutdown
		if (WaitForSingleObject(m_ProcessHandle, 5000) == WAIT_OBJECT_0)
		{
			Reset();
			return true;
		}
	}

	// Fall back to forceful termination if graceful shutdown failed
	if (!TerminateProcess(m_ProcessHandle, 0))
	{
		return false;
	}
#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	int Res = kill(pid_t(m_Pid), SIGTERM);
	if (Res != 0)
	{
		int err = errno;
		if (err != ESRCH)
		{
			return false;
		}
	}
#endif

	Reset();

	return true;
}

bool
ProcessHandle::Terminate(int ExitCode)
{
	if (!IsRunning())
	{
		return true;
	}

#if ZEN_PLATFORM_WINDOWS
	BOOL bTerminated = TerminateProcess(m_ProcessHandle, ExitCode);
	if (!bTerminated)
	{
		return false;
	}
	DWORD WaitResult = WaitForSingleObject(m_ProcessHandle, INFINITE);
	bool  bSuccess	 = (WaitResult == WAIT_OBJECT_0) || (WaitResult == WAIT_ABANDONED_0);
	if (!bSuccess)
	{
		return false;
	}
#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	ZEN_UNUSED(ExitCode);
	int Res = kill(pid_t(m_Pid), SIGKILL);
	if (Res != 0)
	{
		int err = errno;
		if (err != ESRCH)
		{
			return false;
		}
	}
#endif
	Reset();
	return true;
}

void
ProcessHandle::Reset()
{
	if (IsValid())
	{
#if ZEN_PLATFORM_WINDOWS
		CloseHandle(m_ProcessHandle);
#endif
		m_ProcessHandle = nullptr;
		m_Pid			= 0;
	}
}

bool
ProcessHandle::Wait(int TimeoutMs, std::error_code& OutEc)
{
	using namespace std::literals;

#if ZEN_PLATFORM_WINDOWS
	const DWORD Timeout = (TimeoutMs < 0) ? INFINITE : TimeoutMs;

	const DWORD WaitResult = WaitForSingleObject(m_ProcessHandle, Timeout);

	switch (WaitResult)
	{
		case WAIT_OBJECT_0:
			return true;

		case WAIT_TIMEOUT:
			return false;

		case WAIT_ABANDONED_0:
			return true;

		case WAIT_FAILED:
			break;
	}

	OutEc = MakeErrorCodeFromLastError();

	return false;
#endif	// ZEN_PLATFORM_WINDOWS

#if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	const int SleepMs	= 20;
	timespec  SleepTime = {0, SleepMs * 1000 * 1000};
	for (int SleepedTimeMS = 0;; SleepedTimeMS += SleepMs)
	{
		int WaitState = 0;
		if (waitpid(m_Pid, &WaitState, WNOHANG | WCONTINUED | WUNTRACED) != -1)
		{
			if (WIFEXITED(WaitState))
			{
				m_ExitCode = WEXITSTATUS(WaitState);
			}
		}

		if (!IsProcessRunning(m_Pid, OutEc))
		{
			return true;
		}
		else if (OutEc)
		{
			return false;
		}

		if (kill(m_Pid, 0) < 0)
		{
			int32_t LastError = zen::GetLastError();
			if (LastError == ESRCH)
			{
				return true;
			}
			OutEc = MakeErrorCode(LastError);
			return false;
		}
		else if (IsZombieProcess(m_Pid, OutEc))
		{
			ZEN_INFO("Found process {} in zombie state, treating as not running", m_Pid);
			return true;
		}

		if (TimeoutMs >= 0 && SleepedTimeMS >= TimeoutMs)
		{
			return false;
		}

		nanosleep(&SleepTime, nullptr);
	}
	return false;
#endif	// ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
}

bool
ProcessHandle::Wait(int TimeoutMs)
{
	std::error_code Ec;
	if (Wait(TimeoutMs, Ec) && !Ec)
	{
		return true;
	}
	else if (Ec)
	{
		throw std::system_error(Ec, std::string("Process::Wait kill failed"));
	}
	return false;
}

int
ProcessHandle::GetExitCode()
{
#if ZEN_PLATFORM_WINDOWS
	DWORD ExitCode = 0;
	GetExitCodeProcess(m_ProcessHandle, &ExitCode);

	ZEN_ASSERT(ExitCode != STILL_ACTIVE);

	return ExitCode;
#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	return m_ExitCode;
#else
	ZEN_NOT_IMPLEMENTED();

	return 0;
#endif
}

int
ProcessHandle::WaitExitCode()
{
	Wait(-1);
	return GetExitCode();
}

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

#if !ZEN_PLATFORM_WINDOWS || ZEN_WITH_TESTS
static void
BuildArgV(std::vector<char*>& Out, char* CommandLine)
{
	char* Cursor = CommandLine;
	while (true)
	{
		// Skip leading whitespace
		for (; *Cursor == ' '; ++Cursor)
			;

		// Check for nullp terminator
		if (*Cursor == '\0')
		{
			break;
		}

		Out.push_back(Cursor);

		// Extract word
		int QuoteCount = 0;
		do
		{
			QuoteCount += (*Cursor == '\"');
			if (*Cursor == ' ' && !(QuoteCount & 1))
			{
				break;
			}
			++Cursor;
		} while (*Cursor != '\0');

		if (*Cursor == '\0')
		{
			break;
		}

		*Cursor = '\0';
		++Cursor;
	}
}
#endif	// !WINDOWS || TESTS

#if ZEN_PLATFORM_WINDOWS
static CreateProcResult
CreateProcNormal(const std::filesystem::path& Executable, std::string_view CommandLine, const CreateProcOptions& Options)
{
	PROCESS_INFORMATION ProcessInfo{};
	STARTUPINFO			StartupInfo{.cb = sizeof(STARTUPINFO)};

	bool				  InheritHandles	= false;
	void*				  Environment		= nullptr;
	LPSECURITY_ATTRIBUTES ProcessAttributes = nullptr;
	LPSECURITY_ATTRIBUTES ThreadAttributes	= nullptr;

	const bool AssignToJob = Options.AssignToJob && Options.AssignToJob->IsValid();

	DWORD CreationFlags = 0;
	if (Options.Flags & CreateProcOptions::Flag_NewConsole)
	{
		CreationFlags |= CREATE_NEW_CONSOLE;
	}
	if (Options.Flags & CreateProcOptions::Flag_NoConsole)
	{
		CreationFlags |= CREATE_NO_WINDOW;
	}
	if (Options.Flags & CreateProcOptions::Flag_Windows_NewProcessGroup)
	{
		CreationFlags |= CREATE_NEW_PROCESS_GROUP;
	}
	if (AssignToJob)
	{
		CreationFlags |= CREATE_SUSPENDED;
	}

	const wchar_t* WorkingDir = nullptr;
	if (Options.WorkingDirectory != nullptr)
	{
		WorkingDir = Options.WorkingDirectory->c_str();
	}

	ExtendableWideStringBuilder<256> CommandLineZ;
	CommandLineZ << CommandLine;

	if (!Options.StdoutFile.empty())
	{
		SECURITY_ATTRIBUTES sa;
		sa.nLength				= sizeof sa;
		sa.lpSecurityDescriptor = nullptr;
		sa.bInheritHandle		= TRUE;

		StartupInfo.hStdInput  = nullptr;
		StartupInfo.hStdOutput = CreateFileW(Options.StdoutFile.c_str(),
											 GENERIC_READ | GENERIC_WRITE,
											 FILE_SHARE_READ,
											 &sa,
											 CREATE_ALWAYS,
											 FILE_ATTRIBUTE_NORMAL,
											 nullptr);

		const BOOL Success = DuplicateHandle(GetCurrentProcess(),
											 StartupInfo.hStdOutput,
											 GetCurrentProcess(),
											 &StartupInfo.hStdError,
											 0,
											 TRUE,
											 DUPLICATE_SAME_ACCESS);

		if (Success)
		{
			StartupInfo.dwFlags |= STARTF_USESTDHANDLES;
			InheritHandles = true;
		}
		else
		{
			CloseHandle(StartupInfo.hStdOutput);
			StartupInfo.hStdOutput = 0;
		}
	}

	BOOL Success = CreateProcessW(Executable.c_str(),
								  CommandLineZ.Data(),
								  ProcessAttributes,
								  ThreadAttributes,
								  InheritHandles,
								  CreationFlags,
								  Environment,
								  WorkingDir,
								  &StartupInfo,
								  &ProcessInfo);

	if (StartupInfo.dwFlags & STARTF_USESTDHANDLES)
	{
		CloseHandle(StartupInfo.hStdError);
		CloseHandle(StartupInfo.hStdOutput);
	}

	if (!Success)
	{
		return nullptr;
	}

	if (AssignToJob)
	{
		if (!Options.AssignToJob->AssignProcess(ProcessInfo.hProcess))
		{
			ZEN_WARN("Failed to assign newly created process to job object");
		}
		ResumeThread(ProcessInfo.hThread);
	}

	CloseHandle(ProcessInfo.hThread);
	return ProcessInfo.hProcess;
}

static CreateProcResult
CreateProcUnelevated(const std::filesystem::path& Executable, std::string_view CommandLine, const CreateProcOptions& Options)
{
	/* Launches a binary with the shell as its parent. The shell (such as
	   Explorer) should be an unelevated process. */

	// No sense in using this route if we are not elevated in the first place
	if (IsUserAnAdmin() == FALSE)
	{
		return CreateProcNormal(Executable, CommandLine, Options);
	}

	// Get the users' shell process and open it for process creation
	HWND ShellWnd = GetShellWindow();
	if (ShellWnd == nullptr)
	{
		return nullptr;
	}

	DWORD ShellPid;
	GetWindowThreadProcessId(ShellWnd, &ShellPid);

	HANDLE Process = OpenProcess(PROCESS_CREATE_PROCESS, FALSE, ShellPid);
	if (Process == nullptr)
	{
		return nullptr;
	}
	auto $0 = MakeGuard([&] { CloseHandle(Process); });

	// Creating a process as a child of another process is done by setting a
	// thread-attribute list on the startup info passed to CreateProcess()
	SIZE_T AttrListSize;
	InitializeProcThreadAttributeList(nullptr, 1, 0, &AttrListSize);

	auto AttrList = (PPROC_THREAD_ATTRIBUTE_LIST)malloc(AttrListSize);
	auto $1		  = MakeGuard([&] { free(AttrList); });

	if (!InitializeProcThreadAttributeList(AttrList, 1, 0, &AttrListSize))
	{
		return nullptr;
	}

	BOOL bOk =
		UpdateProcThreadAttribute(AttrList, 0, PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, (HANDLE*)&Process, sizeof(Process), nullptr, nullptr);
	if (!bOk)
	{
		return nullptr;
	}

	// By this point we know we are an elevated process. It is not allowed to
	// create a process as a child of another unelevated process that share our
	// elevated console window if we have one. So we'll need to create a new one.
	uint32_t CreateProcFlags = EXTENDED_STARTUPINFO_PRESENT;
	if (GetConsoleWindow() != nullptr)
	{
		CreateProcFlags |= CREATE_NEW_CONSOLE;
	}
	else
	{
		CreateProcFlags |= DETACHED_PROCESS;
	}

	// Everything is set up now so we can proceed and launch the process
	STARTUPINFOEXW StartupInfo = {
		.StartupInfo	 = {.cb = sizeof(STARTUPINFOEXW)},
		.lpAttributeList = AttrList,
	};
	PROCESS_INFORMATION ProcessInfo = {};

	const bool AssignToJob = Options.AssignToJob && Options.AssignToJob->IsValid();

	if (Options.Flags & CreateProcOptions::Flag_NewConsole)
	{
		CreateProcFlags |= CREATE_NEW_CONSOLE;
	}
	if (Options.Flags & CreateProcOptions::Flag_NoConsole)
	{
		CreateProcFlags |= CREATE_NO_WINDOW;
	}
	if (AssignToJob)
	{
		CreateProcFlags |= CREATE_SUSPENDED;
	}

	ExtendableWideStringBuilder<256> CommandLineZ;
	CommandLineZ << CommandLine;

	ExtendableWideStringBuilder<256> CurrentDirZ;
	LPCWSTR							 WorkingDirectoryPtr = nullptr;
	if (Options.WorkingDirectory)
	{
		CurrentDirZ << Options.WorkingDirectory->native();
		WorkingDirectoryPtr = CurrentDirZ.c_str();
	}

	bOk = CreateProcessW(Executable.c_str(),
						 CommandLineZ.Data(),
						 nullptr,
						 nullptr,
						 FALSE,
						 CreateProcFlags,
						 nullptr,
						 WorkingDirectoryPtr,
						 &StartupInfo.StartupInfo,
						 &ProcessInfo);
	if (bOk == FALSE)
	{
		return nullptr;
	}

	if (AssignToJob)
	{
		if (!Options.AssignToJob->AssignProcess(ProcessInfo.hProcess))
		{
			ZEN_WARN("Failed to assign newly created process to job object");
		}
		ResumeThread(ProcessInfo.hThread);
	}

	CloseHandle(ProcessInfo.hThread);
	return ProcessInfo.hProcess;
}

static CreateProcResult
CreateProcElevated(const std::filesystem::path& Executable, std::string_view CommandLine, const CreateProcOptions& Options)
{
	ExtendableWideStringBuilder<256> CommandLineZ;
	CommandLineZ << CommandLine;

	SHELLEXECUTEINFO ShellExecuteInfo;
	ZeroMemory(&ShellExecuteInfo, sizeof(ShellExecuteInfo));
	ShellExecuteInfo.cbSize		  = sizeof(ShellExecuteInfo);
	ShellExecuteInfo.fMask		  = SEE_MASK_UNICODE | SEE_MASK_NOCLOSEPROCESS;
	ShellExecuteInfo.lpFile		  = Executable.c_str();
	ShellExecuteInfo.lpVerb		  = TEXT("runas");
	ShellExecuteInfo.nShow		  = SW_SHOW;
	ShellExecuteInfo.lpParameters = CommandLineZ.c_str();

	if (Options.WorkingDirectory != nullptr)
	{
		ShellExecuteInfo.lpDirectory = Options.WorkingDirectory->c_str();
	}

	if (::ShellExecuteEx(&ShellExecuteInfo))
	{
		return ShellExecuteInfo.hProcess;
	}

	return nullptr;
}
#endif	// ZEN_PLATFORM_WINDOWS

CreateProcResult
CreateProc(const std::filesystem::path& Executable, std::string_view CommandLine, const CreateProcOptions& Options)
{
#if ZEN_PLATFORM_WINDOWS
	if (Options.Flags & CreateProcOptions::Flag_Unelevated)
	{
		return CreateProcUnelevated(Executable, CommandLine, Options);
	}

	if (Options.Flags & CreateProcOptions::Flag_Elevated)
	{
		return CreateProcElevated(Executable, CommandLine, Options);
	}

	return CreateProcNormal(Executable, CommandLine, Options);
#else
	std::vector<char*> ArgV;
	std::string		   CommandLineZ(CommandLine);
	BuildArgV(ArgV, CommandLineZ.data());
	ArgV.push_back(nullptr);

	int ChildPid = fork();
	if (ChildPid < 0)
	{
		ThrowLastError("Failed to fork a new child process");
	}
	else if (ChildPid == 0)
	{
		if (Options.WorkingDirectory != nullptr)
		{
			int Result = chdir(Options.WorkingDirectory->c_str());
			ZEN_UNUSED(Result);
		}

		if (execv(Executable.c_str(), ArgV.data()) < 0)
		{
			ThrowLastError("Failed to exec() a new process image");
		}
	}

	return ChildPid;
#endif
}

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

ProcessMonitor::ProcessMonitor()
{
}

ProcessMonitor::~ProcessMonitor()
{
	RwLock::ExclusiveLockScope _(m_Lock);

	for (HandleType& Proc : m_ProcessHandles)
	{
#if ZEN_PLATFORM_WINDOWS
		CloseHandle(Proc);
#endif
		Proc = 0;
	}
}

bool
ProcessMonitor::IsRunning()
{
	RwLock::ExclusiveLockScope _(m_Lock);

	bool FoundOne = false;

	for (HandleType& Proc : m_ProcessHandles)
	{
		bool ProcIsActive;

#if ZEN_PLATFORM_WINDOWS
		DWORD ExitCode = 0;

		if (Proc)
		{
			GetExitCodeProcess(Proc, &ExitCode);
			ProcIsActive = (ExitCode == STILL_ACTIVE);
			if (!ProcIsActive)
			{
				CloseHandle(Proc);
			}
		}
		else
		{
			ProcIsActive = false;
		}
#else
		int Pid		 = int(intptr_t(Proc));
		ProcIsActive = IsProcessRunning(Pid);
#endif

		if (!ProcIsActive)
		{
			Proc = 0;
		}

		// Still alive
		FoundOne |= ProcIsActive;
	}

	std::erase_if(m_ProcessHandles, [](HandleType Handle) { return Handle == 0; });

	return FoundOne;
}

void
ProcessMonitor::AddPid(int Pid)
{
	HandleType ProcessHandle;

#if ZEN_PLATFORM_WINDOWS
	ProcessHandle = OpenProcess(PROCESS_QUERY_INFORMATION | SYNCHRONIZE, FALSE, Pid);
#else
	ProcessHandle = HandleType(intptr_t(Pid));
#endif

	RwLock::ExclusiveLockScope _(m_Lock);
	m_ProcessHandles.push_back(ProcessHandle);
}

bool
ProcessMonitor::IsActive() const
{
	RwLock::SharedLockScope _(m_Lock);
	return m_ProcessHandles.empty() == false;
}

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

#if ZEN_PLATFORM_WINDOWS
JobObject::JobObject() = default;

JobObject::~JobObject()
{
	if (m_JobHandle)
	{
		CloseHandle(m_JobHandle);
		m_JobHandle = nullptr;
	}
}

void
JobObject::Initialize()
{
	ZEN_ASSERT(m_JobHandle == nullptr, "JobObject already initialized");

	m_JobHandle = CreateJobObjectW(nullptr, nullptr);
	if (!m_JobHandle)
	{
		ZEN_WARN("Failed to create job object: {}", zen::GetLastError());
		return;
	}

	JOBOBJECT_EXTENDED_LIMIT_INFORMATION LimitInfo = {};
	LimitInfo.BasicLimitInformation.LimitFlags	   = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;

	if (!SetInformationJobObject(m_JobHandle, JobObjectExtendedLimitInformation, &LimitInfo, sizeof(LimitInfo)))
	{
		ZEN_WARN("Failed to set job object limits: {}", zen::GetLastError());
		CloseHandle(m_JobHandle);
		m_JobHandle = nullptr;
	}
}

bool
JobObject::AssignProcess(void* ProcessHandle)
{
	ZEN_ASSERT(m_JobHandle != nullptr, "JobObject not initialized");
	ZEN_ASSERT(ProcessHandle != nullptr, "ProcessHandle is null");

	if (!AssignProcessToJobObject(m_JobHandle, ProcessHandle))
	{
		ZEN_WARN("Failed to assign process to job object: {}", zen::GetLastError());
		return false;
	}

	return true;
}

bool
JobObject::IsValid() const
{
	return m_JobHandle != nullptr;
}
#endif	// ZEN_PLATFORM_WINDOWS

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

bool
IsProcessRunning(int pid, std::error_code& OutEc)
{
	// This function is arguably not super useful, a pid can be re-used
	// by the OS so holding on to a pid and polling it over some time
	// period will not necessarily tell you what you probably want to know.

#if ZEN_PLATFORM_WINDOWS
	HANDLE hProc = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid);

	if (!hProc)
	{
		DWORD Error = zen::GetLastError();
		if (Error == ERROR_INVALID_PARAMETER)
		{
			return false;
		}
		if (Error == ERROR_ACCESS_DENIED)
		{
			// Process is running under other user probably, assume it is running
			return true;
		}
		OutEc = MakeErrorCode(Error);
		return false;
	}
	auto _ = MakeGuard([hProc]() { CloseHandle(hProc); });

	bool  bStillActive = true;
	DWORD ExitCode	   = 0;
	if (0 != GetExitCodeProcess(hProc, &ExitCode))
	{
		bStillActive = ExitCode == STILL_ACTIVE;
	}
	else
	{
		DWORD Error = GetLastError();
		OutEc		= MakeErrorCode(Error);
		return false;
	}
	return bStillActive;
#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	int Res		  = kill(pid_t(pid), 0);
	if (Res == 0)
	{
		if (IsZombieProcess(pid, OutEc))
		{
			ZEN_INFO("Found process {} in zombie state, treating as not running", pid);
			return false;
		}
		if (OutEc)
		{
			return false;
		}
		return true;
	}
	int Error = errno;
	if (Error == ESRCH)	 // No such process
	{
		return false;
	}
	else if (Error == ENOENT)
	{
		return false;
	}
	else if (Error == EPERM)
	{
		return true;  // Running under a user we don't have access to, assume it is live
	}
	else
	{
		OutEc = MakeErrorCode(Error);
		return false;
	}
#endif
}

bool
IsProcessRunning(int pid)
{
	std::error_code Ec;
	bool			IsRunning = IsProcessRunning(pid, Ec);
	if (Ec)
	{
		ThrowSystemError(Ec.value(), fmt::format("Failed determining if process with pid {} is running", pid));
	}
	return IsRunning;
}

int
GetCurrentProcessId()
{
#if ZEN_PLATFORM_WINDOWS
	return ::GetCurrentProcessId();
#else
	return int(getpid());
#endif
}

int
GetProcessId(CreateProcResult ProcId)
{
#if ZEN_PLATFORM_WINDOWS
	return static_cast<int>(::GetProcessId(ProcId));
#else
	return ProcId;
#endif
}

std::filesystem::path
GetProcessExecutablePath(int Pid, std::error_code& OutEc)
{
#if ZEN_PLATFORM_WINDOWS
	HANDLE ModuleSnapshotHandle = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, (DWORD)Pid);
	if (ModuleSnapshotHandle != INVALID_HANDLE_VALUE)
	{
		auto		  __ = MakeGuard([&]() { CloseHandle(ModuleSnapshotHandle); });
		MODULEENTRY32 ModuleEntry;
		ModuleEntry.dwSize = sizeof(MODULEENTRY32);
		if (Module32First(ModuleSnapshotHandle, (LPMODULEENTRY32)&ModuleEntry))
		{
			std::filesystem::path ProcessExecutablePath(ModuleEntry.szExePath);
			return ProcessExecutablePath;
		}
		OutEc = MakeErrorCodeFromLastError();
		return {};
	}
	OutEc = MakeErrorCodeFromLastError();
	return {};
#endif	// ZEN_PLATFORM_WINDOWS
#if ZEN_PLATFORM_MAC
	char Buffer[PROC_PIDPATHINFO_MAXSIZE];
	int	 Res = proc_pidpath(Pid, Buffer, sizeof(Buffer));
	if (Res > 0)
	{
		std::filesystem::path ProcessExecutablePath(Buffer);
		return ProcessExecutablePath;
	}
	OutEc = MakeErrorCodeFromLastError();
	return {};
#endif	// ZEN_PLATFORM_MAC
#if ZEN_PLATFORM_LINUX
	std::filesystem::path EntryPath	  = std::filesystem::path("/proc") / fmt::format("{}", Pid);
	std::filesystem::path ExeLinkPath = EntryPath / "exe";
	char				  Link[4096];
	ssize_t				  BytesRead = readlink(ExeLinkPath.c_str(), Link, sizeof(Link) - 1);
	if (BytesRead > 0)
	{
		Link[BytesRead] = '\0';
		std::filesystem::path ProcessExecutablePath(Link);
		return ProcessExecutablePath;
	}
	OutEc = MakeErrorCodeFromLastError();
	return {};
#endif	// ZEN_PLATFORM_LINUX
}

std::string
GetProcessCommandLine(int Pid, std::error_code& OutEc)
{
#if ZEN_PLATFORM_WINDOWS
	HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, static_cast<DWORD>(Pid));
	if (!hProcess)
	{
		OutEc = MakeErrorCodeFromLastError();
		return {};
	}
	auto _ = MakeGuard([hProcess] { CloseHandle(hProcess); });

	// NtQueryInformationProcess is an undocumented NT API; load it dynamically.
	// Info class 60 = ProcessCommandLine, available since Windows 8.1.
	using PFN_NtQIP = LONG(WINAPI*)(HANDLE, UINT, PVOID, ULONG, PULONG);
	static const PFN_NtQIP s_NtQIP =
		reinterpret_cast<PFN_NtQIP>(GetProcAddress(GetModuleHandleW(L"ntdll.dll"), "NtQueryInformationProcess"));
	if (!s_NtQIP)
	{
		return {};
	}

	constexpr UINT ProcessCommandLineClass	= 60;
	constexpr LONG StatusInfoLengthMismatch = static_cast<LONG>(0xC0000004L);

	ULONG ReturnLength = 0;
	LONG  Status	   = s_NtQIP(hProcess, ProcessCommandLineClass, nullptr, 0, &ReturnLength);
	if (Status != StatusInfoLengthMismatch || ReturnLength == 0)
	{
		return {};
	}

	std::vector<char> Buf(ReturnLength);
	Status = s_NtQIP(hProcess, ProcessCommandLineClass, Buf.data(), ReturnLength, &ReturnLength);
	if (Status < 0)
	{
		OutEc = MakeErrorCodeFromLastError();
		return {};
	}

	// Output: UNICODE_STRING header immediately followed by the UTF-16 string data.
	// The UNICODE_STRING.Buffer field points into our Buf.
	struct LocalUnicodeString
	{
		USHORT Length;
		USHORT MaximumLength;
		WCHAR* Buffer;
	};
	if (ReturnLength < sizeof(LocalUnicodeString))
	{
		return {};
	}
	const auto* Us = reinterpret_cast<const LocalUnicodeString*>(Buf.data());
	if (Us->Length == 0 || Us->Buffer == nullptr)
	{
		return {};
	}

	// Skip argv[0]: may be a quoted path ("C:\...\exe.exe") or a bare path
	const WCHAR* p	 = Us->Buffer;
	const WCHAR* End = Us->Buffer + Us->Length / sizeof(WCHAR);
	if (p < End && *p == L'"')
	{
		++p;
		while (p < End && *p != L'"')
		{
			++p;
		}
		if (p < End)
		{
			++p;  // skip closing quote
		}
	}
	else
	{
		while (p < End && *p != L' ')
		{
			++p;
		}
	}
	while (p < End && *p == L' ')
	{
		++p;
	}
	if (p >= End)
	{
		return {};
	}

	int Utf8Size = WideCharToMultiByte(CP_UTF8, 0, p, static_cast<int>(End - p), nullptr, 0, nullptr, nullptr);
	if (Utf8Size <= 0)
	{
		OutEc = MakeErrorCodeFromLastError();
		return {};
	}
	std::string Result(Utf8Size, '\0');
	WideCharToMultiByte(CP_UTF8, 0, p, static_cast<int>(End - p), Result.data(), Utf8Size, nullptr, nullptr);
	return Result;

#elif ZEN_PLATFORM_LINUX
	std::string CmdlinePath = fmt::format("/proc/{}/cmdline", Pid);
	FILE*		F			= fopen(CmdlinePath.c_str(), "rb");
	if (!F)
	{
		OutEc = MakeErrorCodeFromLastError();
		return {};
	}
	auto FGuard = MakeGuard([F] { fclose(F); });

	// /proc/{pid}/cmdline contains null-separated argv entries; read it all
	std::string Raw;
	char		Chunk[4096];
	size_t		BytesRead;
	while ((BytesRead = fread(Chunk, 1, sizeof(Chunk), F)) > 0)
	{
		Raw.append(Chunk, BytesRead);
	}
	if (Raw.empty())
	{
		return {};
	}

	// Skip argv[0] (first null-terminated entry)
	const char* p	= Raw.data();
	const char* End = Raw.data() + Raw.size();
	while (p < End && *p != '\0')
	{
		++p;
	}
	if (p < End)
	{
		++p;  // skip null terminator of argv[0]
	}

	// Build result: remaining entries joined by spaces (inter-arg nulls → spaces)
	std::string Result;
	Result.reserve(static_cast<size_t>(End - p));
	for (const char* q = p; q < End; ++q)
	{
		Result += (*q == '\0') ? ' ' : *q;
	}
	while (!Result.empty() && Result.back() == ' ')
	{
		Result.pop_back();
	}
	return Result;

#elif ZEN_PLATFORM_MAC
	int	   Mib[3]  = {CTL_KERN, KERN_PROCARGS2, Pid};
	size_t BufSize = 0;
	if (sysctl(Mib, 3, nullptr, &BufSize, nullptr, 0) != 0 || BufSize == 0)
	{
		OutEc = MakeErrorCodeFromLastError();
		return {};
	}

	std::vector<char> Buf(BufSize);
	if (sysctl(Mib, 3, Buf.data(), &BufSize, nullptr, 0) != 0)
	{
		OutEc = MakeErrorCodeFromLastError();
		return {};
	}

	// Layout: [int argc][exec_path\0][null padding][argv[0]\0][argv[1]\0]...[envp\0]...
	if (BufSize < sizeof(int))
	{
		return {};
	}
	int Argc = 0;
	memcpy(&Argc, Buf.data(), sizeof(int));
	if (Argc <= 1)
	{
		return {};
	}

	const char* p	= Buf.data() + sizeof(int);
	const char* End = Buf.data() + BufSize;

	// Skip exec_path and any trailing null padding that follows it
	while (p < End && *p != '\0')
	{
		++p;
	}
	while (p < End && *p == '\0')
	{
		++p;
	}

	// Skip argv[0]
	while (p < End && *p != '\0')
	{
		++p;
	}
	if (p < End)
	{
		++p;
	}

	// Collect argv[1..Argc-1]
	std::string Result;
	for (int i = 1; i < Argc && p < End; ++i)
	{
		if (i > 1)
		{
			Result += ' ';
		}
		const char* ArgStart = p;
		while (p < End && *p != '\0')
		{
			++p;
		}
		Result.append(ArgStart, p);
		if (p < End)
		{
			++p;
		}
	}
	return Result;

#else
	ZEN_UNUSED(Pid);
	ZEN_UNUSED(OutEc);
	return {};
#endif
}

std::error_code
FindProcess(const std::filesystem::path& ExecutableImage, ProcessHandle& OutHandle, bool IncludeSelf)
{
#if ZEN_PLATFORM_WINDOWS
	HANDLE ProcessSnapshotHandle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (ProcessSnapshotHandle == INVALID_HANDLE_VALUE)
	{
		return MakeErrorCodeFromLastError();
	}
	auto _ = MakeGuard([&]() { CloseHandle(ProcessSnapshotHandle); });

	const DWORD ThisProcessId = ::GetCurrentProcessId();

	PROCESSENTRY32 Entry;
	Entry.dwSize = sizeof(PROCESSENTRY32);
	if (Process32First(ProcessSnapshotHandle, (LPPROCESSENTRY32)&Entry))
	{
		do
		{
			if ((IncludeSelf || (Entry.th32ProcessID != ThisProcessId)) && (ExecutableImage.filename() == Entry.szExeFile))
			{
				std::error_code		  Ec;
				std::filesystem::path EntryPath = GetProcessExecutablePath(Entry.th32ProcessID, Ec);
				if (!Ec)
				{
					if (EntryPath == ExecutableImage)
					{
						HANDLE Handle =
							OpenProcess(PROCESS_TERMINATE | SYNCHRONIZE | PROCESS_QUERY_INFORMATION, FALSE, Entry.th32ProcessID);
						if (Handle == NULL)
						{
							return MakeErrorCodeFromLastError();
						}
						DWORD ExitCode = 0;
						GetExitCodeProcess(Handle, &ExitCode);
						if (ExitCode == STILL_ACTIVE)
						{
							OutHandle.Initialize((void*)Handle);
							return {};
						}
					}
				}
			}
		} while (::Process32Next(ProcessSnapshotHandle, (LPPROCESSENTRY32)&Entry));
		return {};
	}
	return MakeErrorCodeFromLastError();
#endif	// ZEN_PLATFORM_WINDOWS
#if ZEN_PLATFORM_MAC
	int	   Mib[4]	  = {CTL_KERN, KERN_PROC, KERN_PROC_ALL, 0};
	size_t BufferSize = 0;

	struct kinfo_proc* Processes = nullptr;
	uint32_t		   ProcCount = 0;

	const pid_t ThisProcessId = getpid();

	if (sysctl(Mib, 4, NULL, &BufferSize, NULL, 0) != -1 && BufferSize > 0)
	{
		struct kinfo_proc* Processes = (struct kinfo_proc*)malloc(BufferSize);
		auto			   _		 = MakeGuard([&]() { free(Processes); });
		if (sysctl(Mib, 4, Processes, &BufferSize, NULL, 0) != -1)
		{
			ProcCount = (uint32_t)(BufferSize / sizeof(struct kinfo_proc));
			char Buffer[PROC_PIDPATHINFO_MAXSIZE];
			for (uint32_t ProcIndex = 0; ProcIndex < ProcCount; ProcIndex++)
			{
				pid_t Pid = Processes[ProcIndex].kp_proc.p_pid;
				if (IncludeSelf || (Pid != ThisProcessId))
				{
					std::error_code		  Ec;
					std::filesystem::path EntryPath = GetProcessExecutablePath(Pid, Ec);
					if (!Ec)
					{
						if (EntryPath == ExecutableImage)
						{
							if (Processes[ProcIndex].kp_proc.p_stat != SZOMB)
							{
								OutHandle.Initialize(Pid, Ec);
								return Ec;
							}
						}
					}
					Ec.clear();
				}
			}
			return {};
		}
	}
	return MakeErrorCodeFromLastError();
#endif	// ZEN_PLATFORM_MAC
#if ZEN_PLATFORM_LINUX
	const pid_t ThisProcessId = getpid();

	std::vector<uint32_t> RunningPids;
	DirectoryContent	  ProcList;
	GetDirectoryContent("/proc", DirectoryContentFlags::IncludeDirs, ProcList);
	for (const std::filesystem::path& EntryPath : ProcList.Directories)
	{
		std::string				EntryName = EntryPath.stem();
		std::optional<uint32_t> PidMaybe  = ParseInt<uint32_t>(EntryName);
		if (PidMaybe.has_value())
		{
			if (pid_t Pid = PidMaybe.value(); IncludeSelf || (Pid != ThisProcessId))
			{
				RunningPids.push_back(Pid);
			}
		}
	}

	for (uint32_t Pid : RunningPids)
	{
		std::error_code		  Ec;
		std::filesystem::path EntryPath = GetProcessExecutablePath((int)Pid, Ec);
		if (!Ec)
		{
			if (EntryPath == ExecutableImage)
			{
				char Status = GetPidStatus(Pid, Ec);
				if (!Ec)
				{
					if (Status && (Status != 'Z'))
					{
						OutHandle.Initialize(Pid, Ec);
						return Ec;
					}
				}
			}
		}
		Ec.clear();
	}
	return {};
#endif	// ZEN_PLATFORM_LINUX
}

void
WaitForThreads(uint64_t WaitTimeMs)
{
#if ZEN_PLATFORM_WINDOWS
	auto ThreadSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	if (ThreadSnapshot == INVALID_HANDLE_VALUE)
	{
		return;
	}
	auto _ = MakeGuard([ThreadSnapshot] { CloseHandle(ThreadSnapshot); });

	const DWORD CurrentProcessId = ::GetCurrentProcessId();
	const DWORD CurrentThreadId	 = ::GetCurrentThreadId();

	Stopwatch Timer;

	do
	{
		THREADENTRY32 ThreadEntry;
		ThreadEntry.dwSize = sizeof(THREADENTRY32);

		uint32_t ThreadCount = 0;
		if (Thread32First(ThreadSnapshot, &ThreadEntry))
		{
			do
			{
				if (ThreadEntry.th32OwnerProcessID == CurrentProcessId && ThreadEntry.th32ThreadID != CurrentThreadId)
				{
					ThreadCount++;
				}
			} while (Thread32Next(ThreadSnapshot, &ThreadEntry));
		}

		if (ThreadCount <= 1)
		{
			break;
		}

		const uint64_t SleepMs = 10;
		Sleep(SleepMs);
	} while (Timer.GetElapsedTimeMs() < WaitTimeMs);
#else
	ZEN_UNUSED(WaitTimeMs);
#endif
}

void
GetProcessMetrics(ProcessHandle& Handle, ProcessMetrics& OutMetrics)
{
#if ZEN_PLATFORM_WINDOWS
	FILETIME CreationTime;
	FILETIME ExitTime;
	FILETIME KernelTime;
	FILETIME UserTime;

	if (GetProcessTimes(Handle.Handle(), &CreationTime, &ExitTime, &KernelTime, &UserTime))
	{
		ULARGE_INTEGER KTime;
		KTime.LowPart  = KernelTime.dwLowDateTime;
		KTime.HighPart = KernelTime.dwHighDateTime;

		ULARGE_INTEGER UTime;
		UTime.LowPart  = UserTime.dwLowDateTime;
		UTime.HighPart = UserTime.dwHighDateTime;

		OutMetrics.KernelTimeMs = KTime.QuadPart / 10000;
		OutMetrics.UserTimeMs	= UTime.QuadPart / 10000;
	}

	PROCESS_MEMORY_COUNTERS MemCounters;
	if (GetProcessMemoryInfo(Handle.Handle(), &MemCounters, sizeof(MemCounters)))
	{
		OutMetrics.WorkingSetSize	  = MemCounters.WorkingSetSize;
		OutMetrics.PeakWorkingSetSize = MemCounters.PeakWorkingSetSize;
		OutMetrics.PagefileUsage	  = MemCounters.PagefileUsage;
		OutMetrics.PeakPagefileUsage  = MemCounters.PeakPagefileUsage;
	}
#else
	// TODO: implement for Linux and Mac
	ZEN_UNUSED(Handle);
	ZEN_UNUSED(OutMetrics);
#endif
}

#if ZEN_WITH_TESTS

void
process_forcelink()
{
}

TEST_SUITE_BEGIN("core.process");

TEST_CASE("Process")
{
	int Pid = GetCurrentProcessId();
	CHECK(Pid > 0);
	CHECK(IsProcessRunning(Pid));
}

TEST_CASE("FindProcess")
{
	{
		ProcessHandle	Process;
		std::error_code Ec = FindProcess(GetRunningExecutablePath(), Process, /*IncludeSelf*/ true);
		CHECK(!Ec);
		CHECK(Process.IsValid());
	}
	{
		ProcessHandle	Process;
		std::error_code Ec = FindProcess(GetRunningExecutablePath(), Process, /*IncludeSelf*/ false);
		CHECK(!Ec);
		CHECK(!Process.IsValid());
	}
	{
		ProcessHandle	Process;
		std::error_code Ec = FindProcess("this/does\\not/exist\\123914921929412312312312asdad\\12134.no", Process, /*IncludeSelf*/ false);
		CHECK(!Ec);
		CHECK(!Process.IsValid());
	}
}

TEST_CASE("BuildArgV")
{
	const char* Words[] = {"one", "two", "three", "four", "five"};
	struct
	{
		int			WordCount;
		const char* Input;
	} Cases[] = {
		{0, ""},
		{0, "  "},
		{1, "one"},
		{1, " one"},
		{1, "one  "},
		{2, "one two"},
		{2, " one two"},
		{2, "one two "},
		{2, " one two"},
		{2, "one two "},
		{2, "one two "},
		{3, "one two three"},
		{3, "\"one\" two \"three\""},
		{5, "one two three four five"},
	};

	for (const auto& Case : Cases)
	{
		std::vector<char*> OutArgs;
		StringBuilder<64>  Mutable;
		Mutable << Case.Input;
		BuildArgV(OutArgs, Mutable.Data());

		CHECK_EQ(OutArgs.size(), Case.WordCount);

		for (int i = 0, n = int(OutArgs.size()); i < n; ++i)
		{
			const char* Truth	 = Words[i];
			size_t		TruthLen = strlen(Truth);

			const char* Candidate = OutArgs[i];
			bool		bQuoted	  = (Candidate[0] == '\"');
			Candidate += bQuoted;

			CHECK(strncmp(Truth, Candidate, TruthLen) == 0);

			if (bQuoted)
			{
				CHECK_EQ(Candidate[TruthLen], '\"');
			}
		}
	}
}

TEST_SUITE_END(/* core.process */);

#endif

}  // namespace zen