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

#include <zenovermind/overmindclient.h>

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

ZEN_THIRD_PARTY_INCLUDES_START
#include <json11.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen::overmind {

namespace {

	HttpClient::KeyValueMap MakeOvermindHeaders(const OvermindConfig& Config)
	{
		HttpClient::KeyValueMap Headers;
		if (!Config.AuthToken.empty())
		{
			Headers->emplace("Authorization", "Bearer " + Config.AuthToken);
		}
		return Headers;
	}

}  // namespace

OvermindClient::OvermindClient(const OvermindConfig& Config) : m_Config(Config), m_Log(zen::logging::Get("overmind.client"))
{
}

OvermindClient::~OvermindClient() = default;

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

	HttpClientSettings Settings;
	Settings.LogCategory	= "overmind.http";
	Settings.ConnectTimeout = std::chrono::milliseconds{10000};
	Settings.Timeout		= std::chrono::milliseconds{60000};
	Settings.RetryCount		= 1;

	// Ensure the base URL ends with a slash so path concatenation works correctly
	std::string BaseUrl = m_Config.ServerUrl;
	if (!BaseUrl.empty() && BaseUrl.back() != '/')
	{
		BaseUrl += '/';
	}

	m_Http = std::make_unique<zen::HttpClient>(BaseUrl, Settings);

	return true;
}

std::string
OvermindClient::BuildJobJson(const std::string& JobName,
							 const std::string& OrchestratorEndpoint,
							 const std::string& CoordinatorSession,
							 bool				CleanStart,
							 const std::string& TraceHost) const
{
	ZEN_TRACE_CPU("OvermindClient::BuildJobJson");

	// Build the args array for the zenserver compute command
	json11::Json::array Args;
	Args.push_back("compute");
	Args.push_back("--http=asio");

	if (!OrchestratorEndpoint.empty())
	{
		ExtendableStringBuilder<256> CoordArg;
		CoordArg << "--coordinator-endpoint=" << OrchestratorEndpoint;
		Args.push_back(std::string(CoordArg.ToView()));
	}

	{
		ExtendableStringBuilder<128> IdArg;
		IdArg << "--instance-id=overmind-" << JobName;
		Args.push_back(std::string(IdArg.ToView()));
	}

	if (!CoordinatorSession.empty())
	{
		ExtendableStringBuilder<128> SessionArg;
		SessionArg << "--coordinator-session=" << CoordinatorSession;
		Args.push_back(std::string(SessionArg.ToView()));
	}

	if (CleanStart)
	{
		Args.push_back("--clean");
	}

	if (!TraceHost.empty())
	{
		ExtendableStringBuilder<128> TraceArg;
		TraceArg << "--tracehost=" << TraceHost;
		Args.push_back(std::string(TraceArg.ToView()));
	}

	json11::Json Task = json11::Json::object{
		{"name", "zenserver"},
		{"type", "TASK_TYPE_MAIN"},
		{"command", m_Config.CommandRef},
		{"args", Args},
		{"resources",
		 json11::Json::object{
			 {"memory", m_Config.Memory},
			 {"cpu", m_Config.Cpu},
		 }},
	};

	json11::Json Body = json11::Json::object{
		{"namespace", m_Config.Namespace},
		{"region", m_Config.Region},
		{"definition",
		 json11::Json::object{
			 {"name", JobName},
			 {"os", m_Config.Os},
			 {"arch", m_Config.Arch},
			 {"tasks", json11::Json::array{Task}},
		 }},
	};

	return Body.dump();
}

bool
OvermindClient::ScheduleJob(const std::string& JobJson, OvermindJobInfo& OutJob)
{
	ZEN_TRACE_CPU("OvermindClient::ScheduleJob");

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

	const HttpClient::Response Response = m_Http->Post("v1/jobs", Payload, MakeOvermindHeaders(m_Config));

	if (Response.Error)
	{
		ZEN_WARN("Overmind job schedule failed: {}", Response.Error->ErrorMessage);
		return false;
	}

	if (!Response.IsSuccess())
	{
		ZEN_WARN("Overmind job schedule failed with HTTP/{}", static_cast<int>(Response.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 from Overmind job schedule: {}", Err);
		return false;
	}

	const auto& Job = Json["job"];
	OutJob.Id		= Job["id"].string_value();
	OutJob.Status	= Job["status"].string_value();

	if (OutJob.Id.empty())
	{
		ZEN_WARN("Overmind job schedule response missing job ID");
		return false;
	}

	ZEN_INFO("Overmind job scheduled: id={}", OutJob.Id);

	return true;
}

bool
OvermindClient::GetJobStatus(const std::string& JobId, OvermindJobInfo& OutJob)
{
	ZEN_TRACE_CPU("OvermindClient::GetJobStatus");

	ExtendableStringBuilder<128> Path;
	Path << "v1/jobs/" << JobId;

	const HttpClient::Response Response = m_Http->Get(Path.ToView(), MakeOvermindHeaders(m_Config));

	if (Response.Error)
	{
		ZEN_WARN("Overmind job status query failed for '{}': {}", JobId, Response.Error->ErrorMessage);
		return false;
	}

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

	if (StatusCode == 404)
	{
		ZEN_INFO("Overmind job '{}' not found", JobId);
		OutJob.Status = "STATUS_ERROR";
		return true;
	}

	if (!Response.IsSuccess())
	{
		ZEN_WARN("Overmind job status query 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 in Overmind job status response: {}", Err);
		return false;
	}

	const auto& Job = Json["job"];
	OutJob.Id		= Job["id"].string_value();
	OutJob.Status	= Job["status"].string_value();

	return true;
}

bool
OvermindClient::CancelJob(const std::string& JobId)
{
	ZEN_TRACE_CPU("OvermindClient::CancelJob");

	ExtendableStringBuilder<256> Path;
	Path << "v1/jobs/" << JobId << "?namespace=" << m_Config.Namespace;

	const HttpClient::Response Response = m_Http->Delete(Path.ToView(), MakeOvermindHeaders(m_Config));

	if (Response.Error)
	{
		ZEN_WARN("Overmind job cancel failed for '{}': {}", JobId, Response.Error->ErrorMessage);
		return false;
	}

	if (!Response.IsSuccess())
	{
		ZEN_WARN("Overmind job cancel failed with HTTP/{}", static_cast<int>(Response.StatusCode));
		return false;
	}

	ZEN_INFO("Overmind job '{}' cancelled", JobId);
	return true;
}

}  // namespace zen::overmind