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

#include <zencore/options.h>

#include <zencore/string.h>

#include <algorithm>
#include <charconv>
#include <map>
#include <sstream>
#include <stdexcept>

namespace zen {

//////////////////////////////////////////////////////////////////////////
// OptionValueBase

std::shared_ptr<OptionValueBase>
OptionValueBase::DefaultValue(const std::string& Val)
{
	m_HasDefault   = true;
	m_DefaultValue = Val;
	Parse(Val);
	return shared_from_this();
}

std::shared_ptr<OptionValueBase>
OptionValueBase::ImplicitValue(const std::string& Val)
{
	m_HasImplicit	= true;
	m_ImplicitValue = Val;
	return shared_from_this();
}

//////////////////////////////////////////////////////////////////////////
// OptionValue<bool>

void
OptionValue<bool>::Parse(const std::string& Text)
{
	if (Text == "true" || Text == "1" || Text == "yes" || Text.empty())
	{
		m_Ref = true;
	}
	else if (Text == "false" || Text == "0" || Text == "no")
	{
		m_Ref = false;
	}
	else
	{
		throw std::runtime_error("Invalid boolean value: " + Text);
	}
}

//////////////////////////////////////////////////////////////////////////
// OptionValue<T> parse implementations

template<>
void
OptionValue<std::string>::Parse(const std::string& Text)
{
	m_Ref = Text;
}

template<>
void
OptionValue<std::filesystem::path>::Parse(const std::string& Text)
{
	std::string_view View = Text;

	// Strip surrounding quotes
	if (View.length() > 2 && View[0] == '"' && View[View.length() - 1] == '"')
	{
		View = View.substr(1, View.length() - 2);
	}

	// Strip trailing separators
	if (View.ends_with('/') || View.ends_with('\\') || View.ends_with(std::filesystem::path::preferred_separator))
	{
		View = View.substr(0, View.length() - 1);
	}

	m_Ref = std::filesystem::path(View).make_preferred();
}

template<>
void
OptionValue<int>::Parse(const std::string& Text)
{
	auto Result = ParseInt<int>(Text);
	if (!Result.has_value())
		throw std::runtime_error("Failed to parse integer: " + Text);
	m_Ref = Result.value();
}

template<>
void
OptionValue<unsigned int>::Parse(const std::string& Text)
{
	auto Result = ParseInt<unsigned int>(Text);
	if (!Result.has_value())
		throw std::runtime_error("Failed to parse unsigned integer: " + Text);
	m_Ref = Result.value();
}

template<>
void
OptionValue<uint8_t>::Parse(const std::string& Text)
{
	auto Result = ParseInt<unsigned int>(Text);
	if (!Result.has_value() || Result.value() > 255)
		throw std::runtime_error("Failed to parse uint8_t: " + Text);
	m_Ref = static_cast<uint8_t>(Result.value());
}

template<>
void
OptionValue<uint16_t>::Parse(const std::string& Text)
{
	auto Result = ParseInt<uint16_t>(Text);
	if (!Result.has_value())
		throw std::runtime_error("Failed to parse uint16_t: " + Text);
	m_Ref = Result.value();
}

// Note: uint32_t specialization is omitted because uint32_t is unsigned int
// on all supported platforms, making it a duplicate of the above.

template<>
void
OptionValue<uint64_t>::Parse(const std::string& Text)
{
	auto Result = ParseInt<uint64_t>(Text);
	if (!Result.has_value())
		throw std::runtime_error("Failed to parse uint64_t: " + Text);
	m_Ref = Result.value();
}

void
OptionValue<std::vector<std::string>>::Parse(const std::string& Text)
{
	// Each call to Parse appends a value to the vector.
	// Values containing commas are split into multiple entries.
	size_t Start = 0;
	while (Start < Text.size())
	{
		size_t Comma = Text.find(',', Start);
		if (Comma == std::string::npos)
		{
			m_Ref.push_back(Text.substr(Start));
			break;
		}
		m_Ref.push_back(Text.substr(Start, Comma - Start));
		Start = Comma + 1;
	}
}

//////////////////////////////////////////////////////////////////////////
// ParseResult

size_t
ParseResult::Count(std::string_view Name) const
{
	for (const auto& Entry : m_Counts)
	{
		if (Entry.first == Name)
			return Entry.second;
	}
	return 0;
}

const std::vector<std::string>&
ParseResult::Unmatched() const
{
	return m_Unmatched;
}

void
ParseResult::SetCount(const std::string& Name)
{
	for (auto& Entry : m_Counts)
	{
		if (Entry.first == Name)
		{
			Entry.second++;
			return;
		}
	}
	m_Counts.emplace_back(Name, 1);
}

void
ParseResult::AddUnmatched(const std::string& Arg)
{
	m_Unmatched.push_back(Arg);
}

//////////////////////////////////////////////////////////////////////////
// OptionAdder

OptionAdder::OptionAdder(Options& Owner, std::string Group) : m_Owner(Owner), m_Group(std::move(Group))
{
}

OptionAdder&
OptionAdder::operator()(std::string_view Spec, std::string_view Desc, std::shared_ptr<OptionValueBase> Value, std::string_view ArgHelp)
{
	std::string ShortOpt;
	std::string LongOpt;

	// Parse spec: "s,long" or "long"
	auto CommaPos = Spec.find(',');
	if (CommaPos != std::string_view::npos)
	{
		ShortOpt = std::string(Spec.substr(0, CommaPos));
		LongOpt	 = std::string(Spec.substr(CommaPos + 1));

		// Trim whitespace from both parts
		while (!ShortOpt.empty() && ShortOpt.front() == ' ')
			ShortOpt.erase(ShortOpt.begin());
		while (!ShortOpt.empty() && ShortOpt.back() == ' ')
			ShortOpt.pop_back();
		while (!LongOpt.empty() && LongOpt.front() == ' ')
			LongOpt.erase(LongOpt.begin());
		while (!LongOpt.empty() && LongOpt.back() == ' ')
			LongOpt.pop_back();
	}
	else
	{
		LongOpt = std::string(Spec);
	}

	m_Owner.AddOption(m_Group, ShortOpt, LongOpt, Desc, std::move(Value), ArgHelp);
	return *this;
}

//////////////////////////////////////////////////////////////////////////
// Options

Options::Options(std::string_view Program, std::string_view Description) : m_Program(Program), m_Description(Description)
{
}

OptionAdder
Options::AddOptions(std::string_view Group)
{
	return OptionAdder(*this, std::string(Group));
}

void
Options::AddOption(std::string_view					Group,
				   std::string_view					ShortOpt,
				   std::string_view					LongOpt,
				   std::string_view					Desc,
				   std::shared_ptr<OptionValueBase> Value,
				   std::string_view					ArgHelp)
{
	m_Options.push_back(
		{std::string(Group), std::string(ShortOpt), std::string(LongOpt), std::string(Desc), std::move(Value), std::string(ArgHelp)});
}

const Options::OptionInfo*
Options::FindOption(std::string_view Name) const
{
	for (const auto& Opt : m_Options)
	{
		if (Opt.LongOpt == Name || Opt.ShortOpt == Name)
			return &Opt;
	}
	return nullptr;
}

ParseResult
Options::Parse(int argc, char** argv)
{
	ParseResult Result;

	// Apply defaults before parsing
	for (const auto& Opt : m_Options)
	{
		if (Opt.Value && Opt.Value->HasDefault())
		{
			Opt.Value->Parse(Opt.Value->GetDefault());
		}
	}

	size_t PositionalIndex = 0;
	bool   StopParsing	   = false;

	for (int i = 1; i < argc; ++i)
	{
		std::string Arg(argv[i]);

		if (StopParsing)
		{
			Result.AddUnmatched(Arg);
			continue;
		}

		if (Arg == "--")
		{
			StopParsing = true;
			continue;
		}

		if (Arg.starts_with("--"))
		{
			// Long option
			std::string OptName;
			std::string OptValue;
			bool		HasValue = false;

			auto EqPos = Arg.find('=', 2);
			if (EqPos != std::string::npos)
			{
				OptName	 = Arg.substr(2, EqPos - 2);
				OptValue = Arg.substr(EqPos + 1);
				HasValue = true;
			}
			else
			{
				OptName = Arg.substr(2);
			}

			const OptionInfo* Info = FindOption(OptName);
			if (!Info)
			{
				Result.AddUnmatched(Arg);
				continue;
			}

			Result.SetCount(Info->LongOpt);

			if (Info->Value)
			{
				if (HasValue)
				{
					Info->Value->Parse(OptValue);
				}
				else if (Info->Value->IsBoolean())
				{
					// Booleans: check if next arg looks like a value
					if (i + 1 < argc)
					{
						std::string Next(argv[i + 1]);
						if (Next == "true" || Next == "false" || Next == "1" || Next == "0")
						{
							Info->Value->Parse(Next);
							i++;
						}
						else if (Info->Value->HasImplicit())
						{
							Info->Value->Parse(Info->Value->GetImplicit());
						}
					}
					else if (Info->Value->HasImplicit())
					{
						Info->Value->Parse(Info->Value->GetImplicit());
					}
				}
				else if (Info->Value->HasImplicit() && (i + 1 >= argc || std::string_view(argv[i + 1]).starts_with("-")))
				{
					// Non-boolean with implicit value and no next value arg available
					Info->Value->Parse(Info->Value->GetImplicit());
				}
				else if (i + 1 < argc)
				{
					Info->Value->Parse(argv[++i]);
				}
				else
				{
					throw std::runtime_error("Option '--" + OptName + "' requires a value");
				}
			}
		}
		else if (Arg.starts_with("-") && Arg.size() > 1)
		{
			// Short option
			std::string ShortName(1, Arg[1]);

			const OptionInfo* Info = FindOption(ShortName);
			if (!Info)
			{
				Result.AddUnmatched(Arg);
				continue;
			}

			Result.SetCount(Info->LongOpt);

			if (Info->Value)
			{
				if (Arg.size() > 2)
				{
					// Value attached: -sVALUE
					Info->Value->Parse(Arg.substr(2));
				}
				else if (Info->Value->IsBoolean() && Info->Value->HasImplicit())
				{
					Info->Value->Parse(Info->Value->GetImplicit());
				}
				else if (i + 1 < argc)
				{
					Info->Value->Parse(argv[++i]);
				}
				else
				{
					throw std::runtime_error("Option '-" + ShortName + "' requires a value");
				}
			}
		}
		else
		{
			// Positional argument
			if (PositionalIndex < m_PositionalNames.size())
			{
				const std::string& PosName = m_PositionalNames[PositionalIndex];
				const OptionInfo*  Info	   = FindOption(PosName);
				if (Info && Info->Value)
				{
					Info->Value->Parse(Arg);
					Result.SetCount(Info->LongOpt);
				}
				PositionalIndex++;
			}
			else
			{
				Result.AddUnmatched(Arg);
			}
		}
	}

	return Result;
}

void
Options::ParsePositional(std::initializer_list<std::string> Names)
{
	m_PositionalNames = std::vector<std::string>(Names);
}

void
Options::SetPositionalHelp(std::string Text)
{
	m_PositionalHelp = std::move(Text);
}

void
Options::ShowPositionalHelp()
{
	m_ShowPositionalHelp = true;
}

void
Options::SetWidth(int Width)
{
	m_Width = Width;
}

const std::string&
Options::Program() const
{
	return m_Program;
}

std::string
Options::Help() const
{
	std::string Out;

	// Usage line
	Out += "Usage:\n  " + m_Program + " [OPTION...]";
	if (!m_PositionalHelp.empty())
	{
		Out += " " + m_PositionalHelp;
	}
	Out += "\n\n";

	// Group options by group name, preserving order
	std::vector<std::string>							  GroupOrder;
	std::map<std::string, std::vector<const OptionInfo*>> GroupMap;

	for (const auto& Opt : m_Options)
	{
		if (GroupMap.find(Opt.Group) == GroupMap.end())
		{
			GroupOrder.push_back(Opt.Group);
		}
		GroupMap[Opt.Group].push_back(&Opt);
	}

	const int DescColumn = 30;

	for (const auto& GroupName : GroupOrder)
	{
		if (!GroupName.empty())
		{
			Out += " " + GroupName + ":\n";
		}

		for (const auto* Opt : GroupMap[GroupName])
		{
			std::string Line = "  ";

			if (!Opt->ShortOpt.empty())
			{
				Line += "-" + Opt->ShortOpt + ", ";
			}
			else
			{
				Line += "    ";
			}

			Line += "--" + Opt->LongOpt;

			// Show arg hint
			if (Opt->Value && !Opt->Value->IsBoolean())
			{
				if (!Opt->ArgHelp.empty())
				{
					Line += " " + Opt->ArgHelp;
				}
				else
				{
					Line += " <arg>";
				}
			}

			// Pad or newline for description
			if ((int)Line.size() < DescColumn)
			{
				Line += std::string(DescColumn - Line.size(), ' ');
			}
			else
			{
				Line += "\n" + std::string(DescColumn, ' ');
			}

			Line += Opt->Desc;

			// Show default value
			if (Opt->Value && Opt->Value->HasDefault() && !Opt->Value->GetDefault().empty())
			{
				Line += " (default: " + Opt->Value->GetDefault() + ")";
			}

			Out += Line + "\n";
		}

		Out += "\n";
	}

	return Out;
}

}  // namespace zen

