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

#include <zencore/fmtutils.h>
#include <zencore/iobuffer.h>
#include <zencore/logging.h>
#include <zencore/memoryview.h>
#include <zencore/trace.h>
#include <zenhorde/hordeclient.h>
#include <zenhttp/httpclient.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <json11.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen::horde {

HordeClient::HordeClient(const HordeConfig& Config) : m_Config(Config), m_Log(zen::logging::Get("horde.client"))
{
}

HordeClient::~HordeClient() = default;

bool
HordeClient::Initialize()
{
	ZEN_TRACE_CPU("HordeClient::Initialize");

	HttpClientSettings Settings;
	Settings.LogCategory		= "horde.http";
	Settings.ConnectTimeout		= std::chrono::milliseconds{10000};
	Settings.Timeout			= std::chrono::milliseconds{60000};
	Settings.RetryCount			= 1;
	Settings.ExpectedErrorCodes = {HttpResponseCode::ServiceUnavailable, HttpResponseCode::TooManyRequests};

	if (!m_Config.AuthToken.empty())
	{
		Settings.AccessTokenProvider = [token = m_Config.AuthToken]() -> HttpClientAccessToken {
			HttpClientAccessToken Token;
			Token.Value		 = token;
			Token.ExpireTime = HttpClientAccessToken::Clock::now() + std::chrono::hours{24};
			return Token;
		};
	}

	m_Http = std::make_unique<zen::HttpClient>(m_Config.ServerUrl, Settings);

	if (!m_Config.AuthToken.empty())
	{
		if (!m_Http->Authenticate())
		{
			ZEN_WARN("failed to authenticate with Horde server");
			return false;
		}
	}

	return true;
}

std::string
HordeClient::BuildRequestBody() const
{
	json11::Json::object Requirements;

	if (m_Config.Mode == ConnectionMode::Direct && !m_Config.Pool.empty())
	{
		Requirements["pool"] = m_Config.Pool;
	}

	std::string Condition;
#if ZEN_PLATFORM_WINDOWS
	ExtendableStringBuilder<256> CondBuf;
	CondBuf << "(OSFamily == 'Windows' || WineEnabled == '" << (m_Config.AllowWine ? "true" : "false") << "')";
	Condition = std::string(CondBuf);
#elif ZEN_PLATFORM_MAC
	Condition = "OSFamily == 'MacOS'";
#else
	Condition = "OSFamily == 'Linux'";
#endif

	if (!m_Config.Condition.empty())
	{
		Condition += " ";
		Condition += m_Config.Condition;
	}

	Requirements["condition"] = Condition;
	Requirements["exclusive"] = true;

	json11::Json::object Connection;
	Connection["modePreference"] = ToString(m_Config.Mode);

	if (m_Config.EncryptionMode != Encryption::None)
	{
		Connection["encryption"] = ToString(m_Config.EncryptionMode);
	}

	// Request configured zen service port to be forwarded. The Horde agent will map this
	// to a local port on the provisioned machine and report it back in the response.
	json11::Json::object PortsObj;
	PortsObj["ZenPort"] = json11::Json(m_Config.ZenServicePort);
	Connection["ports"] = PortsObj;

	json11::Json::object Root;
	Root["requirements"] = Requirements;
	Root["connection"]	 = Connection;

	return json11::Json(Root).dump();
}

