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

#include "trace_cmd.h"

#include "browser_launcher.h"
#include "consoleprogress.h"
#include "symbol_resolver.h"
#include "trace_analyze.h"
#include "trace_model.h"
#include "trace_viewer_service.h"
#include "zenserviceclient.h"

#include <zencore/except_fmt.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/string.h>
#include <zencore/thread.h>
#include <zencore/workthreadpool.h>
#include <zenhttp/httpclient.h>
#include <zenhttp/httpcommon.h>
#include <zenhttp/httpserver.h>

#include <filesystem>
#include <numeric>

using namespace std::literals;

namespace zen {

namespace {

#if ZEN_PLATFORM_WINDOWS
	constexpr const char* kSymbolBackendHelp = "Symbol backend: auto (default), pdb, dbghelp, llvm, off";
#elif ZEN_PLATFORM_MAC
	constexpr const char* kSymbolBackendHelp = "Symbol backend: auto (default - prefers llvm, falls back to atos), llvm, atos, off";
#else
	constexpr const char* kSymbolBackendHelp = "Symbol backend: auto (default - uses llvm), llvm, off";
#endif

}  // namespace

//////////////////////////////////////////////////////////////////////////
// TraceAnalyzeSubCmd

TraceAnalyzeSubCmd::TraceAnalyzeSubCmd() : ZenSubCmdBase("analyze", "Analyze a .utrace file")
{
	SubOptions().add_option("", "", "file", "Path to .utrace file", cxxopts::value(m_TraceFilePath), "<filepath>");
	SubOptions().add_option("",
							"",
							"live-allocs",
							"Dump top N live-allocation callstacks (0 = off, default 50)",
							cxxopts::value(m_LiveAllocs)->default_value("50"),
							"<count>");
	SubOptions().add_option("",
							"",
							"churn",
							"Dump top N allocation-churn callstacks (0 = off, default 0)",
							cxxopts::value(m_Churn)->default_value("0"),
							"<count>");
	SubOptions().add_option("",
							"",
							"churn-distance",
							"Max event distance between alloc and free to count as churn (default 1000)",
							cxxopts::value(m_ChurnMin)->default_value("1000"),
							"<events>");
	SubOptions().add_option("", "", "symbols", kSymbolBackendHelp, cxxopts::value(m_Symbols)->default_value("auto"), "<backend>");
	SubOptions().add_option("",
							"",
							"html-report",
							"Write a standalone HTML memory report (all live leaks + top 100 churn sites)",
							cxxopts::value(m_HtmlReportPath),
							"<filepath>");
	SubOptions().add_option("",
							"",
							"callstack-skip",
							"Semicolon-separated wildcard patterns for frames to hide from analyzed callstacks",
							cxxopts::value(m_CallstackSkip),
							"<pattern;...>");
	SubOptions().add_option("",
							"",
							"no-callstack-heuristic",
							"Disable leading third-party frame trimming in analyzed callstacks",
							cxxopts::value(m_NoCallstackHeuristic)->default_value("false"),
							"<no-callstack-heuristic>");
	SubOptions().add_option("",
							"",
							"no-cache",
							"Skip reading/writing the .ucache_z analysis cache",
							cxxopts::value(m_NoCache)->default_value("false"),
							"<no-cache>");
	SubOptions().parse_positional({"file"});
	SubOptions().positional_help("<file.utrace>");
}

void
TraceAnalyzeSubCmd::Run(const ZenCliOptions& GlobalOptions)
{
	ZEN_UNUSED(GlobalOptions);

	std::filesystem::path FilePath = trace_detail::ResolveTraceFile(m_TraceFilePath, SubOptions());

	trace_detail::AnalyzeOptions Options;
	Options.LiveAllocsLimit			 = m_LiveAllocs;
	Options.ChurnLimit				 = m_Churn;
	Options.ChurnDistanceThreshold	 = uint64_t(m_ChurnMin);
	Options.Symbols					 = trace_detail::ParseSymbolBackend(m_Symbols);
	Options.NoCache					 = m_NoCache;
	Options.EnableCallstackHeuristic = !m_NoCallstackHeuristic;
	Options.HtmlReportPath			 = m_HtmlReportPath;
	ForEachStrTok(m_CallstackSkip, ';', [&Options](std::string_view Pattern) {
		if (!Pattern.empty())
		{
			Options.CallstackSkipPatterns.emplace_back(Pattern);
		}
		return true;
	});
	trace_detail::RunAnalyze(FilePath, Options);
}

//////////////////////////////////////////////////////////////////////////
// TraceInspectSubCmd

TraceInspectSubCmd::TraceInspectSubCmd() : ZenSubCmdBase("inspect", "Inspect event schemas in a .utrace file")
{
	SubOptions().add_option("", "", "file", "Path to .utrace file", cxxopts::value(m_TraceFilePath), "<filepath>");
	SubOptions().parse_positional({"file"});
	SubOptions().positional_help("<file.utrace>");
}

void
TraceInspectSubCmd::Run(const ZenCliOptions& GlobalOptions)
{
	ZEN_UNUSED(GlobalOptions);

	std::filesystem::path FilePath = trace_detail::ResolveTraceFile(m_TraceFilePath, SubOptions());
	trace_detail::RunInspect(FilePath);
}

//////////////////////////////////////////////////////////////////////////
// TraceServeSubCmd

TraceServeSubCmd::TraceServeSubCmd() : ZenSubCmdBase("serve", "Serve an interactive viewer for a .utrace file")
{
	AddAlias("view");
	SubOptions().add_option("", "", "file", "Path to .utrace file", cxxopts::value(m_TraceFilePath), "<filepath>");
	SubOptions().add_option("", "p", "port", "Port to listen on", cxxopts::value(m_Port)->default_value("1480"), "<port>");
	SubOptions().add_option("", "", "bind", "Address to bind to", cxxopts::value(m_Bind)->default_value("127.0.0.1"), "<host>");
	SubOptions().add_option("", "", "symbols", kSymbolBackendHelp, cxxopts::value(m_Symbols)->default_value("auto"), "<backend>");
	SubOptions().add_option("",
							"",
							"no-browser",
							"Do not launch a web browser after starting the server",
							cxxopts::value(m_NoBrowser)->default_value("false"),
							"<no-browser>");
	SubOptions().parse_positional({"file"});
	SubOptions().positional_help("<file.utrace>");
}

void
TraceServeSubCmd::Run(const ZenCliOptions& GlobalOptions)
{
	ZEN_UNUSED(GlobalOptions);

	std::filesystem::path FilePath = trace_detail::ResolveTraceFile(m_TraceFilePath, SubOptions());

	WorkerThreadPool ThreadPool(gsl::narrow<int>(GetHardwareConcurrency()));

	uint64_t FileSize = uint64_t(std::filesystem::file_size(FilePath));
	ZEN_CONSOLE("Parsing {} ({})", FilePath.filename().string(), zen::NiceBytes(FileSize));

	std::unique_ptr<ProgressBase>			   ProgressOwner(CreateConsoleProgress(ConsoleProgressMode::Pretty));
	std::unique_ptr<ProgressBase::ProgressBar> Progress = ProgressOwner->CreateProgressBar("Parse");
	trace_detail::TraceModel				   Model =
		trace_detail::BuildTraceModel(FilePath, ThreadPool, [&](uint64_t BytesProcessed, uint64_t TotalBytes, uint64_t EventsSoFar) {
			Progress->UpdateState(
				{
					.Task			= "Parsing trace",
					.Details		= fmt::format("{} events", zen::ThousandsNum(EventsSoFar)),
					.TotalCount		= TotalBytes,
					.RemainingCount = TotalBytes - std::min(BytesProcessed, TotalBytes),
				},
				false);
		});
	Progress->Finish();

	ZEN_CONSOLE("  Events:   {}", zen::ThousandsNum(Model.TotalEvents));
	ZEN_CONSOLE("  Threads:  {}", Model.Threads.size());
	ZEN_CONSOLE(
		"  Scopes:   {}",
		zen::ThousandsNum(std::accumulate(Model.Timelines.begin(),
										  Model.Timelines.end(),
										  size_t(0),
										  [](size_t Acc, const trace_detail::ThreadTimeline& T) { return Acc + T.Scopes.size(); })));
	ZEN_CONSOLE("  Time:     {}", zen::NiceTimeSpanMs(Model.ParseTimeMs));

	std::unique_ptr<trace_detail::SymbolResolver> Symbols = trace_detail::CreateSymbolResolver(trace_detail::ParseSymbolBackend(m_Symbols));
	for (const trace_detail::ModuleInfo& Mod : Model.Modules)
	{
		Symbols->LoadModule(Mod);
	}
	ZEN_CONSOLE("  Symbols:  {} modules loaded", Model.Modules.size());
	ZEN_CONSOLE("");

	HttpServerConfig Config;
	Config.ServerClass		 = "asio";
	Config.IsDedicatedServer = false;
	Config.AllowPortProbing	 = true;
	Config.ForceLoopback	 = (m_Bind == "127.0.0.1" || m_Bind == "localhost" || m_Bind == "::1");

	Ref<HttpServer> Server = CreateHttpServer(Config);

	std::filesystem::path TempDir = std::filesystem::temp_directory_path() / "zen-trace-viewer";
	std::error_code		  Ec;
	std::filesystem::create_directories(TempDir, Ec);

	int EffectivePort = Server->Initialize(m_Port, TempDir);
	if (EffectivePort <= 0)
	{
		throw zen::runtime_error("Failed to initialize HTTP server");
	}

	TraceViewerService ViewerService(Model, std::move(Symbols));
	Server->RegisterService(ViewerService);

	std::string Url = fmt::format("http://{}:{}{}", m_Bind, EffectivePort, ViewerService.BaseUri());
	ZEN_CONSOLE("Serving trace viewer at {}", Url);
	ZEN_CONSOLE("Press Ctrl+C to stop");

	if (!m_NoBrowser)
	{
		try
		{
			LaunchBrowser(Url);
		}
		catch (const std::exception& E)
		{
			ZEN_WARN("Failed to launch browser: {}", E.what());
		}
	}

	Server->Run(/*IsInteractiveSession=*/true);
	Server->Close();
}

//////////////////////////////////////////////////////////////////////////
// TraceTrimSubCmd

TraceTrimSubCmd::TraceTrimSubCmd() : ZenSubCmdBase("trim", "Trim a .utrace file to a time range while preserving important events")
{
	SubOptions().add_option("", "", "file", "Path to input .utrace file", cxxopts::value(m_TraceFilePath), "<filepath>");
	SubOptions().add_option("", "o", "output", "Path to output .utrace file", cxxopts::value(m_OutputPath), "<filepath>");
	SubOptions().add_option("",
							"",
							"start",
							"Start of the time window in seconds from the beginning of the trace",
							cxxopts::value(m_StartSec)->default_value("0"),
							"<seconds>");
	SubOptions().add_option("",
							"",
							"end",
							"End of the time window in seconds from the beginning of the trace",
							cxxopts::value(m_EndSec)->default_value("0"),
							"<seconds>");
	SubOptions().parse_positional({"file"});
	SubOptions().positional_help("<file.utrace>");
}

void
TraceTrimSubCmd::Run(const ZenCliOptions& GlobalOptions)
{
	ZEN_UNUSED(GlobalOptions);

	std::filesystem::path InputPath = trace_detail::ResolveTraceFile(m_TraceFilePath, SubOptions());

	if (m_OutputPath.empty())
	{
		throw zen::OptionParseException("--output is required", SubOptions().help());
	}
	if (m_EndSec <= m_StartSec)
	{
		throw zen::OptionParseException("--end must be greater than --start", SubOptions().help());
	}

	trace_detail::TraceTrimArgs Args;
	Args.InputPath	= InputPath;
	Args.OutputPath = std::filesystem::absolute(m_OutputPath);
	Args.StartSec	= m_StartSec;
	Args.EndSec		= m_EndSec;

	trace_detail::RunTraceTrim(Args);
}

//////////////////////////////////////////////////////////////////////////
// TraceStartSubCmd

TraceStartSubCmd::TraceStartSubCmd() : ZenSubCmdBase("start", "Start zen server realtime tracing")
{
	SubOptions().add_option("", "u", "hosturl", ZenCmdBase::kHostUrlHelp, cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	SubOptions().add_option("", "", "host", "Stream trace data to a remote host", cxxopts::value(m_TraceHost), "<hostip>");
	SubOptions().add_option("", "", "file", "Write trace data to a file", cxxopts::value(m_TraceFile), "<filepath>");
}

void
TraceStartSubCmd::Run(const ZenCliOptions& GlobalOptions)
{
	ZEN_UNUSED(GlobalOptions);

	if (m_TraceHost.empty() && m_TraceFile.empty())
	{
		throw OptionParseException("Either --host or --file is required", SubOptions().help());
	}
	if (!m_TraceHost.empty() && !m_TraceFile.empty())
	{
		throw OptionParseException("--host and --file are mutually exclusive", SubOptions().help());
	}

	std::string StartArg = m_TraceHost.empty() ? fmt::format("file={}", m_TraceFile) : fmt::format("host={}", m_TraceHost);

	ZenServiceClient Service({.HostSpec = m_HostName, .CommandName = "start"});
	HttpClient&		 Http = Service.Http();
	if (HttpClient::Response Response = Http.Post(fmt::format("/admin/trace/start?{}"sv, StartArg)))
	{
		ZEN_CONSOLE("OK: {}", Response.ToText());
	}
	else
	{
		Response.ThrowError("Trace start failed");
	}
}

//////////////////////////////////////////////////////////////////////////
// TraceStopSubCmd

TraceStopSubCmd::TraceStopSubCmd() : ZenSubCmdBase("stop", "Stop zen server realtime tracing")
{
	SubOptions().add_option("", "u", "hosturl", ZenCmdBase::kHostUrlHelp, cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
}

void
TraceStopSubCmd::Run(const ZenCliOptions& GlobalOptions)
{
	ZEN_UNUSED(GlobalOptions);

	ZenServiceClient Service({.HostSpec = m_HostName, .CommandName = "stop"});
	HttpClient&		 Http = Service.Http();
	if (HttpClient::Response Response = Http.Post("/admin/trace/stop"sv))
	{
		ZEN_CONSOLE("OK: {}", Response.ToText());
	}
	else
	{
		Response.ThrowError("Trace stop failed");
	}
}

//////////////////////////////////////////////////////////////////////////
// TraceStatusSubCmd

TraceStatusSubCmd::TraceStatusSubCmd() : ZenSubCmdBase("status", "Report zen server realtime tracing status")
{
	SubOptions().add_option("", "u", "hosturl", ZenCmdBase::kHostUrlHelp, cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
}

void
TraceStatusSubCmd::Run(const ZenCliOptions& GlobalOptions)
{
	ZEN_UNUSED(GlobalOptions);

	ZenServiceClient Service({.HostSpec = m_HostName, .CommandName = "status"});
	HttpClient&		 Http = Service.Http();
	if (HttpClient::Response Response = Http.Get("/admin/trace"sv))
	{
		ZEN_CONSOLE("OK: {}", Response.ToText());
	}
	else
	{
		Response.ThrowError("Trace status failed");
	}
}

//////////////////////////////////////////////////////////////////////////
// TraceCommand

TraceCommand::TraceCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("__hidden__", "", "subcommand", "", cxxopts::value<std::string>(m_SubCommand)->default_value(""), "");
	m_Options.parse_positional({"subcommand"});

	AddSubCommand(m_AnalyzeSubCmd);
	AddSubCommand(m_InspectSubCmd);
	AddSubCommand(m_ServeSubCmd);
	AddSubCommand(m_TrimSubCmd);
	AddSubCommand(m_StartSubCmd);
	AddSubCommand(m_StopSubCmd);
	AddSubCommand(m_StatusSubCmd);
}

TraceCommand::~TraceCommand() = default;

}  // namespace zen