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

#include <zenutil/consul.h>

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

#include <fmt/format.h>

namespace zen::consul {

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

struct ConsulProcess::Impl
{
	Impl(std::string_view BaseUri) : m_HttpClient(BaseUri) {}
	~Impl() = default;

	void SpawnConsulAgent()
	{
		if (m_ProcessHandle.IsValid())
		{
			return;
		}

		CreateProcOptions Options;
		Options.Flags |= CreateProcOptions::Flag_Windows_NewProcessGroup;

		const std::filesystem::path ConsulExe = GetRunningExecutablePath().parent_path() / ("consul" ZEN_EXE_SUFFIX_LITERAL);
		CreateProcResult			Result	  = CreateProc(ConsulExe, "consul" ZEN_EXE_SUFFIX_LITERAL " agent -dev", Options);

		if (Result)
		{
			m_ProcessHandle.Initialize(Result);

			Stopwatch Timer;

			// Poll to check when the agent is ready

			do
			{
				Sleep(100);
				HttpClient::Response Resp = m_HttpClient.Get("v1/status/leader");
				if (Resp)
				{
					ZEN_INFO("Consul agent started successfully (waited {})", NiceTimeSpanMs(Timer.GetElapsedTimeMs()));

					return;
				}
			} while (Timer.GetElapsedTimeMs() < 10000);
		}

		// Report failure!

		ZEN_WARN("Consul agent failed to start within timeout period");
	}

