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

#include "luaconfig.h"

namespace zen::LuaConfig {

void
EscapeBackslash(std::string& InOutString)
{
	std::size_t BackslashPos = InOutString.find('\\');
	if (BackslashPos != std::string::npos)
	{
		std::size_t						  Offset = 0;
		zen::ExtendableStringBuilder<512> PathBuilder;
		while (BackslashPos != std::string::npos)
		{
			PathBuilder.Append(InOutString.substr(Offset, BackslashPos + 1 - Offset));
			PathBuilder.Append('\\');
			Offset		 = BackslashPos + 1;
			BackslashPos = InOutString.find('\\', Offset);
		}
		PathBuilder.Append(InOutString.substr(Offset, BackslashPos));
		InOutString = PathBuilder.ToString();
	}
}

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

BoolOption::BoolOption(bool& Value) : Value(Value)
{
}

void
BoolOption::Print(std::string_view, zen::StringBuilderBase& StringBuilder)
{
	StringBuilder.Append(Value ? "true" : "false");
}

void
BoolOption::Parse(sol::object Object)
{
	Value = Object.as<bool>();
}

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

StringOption::StringOption(std::string& Value) : Value(Value)
{
}

void
StringOption::Print(std::string_view, zen::StringBuilderBase& StringBuilder)
{
	StringBuilder.Append(fmt::format("\"{}\"", Value));
}

void
StringOption::Parse(sol::object Object)
{
	Value = Object.as<std::string>();
}

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

FilePathOption::FilePathOption(std::filesystem::path& Value) : Value(Value)
{
}

void
FilePathOption::Print(std::string_view, zen::StringBuilderBase& StringBuilder)
{
	std::string Path = Value.string();
	EscapeBackslash(Path);
	StringBuilder.Append(fmt::format("\"{}\"", Path));
}

void
FilePathOption::Parse(sol::object Object)
{
	std::string Str = Object.as<std::string>();
	if (!Str.empty())
	{
		Value = MakeSafeAbsolutePath(Str);
	}
}

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

LuaContainerWriter::LuaContainerWriter(zen::StringBuilderBase& StringBuilder, std::string_view Indent)
: StringBuilder(StringBuilder)
, InitialIndent(Indent.length())
, LocalIndent(Indent)
{
	StringBuilder.Append("{\n");
	LocalIndent.push_back('\t');
}

LuaContainerWriter::~LuaContainerWriter()
{
	LocalIndent.pop_back();
	StringBuilder.Append(LocalIndent);
	StringBuilder.Append("}");
}

void
LuaContainerWriter::BeginContainer(std::string_view Name)
{
	StringBuilder.Append(LocalIndent);
	if (!Name.empty())
	{
		StringBuilder.Append(Name);
		StringBuilder.Append(" = {\n");
	}
	else
	{
		StringBuilder.Append("{\n");
	}
	LocalIndent.push_back('\t');
}

void
LuaContainerWriter::WriteValue(std::string_view Name, std::string_view Value)
{
	if (Name.empty())
	{
		StringBuilder.Append(fmt::format("{}\"{}\",\n", LocalIndent, Value));
	}
	else
	{
		StringBuilder.Append(fmt::format("{}{} = \"{}\",\n", LocalIndent, Name, Value));
	}
}

void
LuaContainerWriter::EndContainer()
{
	LocalIndent.pop_back();
	StringBuilder.Append(LocalIndent);
	StringBuilder.Append("}");
	StringBuilder.Append(",\n");
}

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

StringArrayOption::StringArrayOption(std::vector<std::string>& Value) : Value(Value)
{
}

void
StringArrayOption::Print(std::string_view Indent, zen::StringBuilderBase& StringBuilder)
{
	if (Value.empty())
	{
		StringBuilder.Append("{}");
	}
	if (Value.size() == 1)
	{
		StringBuilder.Append(fmt::format("\"{}\"", Value[0]));
	}
	else
	{
		LuaContainerWriter Writer(StringBuilder, Indent);
		for (std::string String : Value)
		{
			Writer.WriteValue("", String);
		}
	}
}

void
StringArrayOption::Parse(sol::object Object)
{
	if (Object.get_type() == sol::type::string)
	{
		Value.push_back(Object.as<std::string>());
	}
	else if (Object.get_type() == sol::type::table)
	{
		for (const auto& Kv : Object.as<sol::table>())
		{
			Value.push_back(Kv.second.as<std::string>());
		}
	}
}

std::shared_ptr<OptionValue>
MakeOption(std::string& Value)
{
	return std::make_shared<StringOption>(Value);
}

std::shared_ptr<OptionValue>
MakeOption(std::filesystem::path& Value)
{
	return std::make_shared<FilePathOption>(Value);
}

std::shared_ptr<OptionValue>
MakeOption(bool& Value)
{
	return std::make_shared<BoolOption>(Value);
}

std::shared_ptr<OptionValue>
MakeOption(std::vector<std::string>& Value)
{
	return std::make_shared<StringArrayOption>(Value);
}

void
Options::Parse(const std::filesystem::path& Path, const cxxopts::ParseResult& CmdLineResult)
{
	zen::IoBuffer LuaScript = zen::IoBufferBuilder::MakeFromFile(Path);

	if (LuaScript)
	{
		sol::state lua;

		lua.open_libraries(sol::lib::base);

		lua.set_function("getenv", [&](const std::string env) -> sol::object {
#if ZEN_PLATFORM_WINDOWS
			std::wstring EnvVarValue;
			size_t		 RequiredSize = 0;
			std::wstring EnvWide	  = zen::Utf8ToWide(env);
			_wgetenv_s(&RequiredSize, nullptr, 0, EnvWide.c_str());

			if (RequiredSize == 0)
				return sol::make_object(lua, sol::lua_nil);

			EnvVarValue.resize(RequiredSize);
			_wgetenv_s(&RequiredSize, EnvVarValue.data(), RequiredSize, EnvWide.c_str());
			return sol::make_object(lua, zen::WideToUtf8(EnvVarValue.c_str()));
#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
			char* EnvVariable = getenv(env.c_str());
			if (EnvVariable == nullptr)
			{
				return sol::make_object(lua, sol::lua_nil);
			}
			return sol::make_object(lua, EnvVariable);
#else
			ZEN_UNUSED(env);
			return sol::make_object(lua, sol::lua_nil);
#endif
		});

		try
		{
			sol::load_result config = lua.load(std::string_view((const char*)LuaScript.Data(), LuaScript.Size()), "zen_cfg");

			if (!config.valid())
			{
				sol::error err = config;

				std::string ErrorString = sol::to_string(config.status());

				throw std::runtime_error(fmt::format("{} ('{}')", ErrorString, err.what()));
			}

			config();
		}
		catch (const std::exception& e)
		{
			throw std::runtime_error(fmt::format("failed to load config script ('{}'): {}", Path, e.what()).c_str());
		}

		Parse(lua, CmdLineResult);
	}
}

void
Options::Parse(const sol::state& LuaState, const cxxopts::ParseResult& CmdLineResult)
{
	for (auto It : LuaState)
	{
		sol::object Key		= It.first;
		sol::type	KeyType = Key.get_type();
		if (KeyType == sol::type::string)
		{
			sol::type ValueType = It.second.get_type();
			switch (ValueType)
			{
				case sol::type::table:
					{
						std::string Name = Key.as<std::string>();
						if (Name.starts_with("_"))
						{
							continue;
						}
						if (Name == "base")
						{
							continue;
						}
						Traverse(It.second.as<sol::table>(), Name, CmdLineResult);
					}
					break;
				default:
					break;
			}
		}
	}
}

void
Options::Touch(std::string_view Key)
{
	UsedKeys.insert(std::string(Key));
}

void
Options::Print(zen::StringBuilderBase& SB, const cxxopts::ParseResult& CmdLineResult)
{
	for (auto It : OptionMap)
	{
		if (CmdLineResult.count(It.second.CommandLineOptionName) != 0)
		{
			UsedKeys.insert(It.first);
		}
	}

	std::vector<std::string> SortedKeys(UsedKeys.begin(), UsedKeys.end());
	std::sort(SortedKeys.begin(), SortedKeys.end());
	auto GetTablePath = [](const std::string& Key) -> std::vector<std::string> {
		std::vector<std::string> Path;
		zen::ForEachStrTok(Key, '.', [&Path](std::string_view Part) {
			Path.push_back(std::string(Part));
			return true;
		});
		return Path;
	};
	std::vector<std::string> CurrentTablePath;
	std::string				 Indent;
	for (const std::string& Key : SortedKeys)
	{
		std::vector<std::string> KeyPath = GetTablePath(Key);
		std::string				 Name	 = KeyPath.back();
		KeyPath.pop_back();
		if (CurrentTablePath != KeyPath)
		{
			size_t EqualCount = 0;
			while (EqualCount < CurrentTablePath.size() && EqualCount < KeyPath.size() &&
				   CurrentTablePath[EqualCount] == KeyPath[EqualCount])
			{
				EqualCount++;
			}
			while (CurrentTablePath.size() > EqualCount)
			{
				CurrentTablePath.pop_back();
				Indent.pop_back();
				SB.Append(Indent);
				SB.Append("}");
				if (CurrentTablePath.size() == EqualCount && !Indent.empty() && KeyPath.size() >= EqualCount)
				{
					SB.Append(",");
				}
				SB.Append("\n");
				if (Indent.empty())
				{
					SB.Append("\n");
				}
			}
			while (EqualCount < KeyPath.size())
			{
				SB.Append(Indent);
				SB.Append(KeyPath[EqualCount]);
				SB.Append(" = {\n");
				Indent.push_back('\t');
				CurrentTablePath.push_back(KeyPath[EqualCount]);
				EqualCount++;
			}
		}

		SB.Append(Indent);
		SB.Append(Name);
		SB.Append(" = ");
		OptionMap[Key].Value->Print(Indent, SB);
		SB.Append(",\n");
	}
	while (!CurrentTablePath.empty())
	{
		Indent.pop_back();
		SB.Append(Indent);
		SB.Append("}\n");
		CurrentTablePath.pop_back();
	}
}

void
Options::Traverse(sol::table Table, std::string_view PathPrefix, const cxxopts::ParseResult& CmdLineResult)
{
	for (auto It : Table)
	{
		sol::object Key		= It.first;
		sol::type	KeyType = Key.get_type();
		if (KeyType == sol::type::string || KeyType == sol::type::number)
		{
			sol::type ValueType = It.second.get_type();
			switch (ValueType)
			{
				case sol::type::table:
				case sol::type::string:
				case sol::type::number:
				case sol::type::boolean:
					{
						std::string Name = Key.as<std::string>();
						if (Name.starts_with("_"))
						{
							continue;
						}
						Name		  = std::string(PathPrefix) + "." + Key.as<std::string>();
						auto OptionIt = OptionMap.find(Name);
						if (OptionIt != OptionMap.end())
						{
							UsedKeys.insert(Name);
							if (CmdLineResult.count(OptionIt->second.CommandLineOptionName) != 0)
							{
								continue;
							}
							OptionIt->second.Value->Parse(It.second);
							continue;
						}
						if (ValueType == sol::type::table)
						{
							if (Name == "base")
							{
								continue;
							}
							Traverse(It.second.as<sol::table>(), Name, CmdLineResult);
						}
					}
					break;
				default:
					break;
			}
		}
	}
}

}  // namespace zen::LuaConfig