//////////////////////////////////////////////////////////////////////////
//
// Unit tests
//

#if ZEN_WITH_TESTS

#	include <zencore/testing.h>

namespace zen {

void
options_forcelink()
{
}

TEST_CASE("options.basic")
{
	std::string Host;
	uint16_t	Port	= 0;
	bool		Verbose = false;

	Options Opts("testprog", "A test program");
	Opts.AddOptions()("h,host", "Host name", MakeValue(Host))("p,port", "Port number", MakeValue(Port))("v,verbose",
																										"Verbose output",
																										MakeValue(Verbose));

	char* Args[] = {(char*)"testprog", (char*)"--host", (char*)"localhost", (char*)"-p", (char*)"8080", (char*)"-v"};
	auto  Result = Opts.Parse(6, Args);

	CHECK_EQ(Host, "localhost");
	CHECK_EQ(Port, 8080);
	CHECK_EQ(Verbose, true);
	CHECK_EQ(Result.Count("host"), 1);
	CHECK_EQ(Result.Count("port"), 1);
	CHECK_EQ(Result.Count("verbose"), 1);
	CHECK_EQ(Result.Count("nonexistent"), 0);
}

TEST_CASE("options.defaults")
{
	std::string Mode;
	bool		Debug = false;
	int			Count = 0;

	Options Opts("test", "test");
	Opts.AddOptions()("mode", "Mode", MakeValue(Mode)->DefaultValue("fast"))("debug", "Debug", MakeValue(Debug)->DefaultValue("false"))(
		"count",
		"Count",
		MakeValue(Count)->DefaultValue("42"));

	char* Args[] = {(char*)"test"};
	auto  Result = Opts.Parse(1, Args);

	CHECK_EQ(Mode, "fast");
	CHECK_EQ(Debug, false);
	CHECK_EQ(Count, 42);
	CHECK_EQ(Result.Count("mode"), 0);
}

TEST_CASE("options.long_with_equals")
{
	std::string Name;
	int			Value = 0;

	Options Opts("test", "test");
	Opts.AddOptions()("name", "Name", MakeValue(Name))("value", "Value", MakeValue(Value));

	char* Args[] = {(char*)"test", (char*)"--name=hello", (char*)"--value=99"};
	auto  Result = Opts.Parse(3, Args);

	CHECK_EQ(Name, "hello");
	CHECK_EQ(Value, 99);
}

TEST_CASE("options.bool_implicit")
{
	bool FlagA = false;
	bool FlagB = false;
	bool FlagC = true;

	Options Opts("test", "test");
	Opts.AddOptions()("a,alpha", "Flag A", MakeValue(FlagA))("b,beta", "Flag B", MakeValue(FlagB))("c,gamma", "Flag C", MakeValue(FlagC));

	SUBCASE("implicit true")
	{
		char* Args[] = {(char*)"test", (char*)"--alpha", (char*)"-b"};
		Opts.Parse(3, Args);

		CHECK_EQ(FlagA, true);
		CHECK_EQ(FlagB, true);
		CHECK_EQ(FlagC, true);
	}

	SUBCASE("explicit false")
	{
		char* Args[] = {(char*)"test", (char*)"--alpha", (char*)"false", (char*)"--gamma", (char*)"0"};
		Opts.Parse(5, Args);

		CHECK_EQ(FlagA, false);
		CHECK_EQ(FlagC, false);
	}
}

TEST_CASE("options.positional")
{
	std::string Source;
	std::string Target;

	Options Opts("test", "test");
	Opts.AddOptions()("source", "Source", MakeValue(Source))("target", "Target", MakeValue(Target));
	Opts.ParsePositional({"source", "target"});

	char* Args[] = {(char*)"test", (char*)"/path/to/src", (char*)"/path/to/dst"};
	auto  Result = Opts.Parse(3, Args);

	CHECK_EQ(Source, "/path/to/src");
	CHECK_EQ(Target, "/path/to/dst");
	CHECK_EQ(Result.Count("source"), 1);
	CHECK_EQ(Result.Count("target"), 1);
}

TEST_CASE("options.double_dash_stop")
{
	std::string Command;
	bool		Debug = false;

	Options Opts("test", "test");
	Opts.AddOptions()("c,command", "Command", MakeValue(Command))("d,debug", "Debug", MakeValue(Debug));

	char* Args[] = {(char*)"test", (char*)"--command", (char*)"run", (char*)"--", (char*)"--debug", (char*)"extra"};
	auto  Result = Opts.Parse(6, Args);

	CHECK_EQ(Command, "run");
	CHECK_EQ(Debug, false);

	const auto& Unmatched = Result.Unmatched();
	CHECK_EQ(Unmatched.size(), 2);
	CHECK_EQ(Unmatched[0], "--debug");
	CHECK_EQ(Unmatched[1], "extra");
}

TEST_CASE("options.unmatched")
{
	bool Verbose = false;

	Options Opts("test", "test");
	Opts.AddOptions()("v,verbose", "Verbose", MakeValue(Verbose));

	char* Args[] = {(char*)"test", (char*)"--verbose", (char*)"--unknown", (char*)"-x"};
	auto  Result = Opts.Parse(4, Args);

	CHECK_EQ(Verbose, true);

	const auto& Unmatched = Result.Unmatched();
	CHECK_EQ(Unmatched.size(), 2);
	CHECK_EQ(Unmatched[0], "--unknown");
	CHECK_EQ(Unmatched[1], "-x");
}

TEST_CASE("options.groups")
{
	std::string Host;
	uint16_t	Port	= 0;
	bool		Verbose = false;

	Options Opts("test", "test");
	Opts.AddOptions("network")("host", "Host", MakeValue(Host))("port", "Port", MakeValue(Port));
	Opts.AddOptions("logging")("verbose", "Verbose", MakeValue(Verbose));

	std::string Help = Opts.Help();

	CHECK(Help.find("network:") != std::string::npos);
	CHECK(Help.find("logging:") != std::string::npos);
	CHECK(Help.find("--host") != std::string::npos);
	CHECK(Help.find("--verbose") != std::string::npos);
}

TEST_CASE("options.add_option_direct")
{
	std::string HostUrl;
	bool		DryRun = false;

	Options Opts("test", "test");
	Opts.AddOption("server", "u", "hosturl", "Host URL", MakeValue(HostUrl)->DefaultValue(""), "<hosturl>");
	Opts.AddOption("server", "n", "dry", "Dry run", MakeValue(DryRun), "<bool>");

	char* Args[] = {(char*)"test", (char*)"-u", (char*)"http://localhost:8080", (char*)"-n"};
	auto  Result = Opts.Parse(4, Args);

	CHECK_EQ(HostUrl, "http://localhost:8080");
	CHECK_EQ(DryRun, true);
	CHECK_EQ(Result.Count("hosturl"), 1);
	CHECK_EQ(Result.Count("dry"), 1);
}

TEST_CASE("options.integer_types")
{
	uint8_t	 U8	 = 0;
	uint16_t U16 = 0;
	uint64_t U64 = 0;
	int		 I32 = 0;
	unsigned U32 = 0;

	Options Opts("test", "test");
	Opts.AddOptions()("u8", "u8", MakeValue(U8))("u16", "u16", MakeValue(U16))("u64", "u64", MakeValue(U64))("i32", "i32", MakeValue(I32))(
		"u32",
		"u32",
		MakeValue(U32));

	char* Args[] = {(char*)"test",
					(char*)"--u8",
					(char*)"255",
					(char*)"--u16",
					(char*)"65535",
					(char*)"--u64",
					(char*)"18446744073709551615",
					(char*)"--i32",
					(char*)"-42",
					(char*)"--u32",
					(char*)"4294967295"};
	auto  Result = Opts.Parse(11, Args);

	CHECK_EQ(U8, 255);
	CHECK_EQ(U16, 65535);
	CHECK_EQ(U64, 18446744073709551615ULL);
	CHECK_EQ(I32, -42);
	CHECK_EQ(U32, 4294967295U);
}

TEST_CASE("options.vector_string")
{
	std::vector<std::string> Urls;

	Options Opts("test", "test");
	Opts.AddOption("", "", "url", "URLs", MakeValue(Urls), "");

	SUBCASE("comma separated")
	{
		char* Args[] = {(char*)"test", (char*)"--url", (char*)"http://a.com,http://b.com"};
		Opts.Parse(3, Args);

		CHECK_EQ(Urls.size(), 2);
		CHECK_EQ(Urls[0], "http://a.com");
		CHECK_EQ(Urls[1], "http://b.com");
	}

	SUBCASE("multiple flags")
	{
		char* Args[] = {(char*)"test", (char*)"--url", (char*)"http://a.com", (char*)"--url", (char*)"http://b.com"};
		Opts.Parse(5, Args);

		CHECK_EQ(Urls.size(), 2);
		CHECK_EQ(Urls[0], "http://a.com");
		CHECK_EQ(Urls[1], "http://b.com");
	}
}

TEST_CASE("options.path")
{
	std::filesystem::path Dir;

	Options Opts("test", "test");
	Opts.AddOptions()("dir", "Directory", MakeValue(Dir));

	char* Args[] = {(char*)"test", (char*)"--dir", (char*)"C:\\Users\\test\\data"};
	Opts.Parse(3, Args);

	CHECK_EQ(Dir, std::filesystem::path("C:\\Users\\test\\data").make_preferred());
}

TEST_CASE("options.implicit_value")
{
	std::string Scrub;

	Options Opts("test", "test");
	Opts.AddOptions()("scrub", "Scrub", MakeValue(Scrub)->ImplicitValue("yes"));

	SUBCASE("with implicit")
	{
		char* Args[] = {(char*)"test", (char*)"--scrub"};
		auto  Result = Opts.Parse(2, Args);

		CHECK_EQ(Scrub, "yes");
		CHECK_EQ(Result.Count("scrub"), 1);
	}

	SUBCASE("with explicit value")
	{
		char* Args[] = {(char*)"test", (char*)"--scrub=nocas"};
		Opts.Parse(2, Args);

		CHECK_EQ(Scrub, "nocas");
	}
}

TEST_CASE("options.missing_value_throws")
{
	std::string Name;

	Options Opts("test", "test");
	Opts.AddOptions()("name", "Name", MakeValue(Name));

	char* Args[] = {(char*)"test", (char*)"--name"};

	CHECK_THROWS_AS(Opts.Parse(2, Args), std::runtime_error);
}

TEST_CASE("options.help_output")
{
	std::string Host;
	uint16_t	Port = 0;

	Options Opts("myprog", "My program description");
	Opts.AddOption("", "h", "host", "Host name", MakeValue(Host), "<hostname>");
	Opts.AddOption("", "p", "port", "Port number", MakeValue(Port)->DefaultValue("8080"), "<port>");

	std::string Help = Opts.Help();

	CHECK(Help.find("myprog") != std::string::npos);
	CHECK(Help.find("-h, --host") != std::string::npos);
	CHECK(Help.find("<hostname>") != std::string::npos);
	CHECK(Help.find("-p, --port") != std::string::npos);
	CHECK(Help.find("(default: 8080)") != std::string::npos);
}

TEST_CASE("options.short_value_attached")
{
	int Port = 0;

	Options Opts("test", "test");
	Opts.AddOption("", "p", "port", "Port", MakeValue(Port), "");

	char* Args[] = {(char*)"test", (char*)"-p8080"};
	Opts.Parse(2, Args);

	CHECK_EQ(Port, 8080);
}

}  // namespace zen

#endif