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

#include "zenserver.h"

#include <zenbase/refcount.h>
#include <zencore/basicfile.h>
#include <zencore/compactbinarybuilder.h>
#include <zencore/compactbinaryvalidation.h>
#include <zencore/config.h>
#include <zencore/except.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/iobuffer.h>
#include <zencore/jobqueue.h>
#include <zencore/logging.h>
#include <zencore/memory/fmalloc.h>
#include <zencore/scopeguard.h>
#include <zencore/sentryintegration.h>
#include <zencore/session.h>
#include <zencore/string.h>
#include <zencore/thread.h>
#include <zencore/timer.h>
#include <zencore/trace.h>
#include <zencore/workthreadpool.h>
#include <zenhttp/httpserver.h>
#include <zentelemetry/otlptrace.h>
#include <zenutil/service.h>
#include <zenutil/workerpools.h>
#include <zenutil/zenserverprocess.h>

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

ZEN_THIRD_PARTY_INCLUDES_START
#include <fmt/format.h>
#include <asio.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

#include <exception>
#include <set>

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

#include "config/config.h"
#include "diag/logging.h"

#include <zencore/memory/llm.h>

namespace zen {

const FLLMTag&
GetZenserverTag()
{
	static FLLMTag _("zenserver");

	return _;
}

namespace utils {
	extern std::atomic_uint32_t SignalCounter[NSIG];
}

using namespace std::literals;

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

ZenServerBase::ZenServerBase()
{
}

ZenServerBase::~ZenServerBase()
{
}

int
ZenServerBase::Initialize(const ZenServerConfig& ServerOptions, ZenServerState::ZenServerEntry* ServerEntry)
{
	ZEN_OTEL_SPAN("ZenServerBase::Initialize");
	ZEN_TRACE_CPU("ZenServerBase::Initialize");
	ZEN_MEMSCOPE(GetZenserverTag());

	m_IsPowerCycle = ServerOptions.IsPowerCycle;

	const std::string MutexName = fmt::format("zen_{}", ServerOptions.BasePort);

	if (NamedMutex::Exists(MutexName))
	{
		ZEN_WARN("Mutex '{}' already exists - is another instance already running?", MutexName);
		return -1;
	}

	m_UseSentry	  = ServerOptions.SentryConfig.Disable == false;
	m_ServerEntry = ServerEntry;

	// Initialize parent (sponsor) process monitoring

	if (const int ParentPid = ServerOptions.OwnerPid)
	{
		std::error_code Ec;
		ProcessHandle	OwnerProcess;
		OwnerProcess.Initialize(ParentPid, /* out */ Ec);

		if (!OwnerProcess.IsValid())
		{
			ZEN_WARN("Unable to initialize process handle for specified parent pid #{}. Reason: '{}'", ParentPid, Ec.message());

			// If the pid is not reachable should we just shut down immediately? the intended owner process
			// could have been killed or somehow crashed already
		}
		else
		{
			ZEN_INFO("Using parent pid #{} to control process lifetime", ParentPid);
		}

		m_ProcessMonitor.AddPid(ParentPid);
	}

	// Initialize/check mutex based on base port

	if (m_ServerMutex.Create(MutexName) == false)
	{
		std::error_code Ec = MakeErrorCodeFromLastError();
		ZEN_WARN("Failed to create server mutex '{}'. Reason: '{}' ({})", MutexName, Ec.message(), Ec.value());
		return -1;
	}

	EnqueueSigIntTimer();

	m_Http				  = CreateHttpServer(ServerOptions.HttpConfig);
	int EffectiveBasePort = m_Http->Initialize(ServerOptions.BasePort, ServerOptions.DataDir);
	if (EffectiveBasePort == 0)
	{
		ZEN_WARN("Failed to initialize http service '{}' using base port {} and data dir '{}'",
				 ServerOptions.HttpConfig.ServerClass,
				 ServerOptions.BasePort,
				 ServerOptions.DataDir);
		return -1;
	}

	ZEN_INFO("Effective concurrency: {} (hw: {})", GetHardwareConcurrency(), std::thread::hardware_concurrency());

	m_StatusService.RegisterHandler("status", *this);
	m_Http->RegisterService(m_StatusService);

	m_StatsReporter.Initialize(ServerOptions.StatsConfig);
	if (ServerOptions.StatsConfig.Enabled)
	{
		EnqueueStatsReportingTimer();
	}

	LogSettingsSummary(ServerOptions);

	return EffectiveBasePort;
}

void
ZenServerBase::Finalize()
{
	// Register health service last so if we return "OK" for health it means all services have been properly initialized

	m_Http->RegisterService(m_HealthService);
}

void
ZenServerBase::GetBuildOptions(StringBuilderBase& OutOptions, char Separator) const
{
	ZEN_MEMSCOPE(GetZenserverTag());

	OutOptions << "ZEN_ADDRESS_SANITIZER=" << (ZEN_ADDRESS_SANITIZER ? "1" : "0");
	OutOptions << Separator;
	OutOptions << "ZEN_USE_SENTRY=" << (ZEN_USE_SENTRY ? "1" : "0");
	OutOptions << Separator;
	OutOptions << "ZEN_WITH_TESTS=" << (ZEN_WITH_TESTS ? "1" : "0");
	OutOptions << Separator;
	OutOptions << "ZEN_USE_MIMALLOC=" << (ZEN_USE_MIMALLOC ? "1" : "0");
	OutOptions << Separator;
	OutOptions << "ZEN_USE_RPMALLOC=" << (ZEN_USE_RPMALLOC ? "1" : "0");
	OutOptions << Separator;
	OutOptions << "ZEN_WITH_HTTPSYS=" << (ZEN_WITH_HTTPSYS ? "1" : "0");
	OutOptions << Separator;
	OutOptions << "ZEN_WITH_MEMTRACK=" << (ZEN_WITH_MEMTRACK ? "1" : "0");
	OutOptions << Separator;
	OutOptions << "ZEN_WITH_TRACE=" << (ZEN_WITH_TRACE ? "1" : "0");
}

void
ZenServerBase::EnsureIoRunner()
{
	ZEN_MEMSCOPE(GetZenserverTag());
	if (!m_IoRunner.joinable())
	{
		m_IoRunner = std::thread{[this] {
			SetCurrentThreadName("timer_io");
			m_IoContext.run();
		}};
	}
}

void
ZenServerBase::OnReady()
{
	if (m_ServerEntry)
	{
		m_ServerEntry->SignalReady();
	}

	if (m_IsReadyFunc)
	{
		m_IsReadyFunc();
	}
}

void
ZenServerBase::RequestExit(int ExitCode)
{
	if (RequestApplicationExit(ExitCode))
	{
		if (m_Http)
		{
			m_Http->RequestExit();
		}
	}
}

std::string_view
ZenServerBase::ToString(ServerState Value)
{
	switch (Value)
	{
		case kInitializing:
			return "initializing"sv;
		case kRunning:
			return "running"sv;
		case kShuttingDown:
			return "shutdown"sv;
		default:
			return "unknown"sv;
	}
}

void
ZenServerBase::EnqueueStatsReportingTimer()
{
	ZEN_MEMSCOPE(GetZenserverTag());
	m_StatsReportingTimer.expires_after(std::chrono::milliseconds(500));
	m_StatsReportingTimer.async_wait([this](const asio::error_code& Ec) {
		if (!Ec)
		{
			m_StatsReporter.ReportStats();
			EnqueueStatsReportingTimer();
		}
	});
	EnsureIoRunner();
}

void
ZenServerBase::EnqueueProcessMonitorTimer()
{
	ZEN_MEMSCOPE(GetZenserverTag());
	m_PidCheckTimer.expires_after(std::chrono::seconds(1));
	m_PidCheckTimer.async_wait([this](const asio::error_code&) { CheckOwnerPid(); });

	EnsureIoRunner();
}

void
ZenServerBase::CheckOwnerPid()
{
	bool IsRunning = UpdateProcessMonitor();

	if (IsRunning)
	{
		m_FoundNoActiveSponsors = false;
		EnqueueProcessMonitorTimer();
	}
	else
	{
		// Delay exit one iteration to avoid race conditions where one process detaches
		// and another attaches
		if (m_FoundNoActiveSponsors)
		{
			ZEN_INFO(ZEN_APP_NAME " exiting since sponsor processes are all gone");
			RequestExit(0);
		}
		else
		{
			m_FoundNoActiveSponsors = true;
			EnqueueProcessMonitorTimer();
		}
	}
}

bool
ZenServerBase::UpdateProcessMonitor()
{
	if (m_ServerEntry)
	{
		// Pick up any new "owner" processes
		std::set<uint32_t> AddedPids;

		for (auto& PidEntry : m_ServerEntry->SponsorPids)
		{
			if (uint32_t ThisPid = PidEntry.load(std::memory_order_relaxed))
			{
				if (PidEntry.compare_exchange_strong(ThisPid, 0))
				{
					if (AddedPids.insert(ThisPid).second)
					{
						m_ProcessMonitor.AddPid(ThisPid);

						ZEN_INFO("added process with pid {} as a sponsor process", ThisPid);
					}
				}
			}
		}
	}
	return m_ProcessMonitor.IsRunning();
}

void
ZenServerBase::EnqueueStateExitFlagTimer()
{
	ZEN_MEMSCOPE(GetZenserverTag());
	m_StateExitFlagTimer.expires_after(std::chrono::milliseconds(500));
	m_StateExitFlagTimer.async_wait([this](const asio::error_code&) { CheckStateExitFlag(); });
	EnsureIoRunner();
}

void
ZenServerBase::CheckStateExitFlag()
{
	if (m_ServerEntry && m_ServerEntry->IsShutdownRequested())
	{
		RequestExit(0);
		return;
	}
	EnqueueStateExitFlagTimer();
}

void
ZenServerBase::EnqueueSigIntTimer()
{
	ZEN_MEMSCOPE(GetZenserverTag());
	m_SigIntTimer.expires_after(std::chrono::milliseconds(500));
	m_SigIntTimer.async_wait([this](const asio::error_code&) { CheckSigInt(); });
	EnsureIoRunner();
}

void
ZenServerBase::CheckSigInt()
{
	if (utils::SignalCounter[SIGINT] > 0)
	{
		ZEN_INFO("SIGINT triggered (Ctrl+C) for process {}, exiting", zen::GetCurrentProcessId());
		RequestExit(128 + SIGINT);
		return;
	}
	if (utils::SignalCounter[SIGTERM] > 0)
	{
		ZEN_INFO("SIGTERM triggered for process {}, exiting", zen::GetCurrentProcessId());
		RequestExit(128 + SIGTERM);
		return;
	}
	EnqueueSigIntTimer();
}

void
ZenServerBase::HandleStatusRequest(HttpServerRequest& Request)
{
	CbObjectWriter Cbo;
	Cbo << "ok" << true;
	Cbo << "state" << ToString(m_CurrentState);
	Request.WriteResponse(HttpResponseCode::OK, Cbo.Save());
}

void
ZenServerBase::LogSettingsSummary(const ZenServerConfig& ServerConfig)
{
	// clang-format off
	std::list<std::pair<std::string_view, std::string>> Settings = {
		{"DataDir"sv, 				ServerConfig.DataDir.string()},
		{"AbsLogFile"sv,			ServerConfig.AbsLogFile.string()},
		{"SystemRootDir"sv, 		ServerConfig.SystemRootDir.string()},
		{"ContentDir"sv, 			ServerConfig.ContentDir.string()},
		{"BasePort"sv, 				fmt::to_string(ServerConfig.BasePort)},
		{"IsDebug"sv, 				fmt::to_string(ServerConfig.IsDebug)},
		{"IsCleanStart"sv, 			fmt::to_string(ServerConfig.IsCleanStart)},
		{"IsPowerCycle"sv, 			fmt::to_string(ServerConfig.IsPowerCycle)},
		{"IsTest"sv, 				fmt::to_string(ServerConfig.IsTest)},
		{"Detach"sv, 				fmt::to_string(ServerConfig.Detach)},
		{"NoConsoleOutput"sv, 		fmt::to_string(ServerConfig.NoConsoleOutput)},
		{"QuietConsole"sv, 			fmt::to_string(ServerConfig.QuietConsole)},
		{"CoreLimit"sv, 			fmt::to_string(ServerConfig.CoreLimit)},
		{"IsDedicated"sv, 			fmt::to_string(ServerConfig.IsDedicated)},
		{"ShouldCrash"sv, 			fmt::to_string(ServerConfig.ShouldCrash)},
		{"ChildId"sv,				ServerConfig.ChildId},
		{"LogId"sv,					ServerConfig.LogId},
		{"Sentry DSN"sv,			ServerConfig.SentryConfig.Dsn.empty() ? "not set" : ServerConfig.SentryConfig.Dsn},
		{"Sentry Environment"sv,	ServerConfig.SentryConfig.Environment},
		{"Statsd Enabled"sv,		fmt::to_string(ServerConfig.StatsConfig.Enabled)},
	};
	// clang-format on

	if (ServerConfig.StatsConfig.Enabled)
	{
		Settings.emplace_back("Statsd Host", ServerConfig.StatsConfig.StatsdHost);
		Settings.emplace_back("Statsd Port", fmt::to_string(ServerConfig.StatsConfig.StatsdPort));
	}

	size_t MaxWidth = 0;
	for (const auto& Setting : Settings)
	{
		MaxWidth = std::max(MaxWidth, Setting.first.size());
	}

	ZEN_INFO("Server settings summary:");

	for (const auto& Setting : Settings)
	{
		if (!Setting.second.empty())
		{
			ZEN_INFO("  {:<{}} : {}", Setting.first, MaxWidth, Setting.second);
		}
	}
}

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

ZenServerMain::ZenServerMain(ZenServerConfig& ServerOptions) : m_ServerOptions(ServerOptions)
{
}

ZenServerMain::~ZenServerMain()
{
}

void
ZenServerMain::InitializeLogging()
{
	InitializeServerLogging(m_ServerOptions, /* WithCacheService */ false);
}

int
ZenServerMain::Run()
{
	// On Linux this has the unfortunate side effect of making `top` and other tools display
	// `main` as the program name since threads and processes have a closer relationship
	// there. So we don't name the main thread explicitly there.

#ifndef ZEN_PLATFORM_LINUX
	zen::SetCurrentThreadName("main");
#endif

#if ZEN_USE_SENTRY
	SentryIntegration Sentry;

	if (m_ServerOptions.SentryConfig.Disable == false)
	{
		ZEN_OTEL_SPAN("SentryInit");

		std::string SentryDatabasePath	 = (m_ServerOptions.DataDir / ".sentry-native").string();
		std::string SentryAttachmentPath = m_ServerOptions.AbsLogFile.string();

		Sentry.Initialize({.DatabasePath	= SentryDatabasePath,
						   .AttachmentsPath = SentryAttachmentPath,
						   .Dsn				= m_ServerOptions.SentryConfig.Dsn,
						   .Environment		= m_ServerOptions.SentryConfig.Environment,
						   .AllowPII		= m_ServerOptions.SentryConfig.AllowPII,
						   .Debug			= m_ServerOptions.SentryConfig.Debug},
						  m_ServerOptions.CommandLine);
	}
#endif

	try
	{
		// Mutual exclusion and synchronization
		ZenServerState ServerState;
		ServerState.Initialize();
		ServerState.Sweep();

		uint32_t						AttachSponsorProcessRetriesLeft = 3;
		ZenServerState::ZenServerEntry* Entry							= ServerState.Lookup(m_ServerOptions.BasePort);
		while (Entry)
		{
			if (m_ServerOptions.OwnerPid)
			{
				std::error_code Ec;
				if (!IsProcessRunning(m_ServerOptions.OwnerPid, Ec))
				{
					if (Ec)
					{
						ZEN_WARN(ZEN_APP_NAME
								 " exiting, sponsor owner pid {} can not be checked for running state, reason: '{}'. Will not add sponsor "
								 "to process "
								 "listening to port {} (pid: {})",
								 m_ServerOptions.OwnerPid,
								 Ec.message(),
								 m_ServerOptions.BasePort,
								 Entry->Pid.load());
					}
					else
					{
						ZEN_WARN(ZEN_APP_NAME
								 " exiting, sponsor owner pid {} is no longer running, will not add sponsor to process listening to port "
								 "{} (pid: {})",
								 m_ServerOptions.OwnerPid,
								 m_ServerOptions.BasePort,
								 Entry->Pid.load());
					}
					std::exit(1);
				}
				ZEN_INFO(
					"Looks like there is already a process listening to this port {} (pid: {}), attaching owner pid {} to running instance",
					m_ServerOptions.BasePort,
					Entry->Pid.load(),
					m_ServerOptions.OwnerPid);

				// Sponsor processes are checked every second, so 2 second wait time should be enough
				if (Entry->AddSponsorProcess(m_ServerOptions.OwnerPid, 2000))
				{
					NotifyReady();
					std::exit(0);
				}
				if (AttachSponsorProcessRetriesLeft-- > 0)
				{
					Entry = ServerState.Lookup(m_ServerOptions.BasePort);
				}
				else
				{
					ZEN_WARN(ZEN_APP_NAME " exiting, failed to add sponsor owner pid {} to process listening to port {} (pid: {})",
							 m_ServerOptions.OwnerPid,
							 m_ServerOptions.BasePort,
							 Entry->Pid.load());
					std::exit(1);
				}
			}
			else
			{
				ZEN_WARN(ZEN_APP_NAME " exiting, there is already a process listening to port {} (pid: {})",
						 m_ServerOptions.BasePort,
						 Entry->Pid.load());
				std::exit(1);
			}
		}

		std::error_code Ec;

		std::filesystem::path LockFilePath = m_ServerOptions.DataDir / ".lock";

		m_LockFile.Create(LockFilePath, MakeLockData(false), Ec);

		if (Ec)
		{
			ZEN_INFO(ZEN_APP_NAME " unable to grab lock at '{}' (reason: '{}'), retrying", LockFilePath, Ec.message());
			Sleep(100);

			m_LockFile.Create(LockFilePath, MakeLockData(false), Ec);
			if (Ec)
			{
				ZEN_INFO(ZEN_APP_NAME " unable to grab lock at '{}' (reason: '{}'), retrying", LockFilePath, Ec.message());
				Sleep(500);
				if (Ec)
				{
					ZEN_WARN(ZEN_APP_NAME " exiting, unable to grab lock at '{}' (reason: '{}')", LockFilePath, Ec.message());
					std::exit(99);
				}
			}
		}

		InitializeLogging();

		ZEN_INFO("Command line: {}", m_ServerOptions.CommandLine);

#if ZEN_USE_SENTRY
		Sentry.LogStartupInformation();
#endif

		MaximizeOpenFileCount();

		ZEN_INFO(ZEN_APP_NAME " - using lock file at '{}'", LockFilePath);

		ZEN_INFO(ZEN_APP_NAME " - starting on port {}, version '{}'", m_ServerOptions.BasePort, ZEN_CFG_VERSION_BUILD_STRING_FULL);

		Entry = ServerState.Register(m_ServerOptions.BasePort);

		if (m_ServerOptions.OwnerPid)
		{
			// We are adding a sponsor process to our own entry, can't wait for pick since the code is not run until later
			Entry->AddSponsorProcess(m_ServerOptions.OwnerPid, 0);
		}

		// Run the actual application logic

		DoRun(Entry);
	}
	catch (const AssertException& AssertEx)
	{
		ZEN_CRITICAL(ZEN_APP_NAME " caught assert exception in main for process {}: {}",
					 zen::GetCurrentProcessId(),
					 AssertEx.FullDescription());
		RequestApplicationExit(1);
	}
	catch (const std::system_error& e)
	{
		ZEN_CRITICAL(ZEN_APP_NAME " caught system error exception in main for process {}: {} ({})",
					 zen::GetCurrentProcessId(),
					 e.what(),
					 e.code().value());
		RequestApplicationExit(1);
	}
	catch (const std::exception& e)
	{
		ZEN_CRITICAL(ZEN_APP_NAME " caught exception in main for process {}: {}", zen::GetCurrentProcessId(), e.what());
		RequestApplicationExit(1);
	}

	ShutdownServerLogging();

	ReportServiceStatus(ServiceStatus::Stopped);

	return ApplicationExitCode();
}

void
ZenServerMain::NotifyReady()
{
	if (!m_ServerOptions.ChildId.empty())
	{
		NamedEvent ParentEvent{m_ServerOptions.ChildId};
		ParentEvent.Set();
	}
}

CbObject
ZenServerMain::MakeLockData(bool IsReady)
{
	return MakeLockFilePayload({.Pid				 = GetCurrentProcessId(),
								.SessionId			 = GetSessionId(),
								.EffectiveListenPort = gsl::narrow<uint16_t>(m_ServerOptions.BasePort),
								.Ready				 = IsReady,
								.DataDir			 = m_ServerOptions.DataDir,
								.ExecutablePath		 = GetRunningExecutablePath()});
};

}  // namespace zen