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

#include "zentelemetry/otlptrace.h"
#include "oteltraceprotozero.h"

#if ZEN_WITH_OTEL

#	include <zencore/endian.h>
#	include <zencore/memory/memoryarena.h>
#	include <zencore/scopeguard.h>
#	include <zencore/session.h>
#	include <zencore/testing.h>
#	include <zencore/uid.h>

#	include <EASTL/list.h>
#	include <EASTL/vector.h>

#	include <protozero/pbf_builder.hpp>
#	include <random>

namespace zen::otel {

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

uint64_t
NowInNanoseconds()
{
	const uint64_t NowNanos =
		std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
	return static_cast<uint64_t>(NowNanos);
}

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

/**
 * xorshiftr128+ random number generator
 *
 * https://en.wikipedia.org/wiki/Xorshift#xorshift+
 */
class XorShiftr128p
{
public:
	XorShiftr128p()
	{
		std::seed_seq SeedSeq{1, 2, 3, 5};
		do
		{
			SeedSeq.generate(reinterpret_cast<uint32_t*>(m_State), reinterpret_cast<uint32_t*>(&m_State[2]));
		} while (m_State[0] == 0 && m_State[1] == 0);
	}

	uint64_t Next()
	{
		uint64_t	   x = m_State[0];
		uint64_t const y = m_State[1];
		m_State[0]		 = y;
		x ^= x << 23;
		x ^= x >> 17;
		x ^= y;
		m_State[1] = x + y;
		return x;
	}

private:
	uint64_t m_State[2];
};

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

struct TraceState
{
	Span*	ActiveSpan = nullptr;
	TraceId CurrentTraceId;
};

thread_local TraceState t_TraceState;

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

Span::Span(MemoryArena& Arena, std::string_view Name, Span* SpanChain)
: m_Arena(Arena)
, m_SpanChain(SpanChain)
, m_SpanId(SpanId::NewSpanId())
, m_Name(m_Arena.DuplicateString(Name))
{
	if (Span* ActiveSpan = t_TraceState.ActiveSpan)
	{
		m_ParentSpan			 = ActiveSpan;
		m_NextSibling			 = ActiveSpan->m_FirstChild;
		ActiveSpan->m_FirstChild = this;
	}
	else
	{
		m_ParentSpan				= nullptr;
		t_TraceState.CurrentTraceId = TraceId::NewTraceId();
	}

	t_TraceState.ActiveSpan = this;
	m_StartTime				= NowInNanoseconds();
}

Span::~Span()
{
	End();
}

void
Span::End()
{
	if (m_EndTime)
	{
		return;
	}

	ZEN_ASSERT_SLOW(t_TraceState.ActiveSpan == this);

	m_EndTime				= NowInNanoseconds();
	t_TraceState.ActiveSpan = m_ParentSpan;
}

Span*
Span::GetCurrentSpan()
{
	return t_TraceState.ActiveSpan;
}

SpanId
Span::GetCurrentSpanId(TraceId& OutTraceId)
{
	if (const auto& State = t_TraceState; State.ActiveSpan)
	{
		OutTraceId = State.CurrentTraceId;
		return State.ActiveSpan->m_SpanId;
	}
	else
	{
		OutTraceId = TraceId();
		return SpanId();
	}
}

void
Span::AddEvent(std::string_view Name)
{
	Event* NewEvent		= new (m_Arena) Event();
	NewEvent->Name		= m_Arena.DuplicateString(Name);
	NewEvent->Timestamp = NowInNanoseconds();
	NewEvent->NextEvent = m_Events;
	m_Events			= NewEvent;
}

void
Span::AddEvent(std::string_view Name, uint64_t Timestamp)
{
	Event* NewEvent		= new (m_Arena) Event();
	NewEvent->Name		= m_Arena.DuplicateString(Name);
	NewEvent->Timestamp = Timestamp;
	NewEvent->NextEvent = m_Events;
	m_Events			= NewEvent;
}

void
Span::AddEvent(std::string_view Name, const AttributeList& Attributes)
{
	Event* NewEvent		= new (m_Arena) Event();
	NewEvent->Name		= m_Arena.DuplicateString(Name);
	NewEvent->Timestamp = NowInNanoseconds();
	NewEvent->NextEvent = m_Events;
	m_Events			= NewEvent;

	for (auto& [K, V] : Attributes)
	{
		AttributePair* NewAttr = new (m_Arena) AttributePair();
		NewAttr->Key		   = m_Arena.DuplicateString(K);
		NewAttr->SetStringValue(m_Arena.DuplicateString(V));
		NewAttr->Next		 = NewEvent->Attributes;
		NewEvent->Attributes = NewAttr;
	}
}

void
Span::AddAttribute(std::string_view Key, std::string_view Value)
{
	AttributePair* NewAttr = new (m_Arena) AttributePair();
	NewAttr->Key		   = m_Arena.DuplicateString(Key);
	NewAttr->SetStringValue(m_Arena.DuplicateString(Value));
	NewAttr->Next = m_Attributes;
	m_Attributes  = NewAttr;
}

void
Span::AddAttribute(std::string_view Key, uint64_t Value)
{
	AttributePair* NewAttr = new (m_Arena) AttributePair();
	NewAttr->Key		   = m_Arena.DuplicateString(Key);
	NewAttr->SetNumericValue(Value);
	NewAttr->Next = m_Attributes;
	m_Attributes  = NewAttr;
}

void
Span::AddAttributes(const AttributeList& Attributes)
{
	for (const auto& [K, V] : Attributes)
	{
		AddAttribute(K, V);
	}
}

void*
Span::operator new(size_t Size, zen::MemoryArena& Arena)
{
	return Arena.Allocate(Size);
}

void
Span::operator delete(void* Ptr, zen::MemoryArena& Arena)
{
	// This exists because operator new with placement arguments
	// requires a matching operator delete, but we don't actually
	// need to do anything here
	ZEN_UNUSED(Ptr, Arena);
}

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

std::atomic<uint32_t> g_TraceCounter;

TraceId
TraceId::NewTraceId()
{
	uint8_t TraceIdBytes[kSize];

	// We use our session ID as the basis for trace IDs
	const Oid SessionId = GetSessionId();
	memcpy(TraceIdBytes, SessionId.OidBits, sizeof SessionId.OidBits);

	const uint32_t TraceCounterBe = ByteSwap(g_TraceCounter.fetch_add(1));

	memcpy(&TraceIdBytes[12], &TraceCounterBe, 4);

	return TraceId(std::span<const uint8_t, kSize>(TraceIdBytes, kSize));
}

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

SpanId
SpanId::NewSpanId()
{
	// Just use a new OID and take the first 8 bytes. We'll probably want
	// a more native solution later

	return SpanId(Oid::NewOid());
}

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

Ref<TraceRecorder> g_TraceRecorder;

void
SetTraceRecorder(Ref<TraceRecorder> Recorder)
{
	g_TraceRecorder = std::move(Recorder);
}

bool
IsRecording()
{
	return !!g_TraceRecorder;
}

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

std::atomic<bool> g_OtlpTraceEnabled{false};

bool
IsOtlpTraceEnabled()
{
	return g_OtlpTraceEnabled.load();
}

thread_local Tracer* t_Tracer;

struct Tracer::Impl
{
	Impl() : m_PrevTracer(t_Tracer) {}
	~Impl() { t_Tracer = m_PrevTracer; }

