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

#include "service_cmd.h"

#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/zencore.h>
#include <zenutil/service.h>
#include <filesystem>

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>
#	include <shellapi.h>
#	include <shlwapi.h>
#	pragma comment(lib, "shlwapi.lib")
#endif

#if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
#	include <unistd.h>
#endif

#if ZEN_PLATFORM_LINUX
#	include <pwd.h>
#	include <grp.h>
#endif

ZEN_THIRD_PARTY_INCLUDES_START
#include <gsl/gsl-lite.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

#include <memory>

namespace zen {

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

namespace {

#if ZEN_PLATFORM_WINDOWS
	BOOL RequiresElevation()
	{
		BOOL   fRet	  = TRUE;
		HANDLE hToken = NULL;
		if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
		{
			TOKEN_ELEVATION Elevation;
			DWORD			cbSize = sizeof(TOKEN_ELEVATION);
			if (GetTokenInformation(hToken, TokenElevation, &Elevation, sizeof(Elevation), &cbSize))
			{
				fRet = !Elevation.TokenIsElevated;
			}
		}
		if (hToken)
		{
			CloseHandle(hToken);
		}
		return fRet;
	}

	void WinRelaunchElevated()
	{
		TCHAR CurrentDir[4096];
		GetCurrentDirectory(4096, CurrentDir);

		ExtendableWideStringBuilder<256> Parameters;
		std::filesystem::path			 ExecutablePath = GetRunningExecutablePath();
		std::wstring					 CommandLine(GetCommandLine());
		std::wstring					 CommandArguments(PathGetArgs(CommandLine.data()));

		ZEN_CONSOLE("Attempting to run '{} {}' elevated...", ExecutablePath, WideToUtf8(CommandArguments));

		SHELLEXECUTEINFO shExInfo = {0};
		shExInfo.cbSize			  = sizeof(shExInfo);
		shExInfo.fMask			  = SEE_MASK_NOCLOSEPROCESS | SEE_MASK_NO_CONSOLE;
		shExInfo.hwnd			  = 0;
		shExInfo.lpVerb			  = TEXT("runas");			   // Operation to perform
		shExInfo.lpFile			  = ExecutablePath.c_str();	   // Application to start
		shExInfo.lpParameters	  = CommandArguments.c_str();  // Additional parameters
		shExInfo.lpDirectory	  = CurrentDir;
		shExInfo.nShow			  = SW_SHOW;
		shExInfo.hInstApp		  = 0;

		DWORD ProcessReturnCode = 1;
		if (ShellExecuteEx(&shExInfo))
		{
			WaitForSingleObject(shExInfo.hProcess, INFINITE);
			GetExitCodeProcess(shExInfo.hProcess, &ProcessReturnCode);
			CloseHandle(shExInfo.hProcess);
			if (ProcessReturnCode == 0)
			{
				ZEN_CONSOLE("Elevated execution completed successfully.");
			}
			else
			{
				throw ErrorWithReturnCode(fmt::format("Elevated execution completed unsuccessfully, return code: '{}'.", ProcessReturnCode),
										  (int)ProcessReturnCode);
			}
		}
		else
		{
			ThrowLastError("Failed to run elevated, operation did not complete");
		}
	}

#elif ZEN_PLATFORM_MAC

	bool RequiresElevation() { return false; }	// Mac service mode commands can run without elevation

#else

	bool RequiresElevation() { return geteuid() != 0; }

#endif	// ZEN_PLATFORM_WINDOWS

