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

#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/logging.h>
#include <zencore/memory/fmalloc.h>
#include <zencore/memory/llm.h>
#include <zencore/memory/memory.h>
#include <zencore/memory/memorytrace.h>
#include <zencore/memory/newdelete.h>
#include <zencore/scopeguard.h>
#include <zencore/sentryintegration.h>
#include <zencore/session.h>
#include <zencore/string.h>
#include <zencore/thread.h>
#include <zencore/trace.h>
#include <zentelemetry/otlptrace.h>
#include <zenutil/config/commandlineoptions.h>
#include <zenutil/service.h>

#include "diag/logging.h"

#include "compute/computeserver.h"

#include "storage/storageconfig.h"
#include "storage/zenstorageserver.h"

#include "hub/zenhubserver.h"

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

//////////////////////////////////////////////////////////////////////////
// We don't have any doctest code in this file but this is needed to bring
// in some shared code into the executable

#if ZEN_WITH_TESTS
#	define ZEN_TEST_WITH_RUNNER 1
#	include <zencore/testing.h>
#endif

namespace zen::utils {
std::atomic_uint32_t SignalCounter[NSIG] = {0};

static void
SignalCallbackHandler(int SigNum)
{
	if (SigNum >= 0 && SigNum < NSIG)
	{
		SignalCounter[SigNum].fetch_add(1);
	}
}
}  // namespace zen::utils

namespace zen {

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

#if ZEN_PLATFORM_WINDOWS

/** Windows Service wrapper for Zen servers
 *
 * This class wraps a Zen server main entry point (the Main template parameter)
 * into a Windows Service by implementing the WindowsService interface.
 *
 * The Main type needs to implement the virtual functions from the ZenServerMain
 * base class, which provides the actual server logic.
 */
template<class Main>
class ZenWindowsService : public WindowsService
{
public:
	ZenWindowsService(typename Main::Config& ServerOptions) : m_EntryPoint(ServerOptions) {}

	ZenWindowsService(const ZenWindowsService&) = delete;
	ZenWindowsService& operator=(const ZenWindowsService&) = delete;

	virtual int Run() override { return m_EntryPoint.Run(); }

private:
	Main m_EntryPoint;
};

#endif	// ZEN_PLATFORM_WINDOWS

}  // namespace zen

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

namespace zen {

/** Application main entry point template
 *
 * This function handles common application startup tasks while allowing
 * different server types to be plugged in via the Main template parameter.
 *
 * On Windows, this function also handles platform-specific service
 * installation and uninstallation.
 *
 * The Main type needs to implement the virtual functions from the ZenServerMain
 * base class, which provides the actual server logic.
 *
 * The Main type is also expected to provide the following members:
 *
 * typedef Config -- Server configuration type, derived from ZenServerConfig
 * typedef Configurator -- Server configuration handler type, implements ZenServerConfiguratorBase
 *
 */
template<class Main>
int
AppMain(int argc, char* argv[])
{
	using namespace std::literals;

	signal(SIGINT, utils::SignalCallbackHandler);
	signal(SIGTERM, utils::SignalCallbackHandler);

#if ZEN_PLATFORM_LINUX
	IgnoreChildSignals();
#endif

	try
	{
		typename Main::Config ServerOptions;

#if ZEN_WITH_TRACE
		TraceInit("zenserver");
		ServerOptions.HasTraceCommandlineOptions = GetTraceOptionsFromCommandline(/* out */ ServerOptions.TraceCmdLineOptions);
		if (ServerOptions.HasTraceCommandlineOptions)
		{
			TraceConfigure(ServerOptions.TraceCmdLineOptions);
		}
#endif	// ZEN_WITH_TRACE

		{
			ZEN_OTEL_SPAN("AppMain::Configure");

			typename Main::Configurator Configurator(ServerOptions);
			Configurator.Configure(argc, argv);
		}

		if (ServerOptions.Detach)
		{
#if ZEN_PLATFORM_LINUX | ZEN_PLATFORM_MAC
			// Detach ourselves from any parent process
			setsid();
#endif
		}

		LimitHardwareConcurrency(ServerOptions.CoreLimit);

		std::string_view DeleteReason;

		if (ServerOptions.IsCleanStart)
		{
			DeleteReason = "clean start requested"sv;
		}
		else if (!ServerOptions.BaseSnapshotDir.empty())
		{
			DeleteReason = "will initialize state from base snapshot"sv;
		}

		if (!DeleteReason.empty())
		{
			if (IsDir(ServerOptions.DataDir))
			{
				ZEN_CONSOLE_INFO("Deleting files from '{}' ({})", ServerOptions.DataDir, DeleteReason);
				DeleteDirectories(ServerOptions.DataDir);
			}
		}

		if (!IsDir(ServerOptions.DataDir))
		{
			ServerOptions.IsFirstRun = true;
			CreateDirectories(ServerOptions.DataDir);
		}

		if (!ServerOptions.BaseSnapshotDir.empty())
		{
			ZEN_CONSOLE_INFO("Copying snapshot from '{}' into '{}", ServerOptions.BaseSnapshotDir, ServerOptions.DataDir);
			CopyTree(ServerOptions.BaseSnapshotDir, ServerOptions.DataDir, {.EnableClone = true});
		}

		ZEN_MEMSCOPE(GetZenserverTag());

#if ZEN_PLATFORM_WINDOWS
		if (ServerOptions.InstallService)
		{
			WindowsService::Install();

			std::exit(0);
		}

		if (ServerOptions.UninstallService)
		{
			WindowsService::Delete();

			std::exit(0);
		}

		ZenWindowsService<Main> App(ServerOptions);
		return App.ServiceMain();
#else
		if (ServerOptions.InstallService || ServerOptions.UninstallService)
		{
			throw std::runtime_error("Service mode is not supported on this platform");
		}

		Main App(ServerOptions);
		return App.Run();
#endif
	}
	catch (const OptionParseException& ParseEx)
	{
		// The parsing error already outputs all the details so no need to output the command line here
		fprintf(stderr, ZEN_APP_NAME " ERROR: %s\n", ParseEx.what());
		return 1;
	}
	catch (const AssertException& AssertEx)
	{
		fprintf(stderr, ZEN_APP_NAME " ERROR: Caught assert exception in main: '%s'\n", AssertEx.FullDescription().c_str());
		return 1;
	}
	catch (const std::exception& Ex)
	{
		fprintf(stderr, ZEN_APP_NAME " ERROR: Caught exception in main: '%s'\n", Ex.what());

		return 1;
	}
}
}  // namespace zen

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

