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

#include "winerunner.h"

#if ZEN_WITH_COMPUTE_SERVICES && ZEN_PLATFORM_LINUX

#	include <zencore/compactbinary.h>
#	include <zencore/compactbinarypackage.h>
#	include <zencore/except.h>
#	include <zencore/filesystem.h>
#	include <zencore/fmtutils.h>
#	include <zencore/iobuffer.h>
#	include <zencore/iohash.h>
#	include <zencore/timer.h>
#	include <zencore/trace.h>

#	include <signal.h>
#	include <sys/wait.h>
#	include <unistd.h>

namespace zen::compute {

using namespace std::literals;

WineProcessRunner::WineProcessRunner(ChunkResolver&				  Resolver,
									 const std::filesystem::path& BaseDir,
									 DeferredDirectoryDeleter&	  Deleter,
									 WorkerThreadPool&			  WorkerPool)
: LocalProcessRunner(Resolver, BaseDir, Deleter, WorkerPool)
{
	// Restore SIGCHLD to default behavior so waitpid() can properly collect
	// child exit status. zenserver/main.cpp sets SIGCHLD to SIG_IGN which
	// causes the kernel to auto-reap children, making waitpid() return
	// -1/ECHILD instead of the exit status we need.
	struct sigaction Action = {};
	sigemptyset(&Action.sa_mask);
	Action.sa_handler = SIG_DFL;
	sigaction(SIGCHLD, &Action, nullptr);
}

SubmitResult
WineProcessRunner::SubmitAction(Ref<RunnerAction> Action)
{
	ZEN_TRACE_CPU("WineProcessRunner::SubmitAction");
	std::optional<PreparedAction> Prepared = PrepareActionSubmission(Action);

	if (!Prepared)
	{
		return SubmitResult{.IsAccepted = false};
	}

	// Build environment array from worker descriptor

	CbObject WorkerDescription = Prepared->WorkerPackage.GetObject();

	std::vector<std::string> EnvStrings;
	for (auto& It : WorkerDescription["environment"sv])
	{
		EnvStrings.emplace_back(It.AsString());
	}

	std::vector<char*> Envp;
	Envp.reserve(EnvStrings.size() + 1);
	for (auto& Str : EnvStrings)
	{
		Envp.push_back(Str.data());
	}
	Envp.push_back(nullptr);

	// Build argv: wine <worker_exe_path> -Build=build.action

	std::string_view	  ExecPath	  = WorkerDescription["path"sv].AsString();
	std::filesystem::path ExePath	  = Prepared->WorkerPath / std::filesystem::path(ExecPath);
	std::string			  ExePathStr  = ExePath.string();
	std::string			  WinePathStr = m_WinePath;
	std::string			  BuildArg	  = "-Build=build.action";

	std::vector<char*> ArgV;
	ArgV.push_back(WinePathStr.data());
	ArgV.push_back(ExePathStr.data());
	ArgV.push_back(BuildArg.data());
	ArgV.push_back(nullptr);

	ZEN_DEBUG("Executing via Wine: {} {} {}", WinePathStr, ExePathStr, BuildArg);

	std::string SandboxPathStr = Prepared->SandboxPath.string();

	pid_t ChildPid = fork();

	if (ChildPid < 0)
	{
		throw std::runtime_error(fmt::format("fork() failed: {}", strerror(errno)));
	}

	if (ChildPid == 0)
	{
		// Child process
		if (chdir(SandboxPathStr.c_str()) != 0)
		{
			_exit(127);
		}

		execve(WinePathStr.c_str(), ArgV.data(), Envp.data());

		// execve only returns on failure
		_exit(127);
	}

	// Parent: store child pid as void* (same convention as zencore/process.cpp)

	Ref<RunningAction> NewAction{new RunningAction()};
	NewAction->Action		 = Action;
	NewAction->ProcessHandle = reinterpret_cast<void*>(static_cast<intptr_t>(ChildPid));
	NewAction->SandboxPath	 = std::move(Prepared->SandboxPath);

	{
		RwLock::ExclusiveLockScope _(m_RunningLock);
		m_RunningMap[Prepared->ActionLsn] = std::move(NewAction);
	}

	Action->SetActionState(RunnerAction::State::Running);

	return SubmitResult{.IsAccepted = true};
}

void
WineProcessRunner::SweepRunningActions()
{
	ZEN_TRACE_CPU("WineProcessRunner::SweepRunningActions");
	std::vector<Ref<RunningAction>> CompletedActions;

	m_RunningLock.WithExclusiveLock([&] {
		for (auto It = begin(m_RunningMap), ItEnd = end(m_RunningMap); It != ItEnd;)
		{
			Ref<RunningAction> Running = It->second;

			pid_t Pid	 = static_cast<pid_t>(reinterpret_cast<intptr_t>(Running->ProcessHandle));
			int	  Status = 0;

			pid_t Result = waitpid(Pid, &Status, WNOHANG);

			if (Result == Pid)
			{
				if (WIFEXITED(Status))
				{
					Running->ExitCode = WEXITSTATUS(Status);
				}
				else if (WIFSIGNALED(Status))
				{
					Running->ExitCode = 128 + WTERMSIG(Status);
				}
				else
				{
					Running->ExitCode = 1;
				}

				Running->ProcessHandle = nullptr;

				CompletedActions.push_back(std::move(Running));
				It = m_RunningMap.erase(It);
			}
			else
			{
				++It;
			}
		}
	});

	ProcessCompletedActions(CompletedActions);
}

void
WineProcessRunner::CancelRunningActions()
{
	ZEN_TRACE_CPU("WineProcessRunner::CancelRunningActions");
	Stopwatch									Timer;
	std::unordered_map<int, Ref<RunningAction>> RunningMap;

	m_RunningLock.WithExclusiveLock([&] { std::swap(RunningMap, m_RunningMap); });

	if (RunningMap.empty())
	{
		return;
	}

	ZEN_INFO("cancelling all running actions");

	// Send SIGTERM to all running processes first

	std::vector<int> TerminatedLsnList;

	for (const auto& Kv : RunningMap)
	{
		Ref<RunningAction> Running = Kv.second;
		pid_t			   Pid	   = static_cast<pid_t>(reinterpret_cast<intptr_t>(Running->ProcessHandle));

		if (kill(Pid, SIGTERM) == 0)
		{
			TerminatedLsnList.push_back(Kv.first);
		}
		else
		{
			ZEN_WARN("kill(SIGTERM) for LSN {} (pid {}) failed: {}", Running->Action->ActionLsn, Pid, strerror(errno));
		}
	}

	// Wait up to 2 seconds for graceful exit, then SIGKILL if needed

	for (int Lsn : TerminatedLsnList)
	{
		if (auto It = RunningMap.find(Lsn); It != RunningMap.end())
		{
			Ref<RunningAction> Running = It->second;
			pid_t			   Pid	   = static_cast<pid_t>(reinterpret_cast<intptr_t>(Running->ProcessHandle));

			// Poll for up to 2 seconds
			bool Exited = false;
			for (int i = 0; i < 20; ++i)
			{
				int	  Status	 = 0;
				pid_t WaitResult = waitpid(Pid, &Status, WNOHANG);
				if (WaitResult == Pid)
				{
					Exited = true;
					ZEN_DEBUG("LSN {}: process exit OK", Running->Action->ActionLsn);
					break;
				}
				usleep(100000);	 // 100ms
			}

			if (!Exited)
			{
				ZEN_WARN("LSN {}: process did not exit after SIGTERM, sending SIGKILL", Running->Action->ActionLsn);
				kill(Pid, SIGKILL);
				waitpid(Pid, nullptr, 0);
			}

			m_DeferredDeleter.Enqueue(Running->Action->ActionLsn, std::move(Running->SandboxPath));
			Running->Action->SetActionState(RunnerAction::State::Failed);
		}
	}

	ZEN_INFO("DONE - cancelled {} running processes (took {})", TerminatedLsnList.size(), NiceTimeSpanMs(Timer.GetElapsedTimeMs()));
}

}  // namespace zen::compute

#endif