	MemoryArena			  m_Arena;
	Tracer*				  m_PrevTracer;
	Span*				  m_SpanChain = nullptr;
	std::atomic<uint64_t> m_SpanCount;
};

Tracer::Tracer() : m_Impl(new Impl())
{
	t_Tracer = this;
}

Tracer::~Tracer()
{
	if (Ref<TraceRecorder> Recorder = g_TraceRecorder)
	{
		const uint64_t SpanCount = m_Impl->m_SpanCount.load();
		const Span**   Spans	 = new (m_Impl->m_Arena) const Span*[SpanCount];

		uint64_t Index = 0;

		for (Span* CurrentSpan = m_Impl->m_SpanChain; CurrentSpan != nullptr; CurrentSpan = CurrentSpan->m_SpanChain)
		{
			Spans[Index++] = CurrentSpan;
		}

		Recorder->RecordSpans(t_TraceState.CurrentTraceId, std::span<const Span*>(Spans, SpanCount));
	}

	delete m_Impl;
}

Tracer*
Tracer::GetTracer()
{
	Tracer* TracerPtr = t_Tracer;

	if (!TracerPtr)
	{
		TracerPtr = new Tracer();
	}

	return TracerPtr;
}

ScopedSpan
Tracer::CreateSpan(std::string_view Name)
{
	return ScopedSpan(Name);
}

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

ScopedSpan::ScopedSpan(std::string_view Name)
{
	if (!IsOtlpTraceEnabled())
	{
		// m_Tracer and m_Span remain null
		return;
	}

	Initialize(Name);
}

void
ScopedSpan::Initialize(std::string_view Name)
{
	Tracer* TracerPtr = Tracer::GetTracer();

	Tracer::Impl* const ImplPtr = TracerPtr->m_Impl;

	Span* NewSpan = new (ImplPtr->m_Arena) Span(ImplPtr->m_Arena, Name, ImplPtr->m_SpanChain);

	ImplPtr->m_SpanChain = NewSpan;
	ImplPtr->m_SpanCount.fetch_add(1);

	m_Tracer = TracerPtr;
	m_Span	 = NewSpan;
}

ScopedSpan::~ScopedSpan()
{
	// this is not inline to avoid code bloat on every use site
}

}  // namespace zen::otel

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

namespace zen::otel {

using namespace std::literals;

void
otlptrace_forcelink()
{
}

#	if ZEN_WITH_TESTS

TEST_CASE("otlp.trace")
{
	// Enable OTLP tracing for the duration of this test
	auto _ = MakeGuard([PreviousState = zen::otel::g_OtlpTraceEnabled.load()]() { zen::otel::g_OtlpTraceEnabled.store(PreviousState); });
	zen::otel::g_OtlpTraceEnabled.store(true);

	{
		ScopedSpan Span0("span0", [](Span& S) {
			S.AddAttribute("attr1"sv, "value1"sv);
			S.AddAttribute("attr2"sv, 42ULL);
			S.AddEvent("TestEvent1"sv);
			S.AddEvent("TestEvent2"sv);
		});

		{
			ScopedSpan Span1 = Tracer::CreateSpan("span1");
			Span1.WithSpan([](Span& S) {
				S.AddEvent("TestEvent3"sv);
				S.AddEvent("TestEvent4"sv);
			});
		}
	}
}

#	endif

}  // namespace zen::otel
#endif