// 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>

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

#if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
#	include <unistd.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 IsElevated()
	{
		BOOL   fRet	  = FALSE;
		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;
	}

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

		ExtendableWideStringBuilder<256> Parameters;
		std::filesystem::path			 ExecutablePath = GetRunningExecutablePath();
		std::wstring					 CommandLine(GetCommandLine());
		std::string::size_type			 ExtraLength	  = CommandLine[0] == '\"' ? 2 : 0;
		std::wstring					 CommandArguments = CommandLine.substr(ExecutablePath.string().length() + ExtraLength + 1);

		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 ReturnCode = 1;
		if (ShellExecuteEx(&shExInfo))
		{
			WaitForSingleObject(shExInfo.hProcess, INFINITE);
			GetExitCodeProcess(shExInfo.hProcess, &ReturnCode);
			CloseHandle(shExInfo.hProcess);
			if (ReturnCode == 0)
			{
				ZEN_CONSOLE("Elevated execution completed successfully.");
			}
			else
			{
				ZEN_CONSOLE("Elevated execution completed unsuccessfully, return code: '{}'.", ReturnCode);
			}
		}
		else
		{
			ZEN_CONSOLE("Failed to run elevated, operation did not complete.");
		}
		return (int)ReturnCode;
	}

#else  // ZEN_PLATFORM_WINDOWS

	bool IsElevated() { return true; /*geteuid() == 0;*/ }

#endif	// ZEN_PLATFORM_WINDOWS

	int RunElevated(bool AllowElevation)
	{
#if ZEN_PLATFORM_WINDOWS
		if (AllowElevation)
		{
			return WinRelaunchElevated();
		}
		else
		{
			ZEN_CONSOLE(
				"This command requires elevated priviliges. Run command with elevated priviliges or add '--allow-elevation' command line "
				"option.");
			return 1;
		}
#endif	// ZEN_PLATFORM_WINDOWS
#if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
		ZEN_UNUSED(AllowElevation);
		ZEN_CONSOLE("This command requires elevated priviliges. Run the command with `sudo`");
		return 1;
#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>");
#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>");
#endif	// ZEN_PLATFORM_WINDOWS
	m_InstallOptions.parse_positional({"executable", "name", "display-name"});
	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("",
							 mak "",
							 "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;
}

int
ServiceCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	using namespace std::literals;

	std::vector<char*> SubCommandArguments;
	cxxopts::Options*  SubOption			 = nullptr;
	int				   ParentCommandArgCount = GetSubCommand(m_Options, argc, argv, m_SubCommands, SubOption, SubCommandArguments);
	if (!ParseOptions(ParentCommandArgCount, argv))
	{
		return 0;
	}

	if (SubOption == nullptr)
	{
		throw zen::OptionParseException("command verb is missing");
	}

	if (!ParseOptions(*SubOption, gsl::narrow<int>(SubCommandArguments.size()), SubCommandArguments.data()))
	{
		return 0;
	}

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

	if (SubOption == &m_InstallOptions)
	{
		ServiceInfo		Info;
		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
		if (!Ec && Info.Status != ServiceStatus::NotInstalled)
		{
			ZEN_CONSOLE("Service '{}' already installed:\n{}", m_ServiceName, FmtServiceInfo(Info, "  "));
			return 1;
		}

		if (!IsElevated())
		{
			return RunElevated(m_AllowElevation);
		}

		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);
		Ec				   = InstallService(
			m_ServiceName,
			ServiceSpec {
				.ExecutablePath = m_ServerExecutable, .CommandLineOptions = GlobalOptions.PassthroughCommandLine
#if ZEN_PLATFORM_WINDOWS
					,
				.DisplayName = m_ServiceDisplayName, .Description = m_ServiceDescription
#endif	// ZEN_PLATFORM_WINDOWS
			});
		if (Ec)
		{
			ZEN_CONSOLE("Failed to install service '{}' using '{}' . Reason: '{}'", m_ServiceName, m_ServerExecutable, Ec.message());
			return 1;
		}
		ZEN_CONSOLE("Installed service '{}' using '{}' successfully", m_ServiceName, m_ServerExecutable);
	}

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

		if (!IsElevated())
		{
			return RunElevated(m_AllowElevation);
		}

		Ec = UninstallService(m_ServiceName);
		if (Ec)
		{
			ZEN_CONSOLE("Failed to uninstall service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
			return 1;
		}
		ZEN_CONSOLE("Uninstalled service {} successfully", m_ServiceName);
	}

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

		if (!IsElevated())
		{
			return RunElevated(m_AllowElevation);
		}

		Ec = StartService(m_ServiceName);
		if (Ec)
		{
			ZEN_CONSOLE("Failed to start service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
			return 1;
		}
		ZEN_CONSOLE("Started service '{}' successfully", m_ServiceName);
	}

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

		if (!IsElevated())
		{
			return RunElevated(m_AllowElevation);
		}

		Ec = StopService(m_ServiceName);
		if (Ec)
		{
			ZEN_CONSOLE("Failed to stop service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
			return 1;
		}
		ZEN_CONSOLE("Stopped service '{}' successfully", m_ServiceName);
	}

	return 0;
}

}  // namespace zen