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

// Zen command line client utility
//

#include "zen.h"

#include "chunk/chunk.h"
#include "cmds/cache.h"
#include "cmds/copy.h"
#include "cmds/dedup.h"
#include "cmds/hash.h"
#include "cmds/print.h"
#include "cmds/projectstore.h"
#include "cmds/rpcreplay.h"
#include "cmds/scrub.h"
#include "cmds/status.h"
#include "cmds/top.h"
#include "cmds/up.h"
#include "cmds/version.h"

#include <zencore/filesystem.h>
#include <zencore/logging.h>
#include <zencore/scopeguard.h>
#include <zencore/string.h>
#include <zenhttp/httpcommon.h>
#include <zenutil/zenserverprocess.h>

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

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

#if ZEN_USE_MIMALLOC
#	include <mimalloc-new-delete.h>
#endif

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

bool
ZenCmdBase::ParseOptions(int argc, char** argv)
{
	cxxopts::Options&	 CmdOptions = Options();
	cxxopts::ParseResult Result		= CmdOptions.parse(argc, argv);

	if (Result.count("help"))
	{
		printf("%s\n", CmdOptions.help().c_str());
		return false;
	}

	if (!Result.unmatched().empty())
	{
		zen::ExtendableStringBuilder<64> StringBuilder;
		for (bool First = true; const auto& Param : Result.unmatched())
		{
			if (!First)
			{
				StringBuilder.Append(", ");
			}
			StringBuilder.Append('"');
			StringBuilder.Append(Param);
			StringBuilder.Append('"');
			First = false;
		}

		throw zen::OptionParseException(fmt::format("Invalid arguments: {}", StringBuilder.ToView()));
	}

	return true;
}

std::string
ZenCmdBase::FormatHttpResponse(const cpr::Response& Response)
{
	if (Response.error.code != cpr::ErrorCode::OK)
	{
		if (Response.error.message.empty())
		{
			return fmt::format("Request '{}' failed, error code {}", Response.url.str(), static_cast<int>(Response.error.code));
		}
		return fmt::format("Request '{}' failed. Reason: '{}' ({})",
						   Response.url.str(),
						   Response.error.message,
						   static_cast<int>(Response.error.code));
	}

	std::string Content;
	if (auto It = Response.header.find("Content-Type"); It != Response.header.end())
	{
		zen::HttpContentType ContentType = zen::ParseContentType(It->second);
		if (ContentType == zen::HttpContentType::kText)
		{
			Content = fmt::format("'{}'", Response.text);
		}
		else if (ContentType == zen::HttpContentType::kJSON)
		{
			Content = fmt::format("\n{}", Response.text);
		}
		else if (!Response.text.empty())
		{
			Content = fmt::format("[{}]", MapContentTypeToString(ContentType));
		}
	}

	std::string_view ResponseString = zen::ReasonStringForHttpResultCode(
		Response.status_code == static_cast<long>(zen::HttpResponseCode::NoContent) ? static_cast<long>(zen::HttpResponseCode::OK)
																					: Response.status_code);
	if (Content.empty())
	{
		return std::string(ResponseString);
	}

	return fmt::format("{}: {}", ResponseString, Content);
}

int
ZenCmdBase::MapHttpToCommandReturnCode(const cpr::Response& Response)
{
	if (zen::IsHttpSuccessCode(Response.status_code))
	{
		return 0;
	}
	if (Response.error.code != cpr::ErrorCode::OK)
	{
		return static_cast<int>(Response.error.code);
	}
	return 1;
}

#if ZEN_WITH_TESTS

class RunTestsCommand : public ZenCmdBase
{
public:
	virtual int Run(const ZenCliOptions& GlobalOptions, int argc, char** argv) override
	{
		ZEN_UNUSED(GlobalOptions);

		// Set output mode to handle virtual terminal sequences
		zen::logging::EnableVTMode();

		return ZEN_RUN_TESTS(argc, argv);
	}

	virtual cxxopts::Options& Options() override { return m_Options; }

private:
	cxxopts::Options m_Options{"runtests", "Run tests"};
};

#endif

//////////////////////////////////////////////////////////////////////////
// TODO: should make this Unicode-aware so we can pass anything in on the
// command line.