	void RunElevated(bool AllowElevation)
	{
#if ZEN_PLATFORM_WINDOWS
		if (AllowElevation)
		{
			WinRelaunchElevated();
		}
		else
		{
			throw std::runtime_error(fmt::format(
				"This command requires elevated priviliges. Run command with elevated priviliges or add '--allow-elevation' command line "
				"option."));
		}
#endif	// ZEN_PLATFORM_WINDOWS
#if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
		ZEN_UNUSED(AllowElevation);
		throw std::runtime_error(fmt::format("This command requires elevated priviliges. Run the command with `sudo`"));
#endif	// ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	}

}  // namespace

ServiceCommand::ServiceCommand()
{
	m_Options.add_option("",
						 "v",
						 "verb",
						 fmt::format("Verb for service - {}, {}, {}, {}, {}. Use '--' ",
									 m_StatusOptions.program(),
									 m_InstallOptions.program(),
									 m_UninstallOptions.program(),
									 m_StartOptions.program(),
									 m_StopOptions.program()),
						 cxxopts::value(m_Verb),
						 "<verb>");
	m_Options.parse_positional({"verb"});
	m_Options.positional_help("verb");

	m_StatusOptions.add_options()("h,help", "Print help");
	m_StatusOptions.add_option("",
							   "n",
							   "name",
							   fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
							   cxxopts::value(m_ServiceName),
							   "<name>");
	m_StatusOptions.parse_positional({"name"});
	m_StatusOptions.positional_help("name");

	m_InstallOptions.add_options()("h,help", "Print help");
	m_InstallOptions.add_option("", "s", "executable", "Path to server executable", cxxopts::value(m_ServerExecutable), "<path>");
	m_InstallOptions.add_option("",
								"n",
								"name",
								fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
								cxxopts::value(m_ServiceName),
								"<name>");

	m_InstallOptions.add_option("",
								"",
								"full",
								fmt::format("Uninstall a running service and update service binaries before installing"),
								cxxopts::value(m_FullInstall),
								"<full>");
	m_InstallOptions.add_option("",
								"",
								"install-path",
								fmt::format("Path in which to install service binaries"),
								cxxopts::value(m_InstallPath),
								"<install-path>");

	m_InstallOptions.add_option("", "u", "user", "User to run service as, defaults to current user", cxxopts::value(m_UserName), "<user>");
#if ZEN_PLATFORM_WINDOWS
	m_InstallOptions.add_option("",
								"d",
								"display-name",
								fmt::format("Service display, defaults to \"{}\"", m_ServiceDisplayName),
								cxxopts::value(m_ServiceDisplayName),
								"<display-name>");
	m_InstallOptions.add_option("",
								"",
								"description",
								fmt::format("Service description", m_ServiceDescription),
								cxxopts::value(m_ServiceDescription),
								"<description>");
	m_InstallOptions.add_option("",
								"",
								"allow-elevation",
								fmt::format("Allow attempt to relauch command using an elevated"),
								cxxopts::value(m_AllowElevation),
								"<allow-elevation>");
	m_InstallOptions.parse_positional({"executable", "name", "display-name"});
#else
	m_InstallOptions.parse_positional({"executable", "name"});
#endif	// ZEN_PLATFORM_WINDOWS
	m_InstallOptions.positional_help("executable name display-name");

	m_UninstallOptions.add_options()("h,help", "Print help");
	m_UninstallOptions.add_option("",
								  "n",
								  "name",
								  fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
								  cxxopts::value(m_ServiceName),
								  "<name>");
	m_UninstallOptions.parse_positional({"name"});
	m_UninstallOptions.positional_help("name");
#if ZEN_PLATFORM_WINDOWS
	m_UninstallOptions.add_option("",
								  "",
								  "allow-elevation",
								  fmt::format("Allow attempt to relauch command using an elevated"),
								  cxxopts::value(m_AllowElevation),
								  "<allow-elevation>");
#endif	// ZEN_PLATFORM_WINDOWS

	m_StartOptions.add_options()("h,help", "Print help");
	m_StartOptions.add_option("",
							  "n",
							  "name",
							  fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
							  cxxopts::value(m_ServiceName),
							  "<name>");
	m_StartOptions.parse_positional({"name"});
	m_StartOptions.positional_help("name");
#if ZEN_PLATFORM_WINDOWS
	m_StartOptions.add_option("",
							  "",
							  "allow-elevation",
							  fmt::format("Allow attempt to relauch command using an elevated"),
							  cxxopts::value(m_AllowElevation),
							  "<allow-elevation>");
#endif	// ZEN_PLATFORM_WINDOWS

	m_StopOptions.add_options()("h,help", "Print help");
	m_StopOptions.add_option("",
							 "n",
							 "name",
							 fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
							 cxxopts::value(m_ServiceName),
							 "<name>");
	m_StopOptions.parse_positional({"name"});
	m_StopOptions.positional_help("name");
#if ZEN_PLATFORM_WINDOWS
	m_StopOptions.add_option("",
							 "",
							 "allow-elevation",
							 fmt::format("Allow attempt to relauch command using an elevated"),
							 cxxopts::value(m_AllowElevation),
							 "<allow-elevation>");
#endif	// ZEN_PLATFORM_WINDOWS
}

ServiceCommand::~ServiceCommand() = default;

std::string
FmtServiceInfo(const ServiceInfo& Info, std::string_view Prefix)
{
	std::string Result = fmt::format(
		"{}Status:             {}\n"
		"{}Executable:         {}\n"
		"{}CommandLineOptions: {}",
		Prefix,
		ToString(Info.Status),
		Prefix,
		Info.Spec.ExecutablePath,
		Prefix,
		Info.Spec.CommandLineOptions);

#if ZEN_PLATFORM_WINDOWS
	Result += fmt::format(
		"\n"
		"{}Display Name:       {}\n"
		"{}Description:        {}",
		Prefix,
		Info.Spec.DisplayName,
		Prefix,
		Info.Spec.Description);
#endif	// ZEN_PLATFORM_WINDOWS

	return Result;
}

enum class ServiceStatusReturnCode
{
	Running = 0,  // successful return indicates proper installation and everything running smoothly
	NotInstalled,
	NotRunning,
	UnknownError
};

void
ServiceCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	using namespace std::literals;

