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

#include "httptrafficrecorder.h"

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

#include <chrono>
#include <filesystem>

namespace zen {

HttpTrafficRecorder::HttpTrafficRecorder(const std::filesystem::path& OutputDir, std::string_view ClientLabel, std::string_view TargetLabel)
: m_Log(logging::Get("proxy.record"))
, m_Dir(OutputDir)
, m_ClientLabel(ClientLabel)
, m_TargetLabel(TargetLabel)
{
	auto Now	  = std::chrono::system_clock::now();
	m_StartTimeMs = uint64_t(std::chrono::duration_cast<std::chrono::milliseconds>(Now.time_since_epoch()).count());

	std::error_code Ec;
	std::filesystem::create_directories(m_Dir, Ec);
	if (Ec)
	{
		ZEN_WARN("failed to create recording directory {} - {}", m_Dir.string(), Ec.message());
		return;
	}

	std::error_code ReqEc;
	m_RequestFile.Open(m_Dir / "request.bin", BasicFile::Mode::kTruncate, ReqEc);
	if (ReqEc)
	{
		ZEN_WARN("failed to open request.bin in {} - {}", m_Dir.string(), ReqEc.message());
		return;
	}

	std::error_code RespEc;
	m_ResponseFile.Open(m_Dir / "response.bin", BasicFile::Mode::kTruncate, RespEc);
	if (RespEc)
	{
		ZEN_WARN("failed to open response.bin in {} - {}", m_Dir.string(), RespEc.message());
		m_RequestFile.Close();
		return;
	}

	m_Valid = true;
	ZEN_DEBUG("recording started in {}", m_Dir.string());
}

HttpTrafficRecorder::~HttpTrafficRecorder()
{
	if (m_Valid && !m_Finalized)
	{
		Finalize(false, Oid::Zero);
	}
}

void
HttpTrafficRecorder::WriteRequest(const char* Data, size_t Length)
{
	if (!m_Valid)
	{
		return;
	}
	m_RequestFile.Write(Data, Length, m_RequestOffset);
	m_RequestOffset += Length;
}

void
HttpTrafficRecorder::WriteResponse(const char* Data, size_t Length)
{
	if (!m_Valid)
	{
		return;
	}
	m_ResponseFile.Write(Data, Length, m_ResponseOffset);
	m_ResponseOffset += Length;
}

void
HttpTrafficRecorder::OnMessageComplete(HttpTrafficInspector::Direction Dir,
									   std::string_view				   Method,
									   std::string_view				   Url,
									   uint16_t						   StatusCode,
									   int64_t /*ContentLength*/)
{
	if (!m_Valid)
	{
		return;
	}

	if (Dir == HttpTrafficInspector::Direction::Request)
	{
		// Record the request boundary. The request spans from m_CurrentRequestStart to m_RequestOffset.
		m_PendingReqOffset	= m_CurrentRequestStart;
		m_PendingReqSize	= m_RequestOffset - m_CurrentRequestStart;
		m_PendingMethod		= Method;
		m_PendingUrl		= Url;
		m_HasPendingRequest = true;

		// Advance start to current offset for the next request.
		m_CurrentRequestStart = m_RequestOffset;
	}
	else
	{
		// Response complete -- pair with pending request.
		RecordedEntry Entry;
		if (m_HasPendingRequest)
		{
			Entry.ReqOffset		= m_PendingReqOffset;
			Entry.ReqSize		= m_PendingReqSize;
			Entry.Method		= std::move(m_PendingMethod);
			Entry.Url			= std::move(m_PendingUrl);
			m_HasPendingRequest = false;
		}

		Entry.RespOffset = m_CurrentResponseStart;
		Entry.RespSize	 = m_ResponseOffset - m_CurrentResponseStart;
		Entry.Status	 = StatusCode;

		m_Entries.push_back(std::move(Entry));

		// Advance start to current offset for the next response.
		m_CurrentResponseStart = m_ResponseOffset;
	}
}

void
HttpTrafficRecorder::Finalize(bool WebSocket, const Oid& SessionId)
{
	if (!m_Valid || m_Finalized)
	{
		return;
	}
	m_Finalized = true;

	m_RequestFile.Close();
	m_ResponseFile.Close();

	// Write index.ucb as a CbObject.
	CbObjectWriter Cbo;

	Cbo << "client" << std::string_view(m_ClientLabel);
	Cbo << "target" << std::string_view(m_TargetLabel);
	Cbo << "startTime" << m_StartTimeMs;
	Cbo << "websocket" << WebSocket;
	if (SessionId != Oid::Zero)
	{
		std::string SessionIdStr = SessionId.ToString();
		Cbo << "sessionId" << std::string_view(SessionIdStr);
	}

	Cbo.BeginArray("entries");
	for (const RecordedEntry& Entry : m_Entries)
	{
		Cbo.BeginObject();
		Cbo << "reqOffset" << Entry.ReqOffset;
		Cbo << "reqSize" << Entry.ReqSize;
		Cbo << "respOffset" << Entry.RespOffset;
		Cbo << "respSize" << Entry.RespSize;
		Cbo << "method" << std::string_view(Entry.Method);
		Cbo << "url" << std::string_view(Entry.Url);
		Cbo << "status" << Entry.Status;
		Cbo.EndObject();
	}
	Cbo.EndArray();

	CbObject IndexObj = Cbo.Save();

	MemoryView		View = IndexObj.GetView();
	std::error_code Ec;
	BasicFile		IndexFile(m_Dir / "index.ucb", BasicFile::Mode::kTruncate, Ec);
	if (!Ec)
	{
		IndexFile.Write(View, 0, Ec);
		if (Ec)
		{
			ZEN_WARN("failed to write index.ucb in {} - {}", m_Dir.string(), Ec.message());
		}
		IndexFile.Close();
	}
	else
	{
		ZEN_WARN("failed to create index.ucb in {} - {}", m_Dir.string(), Ec.message());
	}

	ZEN_DEBUG("recording finalized in {} ({} entries, websocket: {})", m_Dir.string(), m_Entries.size(), WebSocket);
}

}  // namespace zen