	void StopConsulAgent()
	{
		if (!m_ProcessHandle.IsValid())
		{
			return;
		}

		// This waits for the process to exit and also resets the handle
		m_ProcessHandle.Kill();
	}

private:
	ProcessHandle m_ProcessHandle;
	HttpClient	  m_HttpClient;
};

ConsulProcess::ConsulProcess() : m_Impl(std::make_unique<Impl>("http://localhost:8500/"))
{
}

ConsulProcess::~ConsulProcess()
{
}

void
ConsulProcess::SpawnConsulAgent()
{
	m_Impl->SpawnConsulAgent();
}

void
ConsulProcess::StopConsulAgent()
{
	m_Impl->StopConsulAgent();
}

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

ConsulClient::ConsulClient(std::string_view BaseUri) : m_HttpClient(BaseUri)
{
}

ConsulClient::~ConsulClient()
{
}

void
ConsulClient::SetKeyValue(std::string_view Key, std::string_view Value)
{
	IoBuffer ValueBuffer = IoBufferBuilder::MakeFromMemory(MakeMemoryView(Value));
	ValueBuffer.SetContentType(HttpContentType::kText);
	HttpClient::Response Result = m_HttpClient.Put(fmt::format("v1/kv/{}", Key), ValueBuffer, HttpClient::Accept(HttpContentType::kJSON));
	if (!Result)
	{
		throw runtime_error("ConsulClient::SetKeyValue() failed to set key '{}' ({})", Key, Result.ErrorMessage(""));
	}
}

std::string
ConsulClient::GetKeyValue(std::string_view Key)
{
	HttpClient::Response Result = m_HttpClient.Get(fmt::format("v1/kv/{}?raw", Key));
	if (!Result)
	{
		throw runtime_error("ConsulClient::GetKeyValue() failed to get key '{}' ({})", Key, Result.ErrorMessage(""));
	}
	return Result.ToText();
}

void
ConsulClient::DeleteKey(std::string_view Key)
{
	HttpClient::Response Result = m_HttpClient.Delete(fmt::format("v1/kv/{}", Key));
	if (!Result)
	{
		throw runtime_error("ConsulClient::DeleteKey() failed to delete key '{}' ({})", Key, Result.ErrorMessage(""));
	}
}

bool
ConsulClient::RegisterService(const ServiceRegistrationInfo& Info)
{
	using namespace std::literals;

	CbObjectWriter Writer;
	{
		Writer.AddString("ID"sv, Info.ServiceId);
		Writer.AddString("Name"sv, Info.ServiceName);
		if (!Info.Address.empty())
		{
			Writer.AddString("Address"sv, Info.Address);
		}
		Writer.AddInteger("Port"sv, Info.Port);
		if (!Info.Tags.empty())
		{
			Writer.BeginArray("Tags"sv);
			for (const std::pair<std::string, std::string>& Tag : Info.Tags)
			{
				Writer.AddString(fmt::format("{}:{}", Tag.first, Tag.second));
			}
			Writer.EndArray();	// Tags
		}
		Writer.BeginObject("Check"sv);
		{
			Writer.AddString("HTTP"sv, fmt::format("http://{}:{}/{}", Info.Address, Info.Port, Info.HealthEndpoint));
			Writer.AddString("Interval"sv, fmt::format("{}s", Info.HealthIntervalSeconds));
			Writer.AddString("DeregisterCriticalServiceAfter"sv, fmt::format("{}s", Info.DeregisterAfterSeconds));
		}
		Writer.EndObject();	 // Check
	}

	ExtendableStringBuilder<512> SB;
	CompactBinaryToJson(Writer.Save(), SB);

	IoBuffer PayloadBuffer(IoBuffer::Wrap, SB.Data(), SB.Size());
	PayloadBuffer.SetContentType(HttpContentType::kJSON);
	HttpClient::Response Result = m_HttpClient.Put("v1/agent/service/register", PayloadBuffer, HttpClient::Accept(HttpContentType::kJSON));

	if (!Result)
	{
		ZEN_WARN("ConsulClient::RegisterService() failed to register service '{}' ({})", Info.ServiceId, Result.ErrorMessage(""));
		return false;
	}

	return true;
}

bool
ConsulClient::DeregisterService(std::string_view ServiceId)
{
	HttpClient::Response Result =
		m_HttpClient.Put(fmt::format("v1/agent/service/deregister/{}", ServiceId), HttpClient::Accept(HttpContentType::kJSON));

	if (!Result)
	{
		ZEN_WARN("ConsulClient::DeregisterService() failed to deregister service '{}' ({})", ServiceId, Result.ErrorMessage(""));
		return false;
	}

	return true;
}

bool
ConsulClient::FindServiceInJson(std::string_view Json, std::string_view ServiceId)
{
	using namespace std::literals;

	std::string		JsonError;
	CbFieldIterator Root = LoadCompactBinaryFromJson(Json, JsonError);
	if (Root && JsonError.empty())
	{
		for (CbFieldView RootIterator : Root)
		{
			CbObjectView ServiceObject = RootIterator.AsObjectView();
			if (ServiceObject)
			{
				for (CbFieldView ServiceIterator : ServiceObject)
				{
					CbObjectView ObjectIterator = ServiceIterator.AsObjectView();
					if (ObjectIterator["ID"sv].AsString() == ServiceId)
					{
						return true;
					}
				}
			}
		}
	}
	return false;
}

bool
ConsulClient::HasService(std::string_view ServiceId)
{
	HttpClient::Response Result = m_HttpClient.Get("v1/agent/services");
	if (!Result)
	{
		return false;
	}
	return FindServiceInJson(Result.AsText(), ServiceId);
}

bool
ConsulClient::WatchService(std::string_view ServiceId, uint64_t& InOutIndex, int WaitSeconds)
{
	// Note: m_HttpClient uses unlimited HTTP timeout (Timeout{0}); the WaitSeconds parameter
	// governs the server-side bound on the blocking query. Do not add a separate client timeout.
	HttpClient::KeyValueMap Parameters({{"index", std::to_string(InOutIndex)}, {"wait", fmt::format("{}s", WaitSeconds)}});
	HttpClient::Response	Result = m_HttpClient.Get("v1/agent/services", {}, Parameters);
	if (!Result)
	{
		return false;
	}

	auto IndexIt = Result.Header->find("X-Consul-Index");
	if (IndexIt != Result.Header->end())
	{
		std::optional<uint64_t> ParsedIndex = ParseInt<uint64_t>(IndexIt->second);
		if (ParsedIndex.has_value())
		{
			InOutIndex = ParsedIndex.value();
		}
	}

	return FindServiceInJson(Result.AsText(), ServiceId);
}

std::string
ConsulClient::GetAgentServicesJson()
{
	HttpClient::Response Result = m_HttpClient.Get("v1/agent/services");
	if (!Result)
	{
		return "{}";
	}
	return Result.ToText();
}

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

ServiceRegistration::ServiceRegistration(ConsulClient* Client, const ServiceRegistrationInfo& Info) : m_Client(Client), m_Info(Info)
{
	m_RegistrationThread = std::thread(&ServiceRegistration::RegistrationLoop, this);
}

ServiceRegistration::~ServiceRegistration()
{
	try
	{
		m_StopEvent.Set();

		if (m_RegistrationThread.joinable())
		{
			m_RegistrationThread.join();
		}

		if (m_IsRegistered.load())
		{
			if (!m_Client->DeregisterService(m_Info.ServiceId))
			{
				ZEN_WARN("ServiceRegistration: Failed to deregister service '{}' during cleanup", m_Info.ServiceId);
			}
		}
	}
	catch (const std::exception& Ex)
	{
		ZEN_ERROR(":~ServiceRegistration() threw exception: {}", Ex.what());
	}
}

void
ServiceRegistration::WaitForReadyEvent(int MaxWaitMS)
{
	if (m_IsRegistered.load())
	{
		return;
	}
	m_ReadyWakeupEvent.Wait(MaxWaitMS);
}

void
ServiceRegistration::RegistrationLoop()
{
	try
	{
		SetCurrentThreadName("ConsulRegistration");

		// Exponential backoff delays: 100ms, 200ms, 400ms
		constexpr int BackoffDelaysMs[] = {100, 200, 400};
		constexpr int MaxAttempts		= 3;
		constexpr int RetryIntervalMs	= 5000;

		while (true)
		{
			bool Succeeded = false;
			// Try to register with exponential backoff
			for (int Attempt = 0; Attempt < MaxAttempts; ++Attempt)
			{
				if (m_Client->RegisterService(m_Info))
				{
					Succeeded = true;
					break;
				}

				if (m_StopEvent.Wait(BackoffDelaysMs[Attempt]))
				{
					return;
				}
			}

			if (Succeeded || m_Client->RegisterService(m_Info))
			{
				break;
			}

			// All attempts failed, log warning and wait before retrying
			ZEN_WARN("ServiceRegistration: Failed to register service '{}' after {} attempts, retrying in {}s",
					 m_Info.ServiceId,
					 MaxAttempts + 1,
					 RetryIntervalMs / 1000);

			// Wait for 5 seconds before retrying; m_StopEvent.Set() in the destructor
			// wakes this immediately so the thread exits without polling.
			if (m_StopEvent.Wait(RetryIntervalMs))
			{
				return;
			}
		}

		ZEN_INFO("ServiceRegistration: Successfully registered service '{}'", m_Info.ServiceId);
		m_IsRegistered.store(true);
		m_ReadyWakeupEvent.Set();
	}
	catch (const std::exception& Ex)
	{
		ZEN_ERROR("ServiceRegistration::RegistrationLoop failed with exception: {}", Ex.what());
	}
}

}  // namespace zen::consul