	std::vector<char*> SubCommandArguments;
	cxxopts::Options*  SubOption			 = nullptr;
	int				   ParentCommandArgCount = GetSubCommand(m_Options, argc, argv, m_SubCommands, SubOption, SubCommandArguments);

	if (SubOption == nullptr)
	{
		if (!ParseOptions(ParentCommandArgCount, argv))
		{
			return;
		}
		throw OptionParseException("'verb' option is required", m_Options.help());
	}

	// Parse subcommand permissively - forward unrecognised options to the parent parser.
	std::vector<std::string> SubUnmatched;
	if (!ParseOptionsPermissive(*SubOption, gsl::narrow<int>(SubCommandArguments.size()), SubCommandArguments.data(), SubUnmatched))
	{
		return;
	}

	// Build parent arg list: original parent args (without subcommand name) + forwarded unmatched.
	std::vector<char*> ParentArgs;
	ParentArgs.reserve(static_cast<size_t>(ParentCommandArgCount - 1) + SubUnmatched.size());
	ParentArgs.push_back(argv[0]);
	std::copy(argv + 1, argv + ParentCommandArgCount - 1, std::back_inserter(ParentArgs));
	for (std::string& Arg : SubUnmatched)
	{
		ParentArgs.push_back(Arg.data());
	}

	if (!ParseOptions(static_cast<int>(ParentArgs.size()), ParentArgs.data()))
	{
		return;
	}

	if (SubOption == &m_StatusOptions)
	{
		ServiceInfo		Info;
		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
		if (Ec)
		{
			throw std::runtime_error(fmt::format("Can't get information about service '{}'. Reason: '{}'", m_ServiceName, Ec.message()));
		}
		if (Info.Status == ServiceStatus::NotInstalled)
		{
			throw std::runtime_error(fmt::format("Service '{}' is not installed", m_ServiceName));
		}
		else if (Info.Status != ServiceStatus::Running)
		{
			throw std::runtime_error(fmt::format("Service '{}' is not running", m_ServiceName));
		}
		else
		{
			ZEN_CONSOLE("Service '{}':\n{}", m_ServiceName, FmtServiceInfo(Info, "  "));
		}
	}