int
main(int argc, char** argv)
{
	using namespace zen;

#if ZEN_USE_MIMALLOC
	mi_version();
#endif

	zen::logging::InitializeLogging();
	zen::MaximizeOpenFileCount();

	// Set output mode to handle virtual terminal sequences
	zen::logging::EnableVTMode();

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

	auto _ = zen::MakeGuard([] { spdlog::shutdown(); });

	CacheInfoCommand		 CacheInfoCmd;
	CopyCommand				 CopyCmd;
	CreateOplogCommand		 CreateOplogCmd;
	CreateProjectCommand	 CreateProjectCmd;
	DedupCommand			 DedupCmd;
	DownCommand				 DownCmd;
	DropCommand				 DropCmd;
	DropProjectCommand		 ProjectDropCmd;
	ExportOplogCommand		 ExportOplogCmd;
	GcCommand				 GcCmd;
	GcStatusCommand			 GcStatusCmd;
	HashCommand				 HashCmd;
	ImportOplogCommand		 ImportOplogCmd;
	PrintCommand			 PrintCmd;
	PrintPackageCommand		 PrintPkgCmd;
	ProjectInfoCommand		 ProjectInfoCmd;
	PsCommand				 PsCmd;
	RpcReplayCommand		 RpcReplayCmd;
	RpcStartRecordingCommand RpcStartRecordingCmd;
	RpcStopRecordingCommand	 RpcStopRecordingCmd;
	ScrubCommand			 ScrubCmd;
	StatusCommand			 StatusCmd;
	TopCommand				 TopCmd;
	UpCommand				 UpCmd;
	VersionCommand			 VersionCmd;
	CacheStatsCommand		 CacheStatsCmd;
	CacheDetailsCommand		 CacheDetailsCmd;
	ProjectStatsCommand		 ProjectStatsCmd;
	ProjectDetailsCommand	 ProjectDetailsCmd;
#if ZEN_WITH_TESTS
	RunTestsCommand RunTestsCmd;
#endif

	const struct CommandInfo
	{
		const char* CmdName;
		ZenCmdBase* Cmd;
		const char* CmdSummary;
	} Commands[] = {
		// clang-format off
//		{"chunk",				&ChunkCmd,				"Perform chunking"},
		{"cache-info",			&CacheInfoCmd,			"Info on cache, namespace or bucket"},
		{"copy",				&CopyCmd,				"Copy file(s)"},
		{"dedup",				&DedupCmd,				"Dedup files"},
		{"down",				&DownCmd,				"Bring zen server down"},
		{"drop",				&DropCmd,				"Drop cache namespace or bucket"},
		{"gc-status",			&GcStatusCmd,			"Garbage collect zen storage status check"},
		{"gc",					&GcCmd,					"Garbage collect zen storage"},
		{"hash",				&HashCmd,				"Compute file hashes"},
		{"oplog-create",		&CreateOplogCmd,		"Create a project oplog"},
		{"oplog-export",		&ExportOplogCmd,		"Export project store oplog"},
		{"oplog-import",		&ImportOplogCmd,		"Import project store oplog"},
		{"print",				&PrintCmd,				"Print compact binary object"},
		{"printpackage",		&PrintPkgCmd,			"Print compact binary package"},
		{"project-create",		&CreateProjectCmd,		"Create a project"},
		{"project-drop",		&ProjectDropCmd,		"Drop project or project oplog"},
		{"project-info",		&ProjectInfoCmd,		"Info on project or project oplog"},
		{"ps",					&PsCmd,					"Enumerate running zen server instances"},
		{"rpc-record-replay",	&RpcReplayCmd,			"Stops recording of cache rpc requests on a host"},
		{"rpc-record-start",	&RpcStartRecordingCmd,	"Replays a previously recorded session of rpc requests"},
		{"rpc-record-stop",		&RpcStopRecordingCmd,	"Starts recording of cache rpc requests on a host"},
		{"scrub",				&ScrubCmd,				"Scrub zen storage (verify data integrity)"},
		{"status",				&StatusCmd,				"Show zen status"},
		{"top",					&TopCmd,				"Monitor zen server activity"},
		{"up",					&UpCmd,					"Bring zen server up"},
		{"version",				&VersionCmd,			"Get zen server version"},
		{"cache-stats",			&CacheStatsCmd,			"Stats on cache"},
		{"cache-details",		&CacheDetailsCmd,		"Details on cache"},
		{"project-stats",		&ProjectStatsCmd,		"Stats on project store"},
		{"project-details",		&ProjectDetailsCmd,		"Details on project store"},
#if ZEN_WITH_TESTS
		{"runtests",			&RunTestsCmd,			"Run zen tests"},
#endif
		// clang-format on
	};

	// Build set containing available commands

	std::unordered_set<std::string> CommandSet;

	for (const auto& Cmd : Commands)
		CommandSet.insert(Cmd.CmdName);

	// Split command line into options, commands and any pass-through arguments

	std::string				 Passthrough;
	std::vector<std::string> PassthroughV;

	for (int i = 1; i < argc; ++i)
	{
		if (strcmp(argv[i], "--") == 0)
		{
			bool							  IsFirst = true;
			zen::ExtendableStringBuilder<256> Line;

			for (int j = i + 1; j < argc; ++j)
			{
				if (!IsFirst)
				{
					Line.AppendAscii(" ");
				}

				std::string_view ThisArg(argv[j]);
				PassthroughV.push_back(std::string(ThisArg));

				const bool NeedsQuotes = (ThisArg.find(' ') != std::string_view::npos);

				if (NeedsQuotes)
				{
					Line.AppendAscii("\"");
				}

				Line.Append(ThisArg);

				if (NeedsQuotes)
				{
					Line.AppendAscii("\"");
				}

				IsFirst = false;
			}

			Passthrough = Line.c_str();

			// This will "truncate" the arg vector and terminate the loop
			argc = i - 1;
		}
	}

	// Split command line into global vs command options. We do this by simply
	// scanning argv for a string we recognise as a command and split it there

	std::vector<char*> CommandArgVec;
	CommandArgVec.push_back(argv[0]);

	for (int i = 1; i < argc; ++i)
	{
		if (CommandSet.find(argv[i]) != CommandSet.end())
		{
			int commandArgCount = /* exec name */ 1 + argc - (i + 1);
			CommandArgVec.resize(commandArgCount);
			std::copy(argv + i + 1, argv + argc, CommandArgVec.begin() + 1);

			argc = i + 1;

			break;
		}
	}

	// Parse global CLI arguments

	ZenCliOptions GlobalOptions;

	GlobalOptions.PassthroughArgs = Passthrough;
	GlobalOptions.PassthroughV	  = PassthroughV;

	std::string SubCommand = "<None>";

	cxxopts::Options Options("zen", "Zen management tool");

	Options.add_options()("d, debug", "Enable debugging", cxxopts::value<bool>(GlobalOptions.IsDebug));
	Options.add_options()("v, verbose", "Enable verbose logging", cxxopts::value<bool>(GlobalOptions.IsVerbose));
	Options.add_options()("help", "Show command line help");
	Options.add_options()("c, command", "Sub command", cxxopts::value<std::string>(SubCommand));

	Options.parse_positional({"command"});

	const bool IsNullInvoke = (argc == 1);	// If no arguments are passed we want to print usage information

	try
	{
		cxxopts::ParseResult ParseResult = Options.parse(argc, argv);

		if (ParseResult.count("help") || IsNullInvoke == 1)
		{
			std::string Help = Options.help();

			printf("%s\n", Help.c_str());

			printf("available commands:\n");

			for (const CommandInfo& CmdInfo : Commands)
			{
				printf("  %-20s %s\n", CmdInfo.CmdName, CmdInfo.CmdSummary);
			}

			exit(0);
		}

		if (GlobalOptions.IsDebug)
		{
			spdlog::set_level(spdlog::level::debug);
		}

		for (const CommandInfo& CmdInfo : Commands)
		{
			if (StrCaseCompare(SubCommand.c_str(), CmdInfo.CmdName) == 0)
			{
				cxxopts::Options& VerbOptions = CmdInfo.Cmd->Options();

				try
				{
					return CmdInfo.Cmd->Run(GlobalOptions, (int)CommandArgVec.size(), CommandArgVec.data());
				}
				catch (cxxopts::OptionParseException& Ex)
				{
					std::string help = VerbOptions.help();

					printf("Error parsing arguments for command '%s': %s\n\n%s", SubCommand.c_str(), Ex.what(), help.c_str());

					exit(11);
				}
				catch (OptionParseException& Ex)
				{
					std::string help = VerbOptions.help();

					printf("Error parsing arguments for command '%s': %s\n\n%s", SubCommand.c_str(), Ex.what(), help.c_str());

					exit(11);
				}
			}
		}

		printf("Unknown command specified: '%s', exiting\n", SubCommand.c_str());
	}
	catch (cxxopts::OptionParseException& Ex)
	{
		std::string HelpMessage = Options.help();

		printf("Error parsing program arguments: %s\n\n%s", Ex.what(), HelpMessage.c_str());

		return 9;
	}
	catch (OptionParseException& Ex)
	{
		std::string HelpMessage = Options.help();

		printf("Error parsing program arguments: %s\n\n%s", Ex.what(), HelpMessage.c_str());

		return 9;
	}
	catch (std::exception& Ex)
	{
		printf("Exception caught from 'main': %s\n", Ex.what());

		return 10;
	}

	return 0;
}

std::string
ZenCmdBase::ResolveTargetHostSpec(const std::string& InHostSpec)
{
	if (InHostSpec.empty())
	{
		zen::ZenServerState Servers;

		if (Servers.InitializeReadOnly())
		{
			std::string ResolvedSpec = InHostSpec;

			Servers.Snapshot([&](const zen::ZenServerState::ZenServerEntry& Entry) {
				if (ResolvedSpec.empty())
				{
					ResolvedSpec = fmt::format("http://localhost:{}", Entry.EffectiveListenPort);
				}
			});

			return ResolvedSpec;
		}
	}

	return InHostSpec;
}