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

#include "windowsrunner.h"

#if ZEN_WITH_COMPUTE_SERVICES && ZEN_PLATFORM_WINDOWS

#	include <zencore/compactbinary.h>
#	include <zencore/compactbinarypackage.h>
#	include <zencore/except.h>
#	include <zencore/except_fmt.h>
#	include <zencore/filesystem.h>
#	include <zencore/fmtutils.h>
#	include <zencore/scopeguard.h>
#	include <zencore/trace.h>
#	include <zencore/system.h>
#	include <zencore/timer.h>

ZEN_THIRD_PARTY_INCLUDES_START
#	include <userenv.h>
#	include <aclapi.h>
#	include <sddl.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen::compute {

using namespace std::literals;

WindowsProcessRunner::WindowsProcessRunner(ChunkResolver&				Resolver,
										   const std::filesystem::path& BaseDir,
										   DeferredDirectoryDeleter&	Deleter,
										   WorkerThreadPool&			WorkerPool,
										   bool							Sandboxed)
: LocalProcessRunner(Resolver, BaseDir, Deleter, WorkerPool)
, m_Sandboxed(Sandboxed)
{
	if (!m_Sandboxed)
	{
		return;
	}

	// Build a unique profile name per process to avoid collisions
	m_AppContainerName = L"zenserver-sandbox-" + std::to_wstring(GetCurrentProcessId());

	// Clean up any stale profile from a previous crash
	DeleteAppContainerProfile(m_AppContainerName.c_str());

	PSID Sid = nullptr;

	HRESULT Hr = CreateAppContainerProfile(m_AppContainerName.c_str(),
										   m_AppContainerName.c_str(),	// display name
										   m_AppContainerName.c_str(),	// description
										   nullptr,						// no capabilities
										   0,							// capability count
										   &Sid);

	if (FAILED(Hr))
	{
		throw zen::runtime_error("CreateAppContainerProfile failed: HRESULT 0x{:08X}", static_cast<uint32_t>(Hr));
	}

	m_AppContainerSid = Sid;

	ZEN_INFO("AppContainer sandboxing enabled for child processes (profile={})", WideToUtf8(m_AppContainerName));
}

WindowsProcessRunner::~WindowsProcessRunner()
{
	if (m_AppContainerSid)
	{
		FreeSid(m_AppContainerSid);
		m_AppContainerSid = nullptr;
	}

	if (!m_AppContainerName.empty())
	{
		DeleteAppContainerProfile(m_AppContainerName.c_str());
	}
}

