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

#include "run_cmd.h"

#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/process.h>
#include <zencore/string.h>
#include <zencore/timer.h>

using namespace std::literals;

#define ZEN_COLOR_BLACK	  "\033[0;30m"
#define ZEN_COLOR_RED	  "\033[0;31m"
#define ZEN_COLOR_GREEN	  "\033[0;32m"
#define ZEN_COLOR_YELLOW  "\033[0;33m"
#define ZEN_COLOR_BLUE	  "\033[0;34m"
#define ZEN_COLOR_MAGENTA "\033[0;35m"
#define ZEN_COLOR_CYAN	  "\033[0;36m"
#define ZEN_COLOR_WHITE	  "\033[0;37m"

#define ZEN_BRIGHT_COLOR_BLACK	 "\033[1;30m"
#define ZEN_BRIGHT_COLOR_RED	 "\033[1;31m"
#define ZEN_BRIGHT_COLOR_GREEN	 "\033[1;32m"
#define ZEN_BRIGHT_COLOR_YELLOW	 "\033[1;33m"
#define ZEN_BRIGHT_COLOR_BLUE	 "\033[1;34m"
#define ZEN_BRIGHT_COLOR_MAGENTA "\033[1;35m"
#define ZEN_BRIGHT_COLOR_CYAN	 "\033[1;36m"
#define ZEN_BRIGHT_COLOR_WHITE	 "\033[1;37m"

#define ZEN_COLOR_RESET "\033[0m"

namespace zen {

RunCommand::RunCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "n", "count", "Number of times to run command", cxxopts::value(m_RunCount), "<count>");
	m_Options.add_option("", "t", "time", "How long to run command(s) for", cxxopts::value(m_RunTime), "<seconds>");
	m_Options.add_option("",
						 "",
						 "basepath",
						 "Where to run command. Each run will execute in its own numbered subdirectory below this directory. Additionally, "
						 "stdout will be redirected to a file in the provided directory",
						 cxxopts::value(m_BaseDirectory),
						 "<path>");
	m_Options.add_option("",
						 "",
						 "max-dirs",
						 "Number of base directories to retain on rotation",
						 cxxopts::value(m_MaxBaseDirectoryCount),
						 "<count>");
}

RunCommand::~RunCommand()
{
}

void
RunCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	if (!ParseOptions(argc, argv))
	{
		return;
	}

	// Validate arguments

	if (GlobalOptions.PassthroughArgV.empty() || GlobalOptions.PassthroughArgV[0].empty())
		throw OptionParseException("No command specified. The command to run is passed in after a double dash ('--') on the command line",
								   m_Options.help());

	if (m_RunCount < 0)
		throw OptionParseException(fmt::format("'--count' ('{}') is invalid", m_RunCount), m_Options.help());

	if (m_RunTime < -1 || m_RunTime == 0)
		throw OptionParseException(fmt::format("'--time' ('{}') is invalid", m_RunTime), m_Options.help());

	if (m_MaxBaseDirectoryCount < 0)
		throw OptionParseException(fmt::format("'--max-dirs' ('{}') is invalid", m_MaxBaseDirectoryCount), m_Options.help());

	if (m_RunTime > 0 && m_RunCount > 0)
		throw OptionParseException(fmt::format("'--time' ('{}')  conflicts with  '--count' ('{}') ", m_RunTime, m_RunCount),
								   m_Options.help());

	if (m_RunCount == 0)
		m_RunCount = 1;

	std::filesystem::path BaseDirectory;

	if (!m_BaseDirectory.empty())
	{
		BaseDirectory = m_BaseDirectory;

		if (m_MaxBaseDirectoryCount)
		{
			RotateDirectories(BaseDirectory, m_MaxBaseDirectoryCount);
			CreateDirectories(BaseDirectory);
		}
		else
		{
			CleanDirectory(BaseDirectory, /*ForceRemoveReadOnlyFiles*/ false);
		}
	}

	bool TimedRun			 = false;
	auto CommandDeadlineTime = std::chrono::system_clock::now();

	if (m_RunTime > 0)
	{
		m_RunCount = 1'000'000;
		TimedRun   = true;

		CommandDeadlineTime += std::chrono::seconds(m_RunTime);
	}

	struct RunResults
	{
		int										ExitCode = 0;
		std::chrono::duration<long, std::milli> Duration{};
	};

	std::vector<RunResults> Results;
	int						ErrorCount = 0;

	std::filesystem::path ExecutablePath   = SearchPathForExecutable(GlobalOptions.PassthroughArgV[0]);
	std::string			  CommandArguments = GlobalOptions.PassthroughArgs;

	for (int i = 0; i < m_RunCount; ++i)
	{
		std::filesystem::path RunDir;
		if (!BaseDirectory.empty())
		{
			RunDir = BaseDirectory / IntNum(i + 1).c_str();
			CreateDirectories(RunDir);
		}

		Stopwatch Timer;

		CreateProcOptions ProcOptions;

		if (!RunDir.empty())
		{
			ProcOptions.WorkingDirectory = &RunDir;
			ProcOptions.StdoutFile		 = RunDir / "stdout.txt";
		}

		fmt::print(ZEN_BRIGHT_COLOR_WHITE "run #{}" ZEN_COLOR_RESET ": {}\n", i + 1, GlobalOptions.PassthroughCommandLine);

		ProcessHandle Proc;
		Proc.Initialize(CreateProc(ExecutablePath, GlobalOptions.PassthroughCommandLine, ProcOptions));
		if (!Proc.IsValid())
		{
			throw std::runtime_error(fmt::format("failed to launch '{}'", ExecutablePath));
		}

		int ExitCode = Proc.WaitExitCode();

		auto EndTime = std::chrono::system_clock::now();

		if (ExitCode)
			++ErrorCount;

		Results.emplace_back(RunResults{.ExitCode = ExitCode, .Duration = std::chrono::milliseconds(Timer.GetElapsedTimeMs())});

		if (TimedRun)
		{
			if (EndTime >= CommandDeadlineTime)
			{
				m_RunCount = i + 1;
				break;
			}
		}
	}

	fmt::print("{:>5} {:>3} {:>6}\n", "run", "rc", "time");
	int i = 0;
	for (const auto& Entry : Results)
	{
		fmt::print("{:5} {:3} {:>6}\n", ++i, Entry.ExitCode, NiceTimeSpanMs(Entry.Duration.count()));
	}

	if (ErrorCount)
	{
		fmt::print("run complete ({}/{} failed)\n", ErrorCount, m_RunCount);
	}
	else
	{
		fmt::print("run complete, no error exit code\n", m_RunCount);
	}
}

}  // namespace zen