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

#include "statsreporter.h"

#include <zencore/logging.h>
#include <zencore/profiling/counterstrace.h>
#include <zencore/trace.h>
#include <zennet/statsdclient.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <EASTL/functional.h>
#include <EASTL/hash_map.h>
ZEN_THIRD_PARTY_INCLUDES_END

#include <memory>
#include <string>

namespace zen {

// Decorator that forwards every metric call to an inner StatsMetrics (statsd
// in production) and additionally emits a counter trace point so the metric
// shows up in the .utrace stream / `zen trace serve` viewer.
//
// Counter ids are allocated lazily on first sighting of a metric name. The
// name string is owned by the lookup map so the TraceCounter keeps a stable
// pointer to it for late-init re-emission.
class TracingStatsMetrics : public StatsMetrics
{
public:
	void Rebind(StatsMetrics* Inner) { m_Inner = Inner; }

	void Increment(std::string_view Metric) override
	{
		if (m_Inner)
			m_Inner->Increment(Metric);
		Counter(Metric).Increment();
	}

	void Decrement(std::string_view Metric) override
	{
		if (m_Inner)
			m_Inner->Decrement(Metric);
		Counter(Metric).Decrement();
	}

	void Count(std::string_view Metric, int64_t CountDelta) override
	{
		if (m_Inner)
			m_Inner->Count(Metric, CountDelta);
		Counter(Metric).Add(CountDelta);
	}

	void Gauge(std::string_view Metric, uint64_t CurrentValue) override
	{
		if (m_Inner)
			m_Inner->Gauge(Metric, CurrentValue);
		Counter(Metric).SetAlways(int64_t(CurrentValue));
	}

	void Meter(std::string_view Metric, uint64_t IncrementValue) override
	{
		if (m_Inner)
			m_Inner->Meter(Metric, IncrementValue);
		Counter(Metric).Add(int64_t(IncrementValue));
	}

private:
	TraceCounterInt& Counter(std::string_view Metric)
	{
		if (auto It = m_Counters.find_as(Metric, std::hash<std::string_view>(), eastl::equal_to_2<std::string, std::string_view>());
			It != m_Counters.end())
		{
			return *It->second;
		}
		auto [Inserted, _] = m_Counters.try_emplace(std::string(Metric), nullptr);
		Inserted->second   = std::make_unique<TraceCounterInt>(Inserted->first.c_str(), TraceCounterDisplayHint::None);
		return *Inserted->second;
	}

	StatsMetrics*																									   m_Inner = nullptr;
	eastl::hash_map<std::string, std::unique_ptr<TraceCounterInt>, std::hash<std::string>, std::equal_to<std::string>> m_Counters;
};

StatsReporter::StatsReporter() : m_TracingMetrics(std::make_unique<TracingStatsMetrics>())
{
}

StatsReporter::~StatsReporter()
{
}

void
StatsReporter::Initialize(const ZenStatsConfig& Config)
{
	ZEN_TRACE_CPU("StatsReporter::Initialize");
	RwLock::ExclusiveLockScope _(m_Lock);

	if (Config.Enabled)
	{
		ZEN_INFO("initializing stats reporter: {}:{}", Config.StatsdHost, Config.StatsdPort)
		m_Statsd = CreateStatsDaemonClient(Config.StatsdHost, gsl::narrow<uint16_t>(Config.StatsdPort));
		m_Statsd->SetMessageSize(1500, false);
	}
}

void
StatsReporter::Shutdown()
{
	RwLock::ExclusiveLockScope _(m_Lock);
	m_Statsd.reset();
}

void
StatsReporter::AddProvider(StatsProvider* Provider)
{
	if (!Provider)
	{
		return;
	}
	RwLock::ExclusiveLockScope _(m_Lock);
	m_Providers.push_back(Provider);
}

void
StatsReporter::ReportStats()
{
	RwLock::ExclusiveLockScope _(m_Lock);

	// Always run providers through the tracing decorator so counter trace
	// points fire even when statsd is disabled. The decorator no-ops the
	// inner forward when m_Statsd is null.
	m_TracingMetrics->Rebind(m_Statsd.get());
	for (StatsProvider* Provider : m_Providers)
	{
		Provider->ReportMetrics(*m_TracingMetrics);
	}

	if (m_Statsd)
	{
		m_Statsd->Flush();
	}
}

}  // namespace zen