void
WindowsProcessRunner::GrantAppContainerAccess(const std::filesystem::path& Path, DWORD AccessMask)
{
	PACL				 ExistingDacl		= nullptr;
	PSECURITY_DESCRIPTOR SecurityDescriptor = nullptr;

	DWORD Result = GetNamedSecurityInfoW(Path.c_str(),
										 SE_FILE_OBJECT,
										 DACL_SECURITY_INFORMATION,
										 nullptr,
										 nullptr,
										 &ExistingDacl,
										 nullptr,
										 &SecurityDescriptor);

	if (Result != ERROR_SUCCESS)
	{
		throw zen::runtime_error("GetNamedSecurityInfoW failed for '{}': {}", Path.string(), GetSystemErrorAsString(Result));
	}

	auto $0 = MakeGuard([&] { LocalFree(SecurityDescriptor); });

	EXPLICIT_ACCESSW Access{};
	Access.grfAccessPermissions = AccessMask;
	Access.grfAccessMode		= SET_ACCESS;
	Access.grfInheritance		= OBJECT_INHERIT_ACE | CONTAINER_INHERIT_ACE;
	Access.Trustee.TrusteeForm	= TRUSTEE_IS_SID;
	Access.Trustee.TrusteeType	= TRUSTEE_IS_WELL_KNOWN_GROUP;
	Access.Trustee.ptstrName	= static_cast<LPWSTR>(m_AppContainerSid);

	PACL NewDacl = nullptr;

	Result = SetEntriesInAclW(1, &Access, ExistingDacl, &NewDacl);
	if (Result != ERROR_SUCCESS)
	{
		throw zen::runtime_error("SetEntriesInAclW failed for '{}': {}", Path.string(), GetSystemErrorAsString(Result));
	}

	auto $1 = MakeGuard([&] { LocalFree(NewDacl); });

	Result = SetNamedSecurityInfoW(const_cast<LPWSTR>(Path.c_str()),
								   SE_FILE_OBJECT,
								   DACL_SECURITY_INFORMATION,
								   nullptr,
								   nullptr,
								   NewDacl,
								   nullptr);

	if (Result != ERROR_SUCCESS)
	{
		throw zen::runtime_error("SetNamedSecurityInfoW failed for '{}': {}", Path.string(), GetSystemErrorAsString(Result));
	}
}

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

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

	// Set up environment variables

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

	StringBuilder<1024> EnvironmentBlock;

	for (auto& It : WorkerDescription["environment"sv])
	{
		EnvironmentBlock.Append(It.AsString());
		EnvironmentBlock.Append('\0');
	}
	EnvironmentBlock.Append('\0');
	EnvironmentBlock.Append('\0');

	// Execute process - this spawns the child process immediately without waiting
	// for completion

	std::string_view	  ExecPath = WorkerDescription["path"sv].AsString();
	std::filesystem::path ExePath  = Prepared->WorkerPath / std::filesystem::path(ExecPath).make_preferred();

	ExtendableWideStringBuilder<512> CommandLine;
	CommandLine.Append(L'"');
	CommandLine.Append(ExePath.c_str());
	CommandLine.Append(L'"');
	CommandLine.Append(L" -Build=build.action");

	LPSECURITY_ATTRIBUTES lpProcessAttributes = nullptr;
	LPSECURITY_ATTRIBUTES lpThreadAttributes  = nullptr;
	BOOL				  bInheritHandles	  = FALSE;
	DWORD				  dwCreationFlags	  = 0;

	ZEN_DEBUG("Executing: {} (sandboxed={})", WideToUtf8(CommandLine.c_str()), m_Sandboxed);

	CommandLine.EnsureNulTerminated();

	PROCESS_INFORMATION ProcessInformation{};

	if (m_Sandboxed)
	{
		// Grant AppContainer access to sandbox and worker directories
		GrantAppContainerAccess(Prepared->SandboxPath, FILE_ALL_ACCESS);
		GrantAppContainerAccess(Prepared->WorkerPath, FILE_GENERIC_READ | FILE_GENERIC_EXECUTE);

		// Set up extended startup info with AppContainer security capabilities
		SECURITY_CAPABILITIES SecurityCapabilities{};
		SecurityCapabilities.AppContainerSid = m_AppContainerSid;
		SecurityCapabilities.Capabilities	 = nullptr;
		SecurityCapabilities.CapabilityCount = 0;

		SIZE_T AttrListSize = 0;
		InitializeProcThreadAttributeList(nullptr, 1, 0, &AttrListSize);

		auto AttrList = static_cast<PPROC_THREAD_ATTRIBUTE_LIST>(malloc(AttrListSize));
		auto $0		  = MakeGuard([&] { free(AttrList); });

		if (!InitializeProcThreadAttributeList(AttrList, 1, 0, &AttrListSize))
		{
			zen::ThrowLastError("InitializeProcThreadAttributeList failed");
		}

		auto $1 = MakeGuard([&] { DeleteProcThreadAttributeList(AttrList); });

		if (!UpdateProcThreadAttribute(AttrList,
									   0,
									   PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES,
									   &SecurityCapabilities,
									   sizeof(SecurityCapabilities),
									   nullptr,
									   nullptr))
		{
			zen::ThrowLastError("UpdateProcThreadAttribute (SECURITY_CAPABILITIES) failed");
		}

		STARTUPINFOEXW StartupInfoEx{};
		StartupInfoEx.StartupInfo.cb  = sizeof(STARTUPINFOEXW);
		StartupInfoEx.lpAttributeList = AttrList;

		dwCreationFlags |= EXTENDED_STARTUPINFO_PRESENT;

		BOOL Success = CreateProcessW(nullptr,
									  CommandLine.Data(),
									  lpProcessAttributes,
									  lpThreadAttributes,
									  bInheritHandles,
									  dwCreationFlags,
									  (LPVOID)EnvironmentBlock.Data(),
									  Prepared->SandboxPath.c_str(),
									  &StartupInfoEx.StartupInfo,
									  /* out */ &ProcessInformation);

		if (!Success)
		{
			zen::ThrowLastError("Unable to launch sandboxed process");
		}
	}
	else
	{
		STARTUPINFO StartupInfo{};
		StartupInfo.cb = sizeof StartupInfo;

		BOOL Success = CreateProcessW(nullptr,
									  CommandLine.Data(),
									  lpProcessAttributes,
									  lpThreadAttributes,
									  bInheritHandles,
									  dwCreationFlags,
									  (LPVOID)EnvironmentBlock.Data(),
									  Prepared->SandboxPath.c_str(),
									  &StartupInfo,
									  /* out */ &ProcessInformation);

		if (!Success)
		{
			zen::ThrowLastError("Unable to launch process");
		}
	}

	CloseHandle(ProcessInformation.hThread);

	Ref<RunningAction> NewAction{new RunningAction()};
	NewAction->Action		 = Action;
	NewAction->ProcessHandle = ProcessInformation.hProcess;
	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