bool
HordeClient::ResolveCluster(const std::string& RequestBody, ClusterInfo& OutCluster)
{
	ZEN_TRACE_CPU("HordeClient::ResolveCluster");

	const IoBuffer Payload = IoBufferBuilder::MakeFromMemory(MemoryView{RequestBody.data(), RequestBody.size()}, ZenContentType::kJSON);

	const HttpClient::Response Response = m_Http->Post("api/v2/compute/_cluster", Payload);

	if (Response.Error)
	{
		ZEN_WARN("cluster resolution failed: {}", Response.Error->ErrorMessage);
		return false;
	}

	const int StatusCode = static_cast<int>(Response.StatusCode);

	if (StatusCode == 503 || StatusCode == 429)
	{
		ZEN_DEBUG("cluster resolution returned HTTP/{}: no resources", StatusCode);
		return false;
	}

	if (StatusCode == 401)
	{
		ZEN_WARN("cluster resolution returned HTTP/401: token expired");
		return false;
	}

	if (!Response.IsSuccess())
	{
		ZEN_WARN("cluster resolution failed with HTTP/{}", StatusCode);
		return false;
	}

	const std::string  Body(Response.AsText());
	std::string		   Err;
	const json11::Json Json = json11::Json::parse(Body, Err);

	if (!Err.empty())
	{
		ZEN_WARN("invalid JSON response for cluster resolution: {}", Err);
		return false;
	}

	const json11::Json ClusterIdVal = Json["clusterId"];
	if (!ClusterIdVal.is_string() || ClusterIdVal.string_value().empty())
	{
		ZEN_WARN("missing 'clusterId' in cluster resolution response");
		return false;
	}

	OutCluster.ClusterId = ClusterIdVal.string_value();
	return true;
}

bool
HordeClient::ParseHexBytes(std::string_view Hex, uint8_t* Out, size_t OutSize)
{
	if (Hex.size() != OutSize * 2)
	{
		return false;
	}

	for (size_t i = 0; i < OutSize; ++i)
	{
		auto HexToByte = [](char c) -> int {
			if (c >= '0' && c <= '9')
				return c - '0';
			if (c >= 'a' && c <= 'f')
				return c - 'a' + 10;
			if (c >= 'A' && c <= 'F')
				return c - 'A' + 10;
			return -1;
		};

		const int Hi = HexToByte(Hex[i * 2]);
		const int Lo = HexToByte(Hex[i * 2 + 1]);
		if (Hi < 0 || Lo < 0)
		{
			return false;
		}
		Out[i] = static_cast<uint8_t>((Hi << 4) | Lo);
	}

	return true;
}