	if (SubOption == &m_InstallOptions)
	{
		if (RequiresElevation())
		{
			RunElevated(m_AllowElevation);
			return;
		}

		ServiceInfo		Info;
		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
		if (!Ec && Info.Status != ServiceStatus::NotInstalled)
		{
			if (m_FullInstall)
			{
				if (Info.Status == ServiceStatus::Running)
				{
					Ec = StopService(m_ServiceName);
					if (Ec)
					{
						throw std::runtime_error(fmt::format("Failed to stop service '{}' using '{}'. Reason: '{}'",
															 m_ServiceName,
															 m_ServerExecutable,
															 Ec.message()));
					}

					int Timeout = 30000;  // Wait up to 30 seconds for the service to fully stop before uninstalling
					while (Timeout > 0)
					{
						Ec = QueryInstalledService(m_ServiceName, Info);
						if (Ec)
						{
							throw std::runtime_error(fmt::format("Failed to wait for service to stop: '{}'", Ec.message()));
						}

						if (Info.Status == ServiceStatus::Stopped)
						{
							break;
						}

						Sleep(100);
						Timeout -= 100;
					}

					if (Info.Status != ServiceStatus::Stopped)
					{
						throw std::runtime_error("Timed out waiting for service to stop");
					}
				}

				Ec = UninstallService(m_ServiceName);
				if (Ec)
				{
					throw std::runtime_error(fmt::format("Failed to uninstall running service '{}' using '{}'. Reason: '{}'",
														 m_ServiceName,
														 m_ServerExecutable,
														 Ec.message()));
				}
			}
			else
			{
				ZEN_CONSOLE("Service '{}' already installed:\n{}", m_ServiceName, FmtServiceInfo(Info, "  "));
				return;
			}
		}
		if (m_ServerExecutable.empty())
		{
			std::filesystem::path ExePath = zen::GetRunningExecutablePath();
			ExePath.replace_filename("zenserver" + ExePath.extension().string());
			m_ServerExecutable = ExePath;
		}
		m_ServerExecutable = std::filesystem::absolute(m_ServerExecutable);

		if (m_FullInstall)
		{
			if (m_InstallPath.empty())
			{
				throw OptionParseException("'--full' requires '--install-path'", SubOption->help());
			}

			std::filesystem::path ExePath = zen::GetRunningExecutablePath();

			std::filesystem::path FilesToCopy[] = {
				ExePath,
				m_ServerExecutable,
#if ZEN_PLATFORM_WINDOWS
				ExePath.parent_path() / ExePath.stem().replace_extension("pdb"),
				m_ServerExecutable.parent_path() / m_ServerExecutable.stem().replace_extension("pdb"),
				ExePath.parent_path() /
					std::filesystem::path("crashpad_handler.exe")
#endif
#if ZEN_PLATFORM_MAC
						ExePath.parent_path() /
					std::filesystem::path("crashpad_handler")
#endif
			};

			std::filesystem::path DestinationExePath	= m_InstallPath / ExePath.filename();
			std::filesystem::path DestinationServerPath = m_InstallPath / m_ServerExecutable.filename();

			if (!std::filesystem::is_directory(m_InstallPath) && !CreateDirectories(m_InstallPath))
			{
				throw std::runtime_error(fmt::format("Unable to create install directory '{}'", m_InstallPath));
			}

#if ZEN_PLATFORM_LINUX
			uid_t UserId  = 0;
			gid_t GroupId = 0;
			if (!m_UserName.empty())
			{
				struct passwd* Passwd = getpwnam(m_UserName.c_str());
				if (Passwd == NULL)
				{
					throw std::runtime_error(fmt::format("Unable to determine user ID for user '{}'", m_UserName));
				}

				UserId = Passwd->pw_uid;

				struct group* Grp = getgrnam(m_UserName.c_str());
				if (Grp == NULL)
				{
					throw std::runtime_error(fmt::format("Unable to determine group ID for user '{}'", m_UserName));
				}

				GroupId = Grp->gr_gid;
				if (chown(m_InstallPath.c_str(), UserId, GroupId) != 0)
				{
					throw std::runtime_error(
						fmt::format("Unable to set ownership of directory '{}' for user '{}'", m_InstallPath, m_UserName));
				}
			}
#endif

			for (const std::filesystem::path& File : FilesToCopy)
			{
				std::filesystem::path Destination = m_InstallPath / File.filename();

				if (std::error_code CopyEc = CopyFile(File, Destination, {.EnableClone = false}); CopyEc)
				{
					throw std::system_error(CopyEc, fmt::format("Failed to copy '{}' to '{}'", File, Destination));
				}

				ZEN_INFO("Copied '{}' to '{}'", File, Destination);

				if (File.extension() != "pdb")
				{
					std::filesystem::permissions(Destination, std::filesystem::perms::owner_exec, std::filesystem::perm_options::add);
				}

				std::filesystem::permissions(Destination,
											 std::filesystem::perms::owner_write | std::filesystem::perms::group_write,
											 std::filesystem::perm_options::add);

#if ZEN_PLATFORM_LINUX
				if (UserId != 0)
				{
					if (chown(Destination.c_str(), UserId, GroupId) != 0)
					{
						throw std::runtime_error(
							fmt::format("Unable to set ownership of file '{}' for user '{}'", Destination, m_UserName));
					}
				}
#endif
			}

			m_ServerExecutable = m_InstallPath / m_ServerExecutable.filename();
		}

		Ec = InstallService(
			m_ServiceName,
			ServiceSpec {
				.ExecutablePath = m_ServerExecutable, .CommandLineOptions = GlobalOptions.PassthroughCommandLine, .UserName = m_UserName
#if ZEN_PLATFORM_WINDOWS
					,
				.DisplayName = m_ServiceDisplayName, .Description = m_ServiceDescription
#endif	// ZEN_PLATFORM_WINDOWS
			});
		if (Ec)
		{
			throw std::runtime_error(
				fmt::format("Failed to install service '{}' using '{}' . Reason: '{}'", m_ServiceName, m_ServerExecutable, Ec.message()));
		}
		ZEN_CONSOLE("Installed service '{}' using '{}' successfully", m_ServiceName, m_ServerExecutable);

		if (m_FullInstall)
		{
			Ec = StartService(m_ServiceName);
			if (Ec)
			{
				throw std::runtime_error(fmt::format("Failed to start service '{}'. Reason: '{}'", m_ServiceName, Ec.message()));
			}
		}
	}

