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

#include "zenserverprocess.h"

#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/string.h>

#include <shellapi.h>
#include <spdlog/spdlog.h>

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

std::atomic<int> TestCounter{0};

ZenServerEnvironment::ZenServerEnvironment()
{
}

ZenServerEnvironment::~ZenServerEnvironment()
{
}

void
ZenServerEnvironment::Initialize(std::filesystem::path ProgramBaseDir)
{
	m_ProgramBaseDir = ProgramBaseDir;

	spdlog::debug("Program base dir is '{}'", ProgramBaseDir);

	m_IsInitialized = true;
}

void
ZenServerEnvironment::InitializeForTest(std::filesystem::path ProgramBaseDir, std::filesystem::path TestBaseDir)
{
	m_ProgramBaseDir = ProgramBaseDir;
	m_TestBaseDir	 = TestBaseDir;

	spdlog::info("Program base dir is '{}'", ProgramBaseDir);
	spdlog::info("Cleaning test base dir '{}'", TestBaseDir);
	zen::DeleteDirectories(TestBaseDir.c_str());

	m_IsTestInstance = true;
	m_IsInitialized	 = true;
}

std::filesystem::path
ZenServerEnvironment::CreateNewTestDir()
{
	using namespace std::literals;

	zen::ExtendableWideStringBuilder<256> TestDir;
	TestDir << "test"sv << int64_t(++TestCounter);

	std::filesystem::path TestPath = m_TestBaseDir / TestDir.c_str();

	spdlog::info("Creating new test dir @ '{}'", TestPath);

	zen::CreateDirectories(TestPath.c_str());

	return TestPath;
}

std::filesystem::path
ZenServerEnvironment::GetTestRootDir(std::string_view Path)
{
	std::filesystem::path Root = m_ProgramBaseDir.parent_path().parent_path();

	std::filesystem::path Relative{Path};

	return Root / Relative;
}

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

std::atomic<int> ChildIdCounter{0};

ZenServerInstance::ZenServerInstance(ZenServerEnvironment& TestEnvironment) : m_Env(TestEnvironment)
{
	ZEN_ASSERT(TestEnvironment.IsInitialized());
}

ZenServerInstance::~ZenServerInstance()
{
	if (m_Process.IsValid())
	{
		if (m_Terminate)
		{
			spdlog::info("Terminating zenserver process");
			m_Process.Terminate(111);
		}
		else
		{
			SignalShutdown();
			m_Process.Wait();
		}
	}
}

void
ZenServerInstance::SpawnServer(int BasePort)
{
	ZEN_ASSERT(!m_Process.IsValid());  // Only spawn once

	const std::filesystem::path BaseDir	   = m_Env.ProgramBaseDir();
	const std::filesystem::path Executable = BaseDir / "zenserver.exe";

	const int MyPid	  = _getpid();
	const int ChildId = ++ChildIdCounter;

	zen::ExtendableStringBuilder<32> ChildEventName;
	ChildEventName << "Zen_Child_" << ChildId;
	zen::NamedEvent ChildEvent{ChildEventName};

	zen::ExtendableStringBuilder<32> ChildShutdownEventName;
	ChildShutdownEventName << "Zen_Child_" << ChildId;
	ChildShutdownEventName << "_Shutdown";
	zen::NamedEvent ChildShutdownEvent{ChildShutdownEventName};

	zen::ExtendableStringBuilder<32> LogId;
	LogId << "Zen" << ChildId;

	zen::ExtendableWideStringBuilder<128> CommandLine;
	CommandLine << "\"";
	CommandLine.Append(Executable.c_str());
	CommandLine << "\"";

	const bool IsTest = m_Env.IsTestEnvironment();

	if (IsTest)
	{
		CommandLine << " --test --owner-pid " << MyPid << " --log-id " << LogId;
	}

	CommandLine << " --child-id " << ChildEventName;

	if (BasePort)
	{
		CommandLine << " --port " << BasePort;
	}

	if (!m_TestDir.empty())
	{
		CommandLine << " --data-dir ";
		CommandLine << m_TestDir.c_str();
	}

	if (m_MeshEnabled)
	{
		CommandLine << " --mesh";
	}

	std::filesystem::path CurrentDirectory = std::filesystem::current_path();

	spdlog::debug("Spawning server");

	PROCESS_INFORMATION ProcessInfo{};
	STARTUPINFO			StartupInfo{.cb = sizeof(STARTUPINFO)};

	DWORD CreationFlags = 0;

	if (!IsTest)
	{
		CreationFlags |= CREATE_NEW_CONSOLE;
	}

	HANDLE hProcess = NULL;

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

		BOOL Success = CreateProcessW(Executable.c_str(),
									  (LPWSTR)CommandLine.c_str(),
									  ProcessAttributes,
									  ThreadAttributes,
									  InheritHandles,
									  CreationFlags,
									  Environment,
									  CurrentDirectory.c_str(),
									  &StartupInfo,
									  &ProcessInfo);

		if (Success)
		{
			hProcess = ProcessInfo.hProcess;
			CloseHandle(ProcessInfo.hThread);
		}
		else
		{
			DWORD WinError = ::GetLastError();

			if (WinError == ERROR_ELEVATION_REQUIRED)
			{
				// Try launching elevated process

				spdlog::debug("Regular spawn failed - spawning elevated server");

				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 = CommandLine.c_str();

				if (::ShellExecuteEx(&ShellExecuteInfo))
				{
					WinError = NO_ERROR;

					hProcess = ShellExecuteInfo.hProcess;
				}
			}

			if (WinError != NO_ERROR)
			{
				std::error_code err(WinError, std::system_category());

				spdlog::error("Server spawn failed: {}", err.message());

				throw std::system_error(err, "failed to create server process");
			}
		}
	}

	spdlog::debug("Server spawned OK");

	if (IsTest)
	{
		m_Process.Initialize(hProcess);
		m_ShutdownEvent = std::move(ChildShutdownEvent);
	}
	else
	{
		CloseHandle(hProcess);
	}

	m_ReadyEvent = std::move(ChildEvent);
}

void
ZenServerInstance::WaitUntilReady()
{
	m_ReadyEvent.Wait();
}

bool
ZenServerInstance::WaitUntilReady(int Timeout)
{
	return m_ReadyEvent.Wait(Timeout);
}