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

#include "httptrafficinspector.h"

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

#include <charconv>

namespace zen {

// clang-format off
llhttp_settings_t HttpTrafficInspector::s_RequestSettings = []() {
	llhttp_settings_t S;
	llhttp_settings_init(&S);
	S.on_message_begin    = [](llhttp_t*) { return 0; };
	S.on_url              = [](llhttp_t* p, const char* Data, size_t Len) { return GetThis(p)->OnUrl(Data, Len); };
	S.on_status           = [](llhttp_t*, const char*, size_t) { return 0; };
	S.on_header_field     = [](llhttp_t* p, const char* Data, size_t Len) { return GetThis(p)->OnHeaderField(Data, Len); };
	S.on_header_value     = [](llhttp_t* p, const char* Data, size_t Len) { return GetThis(p)->OnHeaderValue(Data, Len); };
	S.on_headers_complete = [](llhttp_t* p) { return GetThis(p)->OnHeadersComplete(); };
	S.on_body             = [](llhttp_t*, const char*, size_t) { return 0; };
	S.on_message_complete = [](llhttp_t* p) { return GetThis(p)->OnMessageComplete(); };
	return S;
}();

llhttp_settings_t HttpTrafficInspector::s_ResponseSettings = []() {
	llhttp_settings_t S;
	llhttp_settings_init(&S);
	S.on_message_begin    = [](llhttp_t*) { return 0; };
	S.on_url              = [](llhttp_t*, const char*, size_t) { return 0; };
	S.on_status           = [](llhttp_t*, const char*, size_t) { return 0; };
	S.on_header_field     = [](llhttp_t* p, const char* Data, size_t Len) { return GetThis(p)->OnHeaderField(Data, Len); };
	S.on_header_value     = [](llhttp_t* p, const char* Data, size_t Len) { return GetThis(p)->OnHeaderValue(Data, Len); };
	S.on_headers_complete = [](llhttp_t* p) { return GetThis(p)->OnHeadersComplete(); };
	S.on_body             = [](llhttp_t*, const char*, size_t) { return 0; };
	S.on_message_complete = [](llhttp_t* p) { return GetThis(p)->OnMessageComplete(); };
	return S;
}();
// clang-format on

HttpTrafficInspector::HttpTrafficInspector(Direction Dir, std::string_view SessionLabel)
: m_Log(logging::Get("proxy.http"))
, m_Direction(Dir)
, m_SessionLabel(SessionLabel)
{
	llhttp_settings_t* Settings = (Dir == Direction::Request) ? &s_RequestSettings : &s_ResponseSettings;
	llhttp_init(&m_Parser, Dir == Direction::Request ? HTTP_REQUEST : HTTP_RESPONSE, Settings);
	m_Parser.data = this;
}

void
HttpTrafficInspector::Inspect(const char* Data, size_t Length)
{
	if (m_Disabled)
	{
		return;
	}

	llhttp_errno_t Err = llhttp_execute(&m_Parser, Data, Length);

	if (Err == HPE_PAUSED_UPGRADE)
	{
		if (m_Direction == Direction::Request)
		{
			ZEN_DEBUG("[{}] >> {} {} (upgrade to WebSocket)", m_SessionLabel, m_Method, m_Url);
		}
		else
		{
			ZEN_DEBUG("[{}] << {} (upgrade to WebSocket)", m_SessionLabel, m_StatusCode);
		}
		ResetMessageState();
		m_Upgraded.store(true, std::memory_order_relaxed);
		m_Disabled = true;
		return;
	}

	if (Err != HPE_OK)
	{
		ZEN_DEBUG("[{}] non-HTTP traffic detected ({}), disabling inspection", m_SessionLabel, llhttp_errno_name(Err));
		m_Disabled = true;
	}
}

int
HttpTrafficInspector::OnUrl(const char* Data, size_t Length)
{
	m_Url.append(Data, Length);
	return 0;
}

int
HttpTrafficInspector::OnHeaderField(const char* Data, size_t Length)
{
	m_CurrentHeaderField.assign(Data, Length);
	return 0;
}

int
HttpTrafficInspector::OnHeaderValue(const char* Data, size_t Length)
{
	if (m_CurrentHeaderField.size() == 14 && StrCaseCompare(m_CurrentHeaderField.c_str(), "Content-Length", 14) == 0)
	{
		int64_t Value = 0;
		std::from_chars(Data, Data + Length, Value);
		m_ContentLength = Value;
	}
	else if (!m_SessionIdCaptured && m_CurrentHeaderField.size() == 10 &&
			 StrCaseCompare(m_CurrentHeaderField.c_str(), "UE-Session", 10) == 0)
	{
		Oid Parsed;
		if (Oid::TryParse(std::string_view(Data, Length), Parsed))
		{
			m_SessionId			= Parsed;
			m_SessionIdCaptured = true;
		}
	}
	m_CurrentHeaderField.clear();
	return 0;
}

int
HttpTrafficInspector::OnHeadersComplete()
{
	if (m_Direction == Direction::Request)
	{
		m_Method = llhttp_method_name(static_cast<llhttp_method_t>(llhttp_get_method(&m_Parser)));
	}
	else
	{
		m_StatusCode = static_cast<uint16_t>(llhttp_get_status_code(&m_Parser));
	}
	return 0;
}

int
HttpTrafficInspector::OnMessageComplete()
{
	if (m_Direction == Direction::Request)
	{
		if (m_ContentLength >= 0)
		{
			ZEN_DEBUG("[{}] >> {} {} (content-length: {})", m_SessionLabel, m_Method, m_Url, m_ContentLength);
		}
		else
		{
			ZEN_DEBUG("[{}] >> {} {}", m_SessionLabel, m_Method, m_Url);
		}
	}
	else
	{
		if (m_ContentLength >= 0)
		{
			ZEN_DEBUG("[{}] << {} (content-length: {})", m_SessionLabel, m_StatusCode, m_ContentLength);
		}
		else
		{
			ZEN_DEBUG("[{}] << {}", m_SessionLabel, m_StatusCode);
		}
	}

	if (m_Observer)
	{
		m_Observer->OnMessageComplete(m_Direction, m_Method, m_Url, m_StatusCode, m_ContentLength);
	}

	m_MessageCount.fetch_add(1, std::memory_order_relaxed);
	ResetMessageState();
	return 0;
}

Oid
HttpTrafficInspector::GetSessionId() const
{
	return m_SessionId;
}

bool
HttpTrafficInspector::HasSessionId() const
{
	return m_SessionIdCaptured;
}

void
HttpTrafficInspector::ResetMessageState()
{
	m_Url.clear();
	m_Method.clear();
	m_StatusCode	= 0;
	m_ContentLength = -1;
	m_CurrentHeaderField.clear();
}

}  // namespace zen