bool
HordeClient::RequestMachine(const std::string& RequestBody, const std::string& ClusterId, MachineInfo& OutMachine)
{
	ZEN_TRACE_CPU("HordeClient::RequestMachine");

	ZEN_INFO("requesting machine from Horde with cluster '{}'", ClusterId.empty() ? "default" : ClusterId.c_str());

	ExtendableStringBuilder<128> ResourcePath;
	ResourcePath << "api/v2/compute/" << (ClusterId.empty() ? "default" : ClusterId.c_str());

	const IoBuffer Payload = IoBufferBuilder::MakeFromMemory(MemoryView{RequestBody.data(), RequestBody.size()}, ZenContentType::kJSON);
	const HttpClient::Response Response = m_Http->Post(ResourcePath.ToView(), Payload);

	// Reset output to invalid state
	OutMachine		= {};
	OutMachine.Port = 0xFFFF;

	if (Response.Error)
	{
		ZEN_WARN("machine request failed: {}", Response.Error->ErrorMessage);
		return false;
	}

	const int StatusCode = static_cast<int>(Response.StatusCode);

	if (StatusCode == 404 || StatusCode == 503 || StatusCode == 429)
	{
		ZEN_DEBUG("machine request returned HTTP/{}: no resources", StatusCode);
		return false;
	}

	if (StatusCode == 401)
	{
		ZEN_WARN("machine request returned HTTP/401: token expired");
		return false;
	}

	if (!Response.IsSuccess())
	{
		ZEN_WARN("machine request failed with HTTP/{}", StatusCode);
		return false;
	}

	const std::string  Body(Response.AsText());
	std::string		   Err;
	const json11::Json Json = json11::Json::parse(Body, Err);

	if (!Err.empty())
	{
		ZEN_WARN("invalid JSON response for machine request: {}", Err);
		return false;
	}

	// Required fields
	const json11::Json NonceVal = Json["nonce"];
	const json11::Json IpVal	= Json["ip"];
	const json11::Json PortVal	= Json["port"];

	if (!NonceVal.is_string() || !IpVal.is_string() || !PortVal.is_number())
	{
		ZEN_WARN("missing 'nonce', 'ip', or 'port' in machine response");
		return false;
	}

	OutMachine.Ip	= IpVal.string_value();
	OutMachine.Port = static_cast<uint16_t>(PortVal.int_value());

	if (!ParseHexBytes(NonceVal.string_value(), OutMachine.Nonce, NonceSize))
	{
		ZEN_WARN("invalid nonce hex string in machine response");
		return false;
	}

	if (const json11::Json PortsVal = Json["ports"]; PortsVal.is_object())
	{
		for (const auto& [Key, Val] : PortsVal.object_items())
		{
			PortInfo Info;
			if (Val["port"].is_number())
			{
				Info.Port = static_cast<uint16_t>(Val["port"].int_value());
			}
			if (Val["agentPort"].is_number())
			{
				Info.AgentPort = static_cast<uint16_t>(Val["agentPort"].int_value());
			}
			OutMachine.Ports[Key] = Info;
		}
	}

	if (const json11::Json ConnectionModeVal = Json["connectionMode"]; ConnectionModeVal.is_string())
	{
		if (FromString(OutMachine.Mode, ConnectionModeVal.string_value()))
		{
			if (const json11::Json ConnectionAddressVal = Json["connectionAddress"]; ConnectionAddressVal.is_string())
			{
				OutMachine.ConnectionAddress = ConnectionAddressVal.string_value();
			}
		}
	}

	// Properties are a flat string array of "Key=Value" pairs describing the machine.
	// We extract OS family and core counts for sizing decisions. If neither core count
	// is available, we fall back to 16 as a conservative default.
	uint16_t LogicalCores  = 0;
	uint16_t PhysicalCores = 0;

	if (const json11::Json PropertiesVal = Json["properties"]; PropertiesVal.is_array())
	{
		for (const json11::Json& PropVal : PropertiesVal.array_items())
		{
			if (!PropVal.is_string())
			{
				continue;
			}

			const std::string Prop = PropVal.string_value();
			if (Prop.starts_with("OSFamily="))
			{
				if (Prop.substr(9) == "Windows")
				{
					OutMachine.IsWindows = true;
				}
			}
			else if (Prop.starts_with("LogicalCores="))
			{
				LogicalCores = static_cast<uint16_t>(std::atoi(Prop.c_str() + 13));
			}
			else if (Prop.starts_with("PhysicalCores="))
			{
				PhysicalCores = static_cast<uint16_t>(std::atoi(Prop.c_str() + 14));
			}
		}
	}

	if (LogicalCores > 0)
	{
		OutMachine.LogicalCores = LogicalCores;
	}
	else if (PhysicalCores > 0)
	{
		OutMachine.LogicalCores = PhysicalCores * 2;
	}
	else
	{
		OutMachine.LogicalCores = 16;
	}

	if (const json11::Json EncryptionVal = Json["encryption"]; EncryptionVal.is_string())
	{
		if (FromString(OutMachine.EncryptionMode, EncryptionVal.string_value()))
		{
			if (OutMachine.EncryptionMode == Encryption::AES)
			{
				const json11::Json KeyVal = Json["key"];
				if (KeyVal.is_string() && !KeyVal.string_value().empty())
				{
					if (!ParseHexBytes(KeyVal.string_value(), OutMachine.Key, KeySize))
					{
						ZEN_WARN("invalid AES key in machine response");
					}
				}
				else
				{
					ZEN_WARN("AES encryption requested but no key provided");
				}
			}
		}
	}

	if (const json11::Json LeaseIdVal = Json["leaseId"]; LeaseIdVal.is_string())
	{
		OutMachine.LeaseId = LeaseIdVal.string_value();
	}

	ZEN_INFO("Horde machine assigned [{}:{}] cores={} lease={}",
			 OutMachine.GetConnectionAddress(),
			 OutMachine.GetConnectionPort(),
			 OutMachine.LogicalCores,
			 OutMachine.LeaseId);

	return true;
}

}  // namespace zen::horde