#if ZEN_WITH_TESTS
int
test_main(int argc, char** argv)
{
#	if ZEN_PLATFORM_WINDOWS
	setlocale(LC_ALL, "en_us.UTF8");
#	endif	// ZEN_PLATFORM_WINDOWS

	zen::logging::InitializeLogging();
	zen::logging::SetLogLevel(zen::logging::level::Debug);

	zen::MaximizeOpenFileCount();

	return ZEN_RUN_TESTS(argc, argv);
}
#endif

int
main(int argc, char* argv[])
{
#if ZEN_PLATFORM_WINDOWS
	setlocale(LC_ALL, "en_us.UTF8");
#endif	// ZEN_PLATFORM_WINDOWS

	zen::CommandLineConverter ArgConverter(argc, argv);

	using namespace zen;
	using namespace std::literals;

	auto _ = zen::MakeGuard([] {
		// Allow some time for worker threads to unravel, in an effort
		// to prevent shutdown races in TLS object destruction, mainly due to
		// threads which we don't directly control (Windows thread pool) and
		// therefore can't join.
		//
		// This isn't a great solution, but for now it seems to help reduce
		// shutdown crashes observed in some situations.
		WaitForThreads(1000);
	});

	enum
	{
		kHub,
		kStore,
		kCompute,
		kTest
	} ServerMode = kStore;

	if (argc >= 2)
	{
		if (argv[1] == "hub"sv)
		{
			ServerMode = kHub;
		}
		else if ((argv[1] == "store"sv) || (argv[1] == "storage"sv))
		{
			ServerMode = kStore;
		}
		else if (argv[1] == "compute"sv)
		{
			ServerMode = kCompute;
		}
		else if (argv[1] == "test"sv)
		{
			ServerMode = kTest;
		}
	}

	switch (ServerMode)
	{
		case kTest:
#if ZEN_WITH_TESTS
			return test_main(argc, argv);
#else
			fprintf(stderr, "test option not available in release mode!\n");
			exit(5);
#endif
			break;
		case kHub:
			return AppMain<ZenHubServerMain>(argc, argv);
		case kCompute:
#if ZEN_WITH_COMPUTE_SERVICES
			return AppMain<ZenComputeServerMain>(argc, argv);
#else
			fprintf(stderr, "compute services are not compiled in!\n");
			exit(5);
#endif
		default:
		case kStore:
			return AppMain<ZenStorageServerMain>(argc, argv);
	}
}