	if (SubOption == &m_UninstallOptions)
	{
		ServiceInfo		Info;
		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
		if (Ec)
		{
			throw std::runtime_error(fmt::format("Failed to inspect installed service '{}'. Reason: '{}'", m_ServiceName, Ec.message()));
		}
		if (Info.Status == ServiceStatus::NotInstalled)
		{
			ZEN_CONSOLE("Service '{}' is not installed", m_ServiceName);
			return;
		}
		if (Info.Status != ServiceStatus::Stopped)
		{
			throw std::runtime_error(fmt::format("Service '{}' is running, stop before uninstalling", m_ServiceName));
		}

		if (RequiresElevation())
		{
			RunElevated(m_AllowElevation);
			return;
		}

		Ec = UninstallService(m_ServiceName);
		if (Ec)
		{
			throw std::runtime_error(fmt::format("Failed to uninstall service '{}'. Reason: '{}'", m_ServiceName, Ec.message()));
		}
		ZEN_CONSOLE("Uninstalled service {} successfully", m_ServiceName);
	}

	if (SubOption == &m_StartOptions)
	{
		ServiceInfo		Info;
		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
		if (Ec)
		{
			throw std::runtime_error(fmt::format("Failed to inspect installed service '{}'. Reason: '{}'", m_ServiceName, Ec.message()));
		}
		if (Info.Status == ServiceStatus::NotInstalled)
		{
			throw std::runtime_error(fmt::format("Service '{}' is not installed", m_ServiceName));
		}
		if (Info.Status != ServiceStatus::Stopped)
		{
			ZEN_CONSOLE("Service '{}' is already running", m_ServiceName);
			return;
		}

		if (RequiresElevation())
		{
			RunElevated(m_AllowElevation);
			return;
		}

		Ec = StartService(m_ServiceName);
		if (Ec)
		{
			throw std::runtime_error(fmt::format("Failed to start service '{}'. Reason: '{}'", m_ServiceName, Ec.message()));
		}
		ZEN_CONSOLE("Started service '{}' successfully", m_ServiceName);
	}

	if (SubOption == &m_StopOptions)
	{
		ServiceInfo		Info;
		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
		if (Ec)
		{
			throw std::runtime_error(fmt::format("Failed to inspect installed service '{}'. Reason: '{}'", m_ServiceName, Ec.message()));
		}
		if (Info.Status == ServiceStatus::NotInstalled)
		{
			throw std::runtime_error(fmt::format("Service '{}' is not installed", m_ServiceName));
		}
		if (Info.Status != ServiceStatus::Running)
		{
			ZEN_CONSOLE("Service '{}' is not running", m_ServiceName);
			return;
		}

		if (RequiresElevation())
		{
			RunElevated(m_AllowElevation);
			return;
		}

		Ec = StopService(m_ServiceName);
		if (Ec)
		{
			throw std::runtime_error(fmt::format("Failed to stop service '{}'. Reason: '{}'", m_ServiceName, Ec.message()));
		}
		ZEN_CONSOLE("Stopped service '{}' successfully", m_ServiceName);
	}
}

}  // namespace zen