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

#include "httpproxystats.h"

#include "tcpproxy.h"

#include <zencore/compactbinarybuilder.h>
#include <zencore/fmtutils.h>

#include <chrono>
#include <filesystem>

namespace zen {

HttpProxyStatsService::HttpProxyStatsService(const std::vector<std::unique_ptr<TcpProxyService>>& ProxyServices,
											 IHttpStatsService&									  StatsService,
											 std::string										  DefaultRecordDir)
: m_ProxyServices(ProxyServices)
, m_StatsService(StatsService)
, m_DefaultRecordDir(std::move(DefaultRecordDir))
{
	m_StatsService.RegisterHandler("proxy", *this);
}

HttpProxyStatsService::~HttpProxyStatsService()
{
	m_StatsService.UnregisterHandler("proxy", *this);
}

const char*
HttpProxyStatsService::BaseUri() const
{
	return "/proxy/";
}

void
HttpProxyStatsService::HandleRequest(HttpServerRequest& Request)
{
	std::string_view Uri = Request.RelativeUri();

	if (Uri == "stats" || Uri == "stats/")
	{
		HandleStatsRequest(Request);
	}
	else if (Uri == "record/start" || Uri == "record/start/")
	{
		HandleRecordStart(Request);
	}
	else if (Uri == "record/stop" || Uri == "record/stop/")
	{
		HandleRecordStop(Request);
	}
	else if (Uri == "record" || Uri == "record/")
	{
		HandleRecordStatus(Request);
	}
	else
	{
		Request.WriteResponse(HttpResponseCode::NotFound);
	}
}

void
HttpProxyStatsService::HandleRecordStart(HttpServerRequest& Request)
{
	if (Request.RequestVerb() != HttpVerb::kPost)
	{
		Request.WriteResponse(HttpResponseCode::MethodNotAllowed);
		return;
	}

	auto			 Params = Request.GetQueryParams();
	std::string_view Dir	= Params.GetValue("dir");

	std::string RecordDir;
	if (Dir.empty())
	{
		RecordDir = m_DefaultRecordDir;
	}
	else
	{
		// Treat dir as a subdirectory name within the default record directory.
		// Reject path separators and parent references to prevent path traversal.
		if (Dir.find("..") != std::string_view::npos || Dir.find('/') != std::string_view::npos || Dir.find('\\') != std::string_view::npos)
		{
			Request.WriteResponse(HttpResponseCode::BadRequest);
			return;
		}
		RecordDir = (std::filesystem::path(m_DefaultRecordDir) / std::string(Dir)).string();
	}

	for (const std::unique_ptr<TcpProxyService>& Service : m_ProxyServices)
	{
		Service->SetRecording(true, RecordDir);
	}

	CbObjectWriter Cbo;
	Cbo << "recording" << true;
	Cbo << "dir" << std::string_view(RecordDir);
	Request.WriteResponse(HttpResponseCode::OK, Cbo.Save());
}

void
HttpProxyStatsService::HandleRecordStop(HttpServerRequest& Request)
{
	if (Request.RequestVerb() != HttpVerb::kPost)
	{
		Request.WriteResponse(HttpResponseCode::MethodNotAllowed);
		return;
	}

	for (const std::unique_ptr<TcpProxyService>& Service : m_ProxyServices)
	{
		Service->SetRecording(false, Service->GetRecordDir());
	}

	CbObjectWriter Cbo;
	Cbo << "recording" << false;
	Request.WriteResponse(HttpResponseCode::OK, Cbo.Save());
}

void
HttpProxyStatsService::HandleRecordStatus(HttpServerRequest& Request)
{
	bool		IsRecording = false;
	std::string RecordDir;
	for (const std::unique_ptr<TcpProxyService>& Service : m_ProxyServices)
	{
		if (Service->IsRecording())
		{
			IsRecording = true;
			RecordDir	= Service->GetRecordDir();
			break;
		}
	}

	CbObjectWriter Cbo;
	Cbo << "recording" << IsRecording;
	Cbo << "dir" << std::string_view(RecordDir);
	Request.WriteResponse(HttpResponseCode::OK, Cbo.Save());
}

CbObject
HttpProxyStatsService::CollectStats()
{
	CbObjectWriter Cbo;

	// Include recording status in stats output.
	{
		bool		IsRecording = false;
		std::string RecordDir;
		for (const std::unique_ptr<TcpProxyService>& Service : m_ProxyServices)
		{
			if (Service->IsRecording())
			{
				IsRecording = true;
				RecordDir	= Service->GetRecordDir();
				break;
			}
		}
		Cbo << "recording" << IsRecording;
		Cbo << "recordDir" << std::string_view(RecordDir);
	}

	Cbo.BeginArray("mappings");
	for (const std::unique_ptr<TcpProxyService>& Service : m_ProxyServices)
	{
		const ProxyMapping& Mapping = Service->GetMapping();

		Cbo.BeginObject();
		{
			std::string ListenAddr = Mapping.ListenDescription();
			Cbo << "listen" << std::string_view(ListenAddr);

			std::string TargetAddr = Mapping.TargetDescription();
			Cbo << "target" << std::string_view(TargetAddr);

			Cbo << "activeConnections" << Service->GetActiveConnections();
			Cbo << "peakActiveConnections" << Service->GetPeakActiveConnections();
			Cbo << "totalConnections" << Service->GetTotalConnections();
			Cbo << "bytesFromClient" << Service->GetTotalBytesFromClient();
			Cbo << "bytesToClient" << Service->GetTotalBytesToClient();

			Cbo << "requestRate1" << Service->GetRequestMeter().Rate1();
			Cbo << "byteRate1" << Service->GetBytesMeter().Rate1();
			Cbo << "byteRate5" << Service->GetBytesMeter().Rate5();

			auto Now	  = std::chrono::steady_clock::now();
			auto Sessions = Service->GetActiveSessions();

			Cbo.BeginArray("connections");
			for (const auto& Session : Sessions)
			{
				Cbo.BeginObject();
				{
					std::string ClientLabel = Session->GetClientLabel();
					Cbo << "client" << std::string_view(ClientLabel);

					std::string TargetLabel = Mapping.TargetDescription();
					Cbo << "target" << std::string_view(TargetLabel);

					Cbo << "bytesFromClient" << Session->GetBytesFromClient();
					Cbo << "bytesToClient" << Session->GetBytesToClient();

					Cbo << "requests" << Session->GetRequestCount();
					Cbo << "websocket" << Session->IsWebSocket();

					if (Session->HasSessionId())
					{
						std::string SessionId = Session->GetSessionId().ToString();
						Cbo << "sessionId" << std::string_view(SessionId);
					}

					double DurationMs = std::chrono::duration<double, std::milli>(Now - Session->GetStartTime()).count();
					Cbo << "durationMs" << DurationMs;
				}
				Cbo.EndObject();
			}
			Cbo.EndArray();
		}
		Cbo.EndObject();
	}
	Cbo.EndArray();

	return Cbo.Save();
}

void
HttpProxyStatsService::HandleStatsRequest(HttpServerRequest& Request)
{
	Request.WriteResponse(HttpResponseCode::OK, CollectStats());
}

}  // namespace zen