WindowsProcessRunner::SweepRunningActions()
{
	ZEN_TRACE_CPU("WindowsProcessRunner::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;

			DWORD ExitCode	= 0;
			BOOL  IsSuccess = GetExitCodeProcess(Running->ProcessHandle, &ExitCode);

			if (IsSuccess && ExitCode != STILL_ACTIVE)
			{
				CloseHandle(Running->ProcessHandle);
				Running->ProcessHandle = INVALID_HANDLE_VALUE;
				Running->ExitCode	   = ExitCode;

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

	ProcessCompletedActions(CompletedActions);
}

void
WindowsProcessRunner::CancelRunningActions()
{
	ZEN_TRACE_CPU("WindowsProcessRunner::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");

	// For expedience we initiate the process termination for all known
	// processes before attempting to wait for them to exit.

	std::vector<int> TerminatedLsnList;

	for (const auto& Kv : RunningMap)
	{
		Ref<RunningAction> Action = Kv.second;

		BOOL TermSuccess = TerminateProcess(Action->ProcessHandle, 222);

		if (TermSuccess)
		{
			TerminatedLsnList.push_back(Kv.first);
		}
		else
		{
			DWORD LastError = GetLastError();

			if (LastError != ERROR_ACCESS_DENIED)
			{
				ZEN_WARN("TerminateProcess for LSN {} not successful: {}", Action->Action->ActionLsn, GetSystemErrorAsString(LastError));
			}
		}
	}

	// We only post results for processes we have terminated, in order
	// to avoid multiple results getting posted for the same action

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

			if (Running->ProcessHandle != INVALID_HANDLE_VALUE)
			{
				DWORD WaitResult = WaitForSingleObject(Running->ProcessHandle, 2000);

				if (WaitResult != WAIT_OBJECT_0)
				{
					ZEN_WARN("wait for LSN {}: process exit did not succeed, result = {}", Running->Action->ActionLsn, WaitResult);
				}
				else
				{
					ZEN_DEBUG("LSN {}: process exit OK", Running->Action->ActionLsn);
				}
			}

			// Clean up and post error result

			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