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

#include "timeline_query.h"

#include <algorithm>

namespace zen::trace_detail {

namespace {

	// Pick the LOD level a given resolution should read from. Mirrors the
	// historical selection in trace_viewer_service.cpp: resolution 0 reads the
	// raw LOD 0; otherwise the smallest LOD whose ResolutionUs >= the request
	// wins, falling back to the coarsest level if none qualify.
	//
	// Returned values: 0 == raw scopes (LOD 0), 1..kTimelineLodCount == DetailLevels[lod-1].
	size_t SelectLodIndex(uint32_t ResolutionUs)
	{
		if (ResolutionUs == 0)
		{
			return 0;
		}
		for (size_t I = 0; I < kTimelineLodCount; ++I)
		{
			if (kTimelineLodResolutions[I] >= ResolutionUs)
			{
				return I + 1;
			}
		}
		return kTimelineLodCount;
	}

	const eastl::vector<TimelineScope>& LodScopes(const ThreadTimeline& Timeline, size_t LodIndex)
	{
		if (LodIndex == 0)
		{
			return Timeline.Scopes;
		}
		return Timeline.DetailLevels[LodIndex - 1].Scopes;
	}

	void ExtractScopesInto(const ThreadTimeline& Timeline, const TimelineQueryRequest& Req, std::vector<TimelineScopeView>& Out)
	{
		const eastl::vector<TimelineScope>& Scopes = LodScopes(Timeline, SelectLodIndex(Req.ResolutionUs));

		auto MidIt =
			std::lower_bound(Scopes.begin(), Scopes.end(), Req.StartUs, [](const TimelineScope& S, uint32_t V) { return S.BeginUs < V; });

		for (auto It = Scopes.begin(); It != MidIt; ++It)
		{
			if ((It->BeginUs + It->DurationUs) < Req.StartUs || It->DurationUs < Req.MinDurUs)
			{
				continue;
			}
			Out.push_back({It->BeginUs, It->DurationUs, It->NameId, It->Depth, It->MergeCount});
		}
		for (auto It = MidIt; It != Scopes.end(); ++It)
		{
			if (It->BeginUs > Req.EndUs)
			{
				break;
			}
			if (It->DurationUs < Req.MinDurUs)
			{
				continue;
			}
			Out.push_back({It->BeginUs, It->DurationUs, It->NameId, It->Depth, It->MergeCount});
		}
	}

	const ThreadTimeline* FindThread(const TraceModel& Model, uint32_t ThreadId)
	{
		auto It = std::find_if(Model.Timelines.begin(), Model.Timelines.end(), [ThreadId](const ThreadTimeline& T) {
			return T.ThreadId == ThreadId;
		});
		return (It != Model.Timelines.end()) ? &*It : nullptr;
	}

	class InMemoryTimelineQuery final : public TimelineQuery
	{
	public:
		explicit InMemoryTimelineQuery(const TraceModel& Model) : m_Model(Model) {}

		void QueryThread(uint32_t ThreadId, const TimelineQueryRequest& Req, std::vector<TimelineScopeView>& Out) const override
		{
			const ThreadTimeline* Timeline = FindThread(m_Model, ThreadId);
			if (Timeline)
			{
				ExtractScopesInto(*Timeline, Req, Out);
			}
		}

		void QueryBatch(std::span<const uint32_t> ThreadIds, const TimelineQueryRequest& Req, BatchResult& Out) const override
		{
			Out.Scopes.clear();
			Out.Ranges.clear();
			Out.Ranges.reserve(ThreadIds.size());

			for (uint32_t ThreadId : ThreadIds)
			{
				const uint32_t		  Begin	   = uint32_t(Out.Scopes.size());
				const ThreadTimeline* Timeline = FindThread(m_Model, ThreadId);
				if (Timeline)
				{
					ExtractScopesInto(*Timeline, Req, Out.Scopes);
				}
				Out.Ranges.push_back({Begin, uint32_t(Out.Scopes.size())});
			}
		}

	private:
		const TraceModel& m_Model;
	};

}  // namespace

std::unique_ptr<TimelineQuery>
MakeInMemoryTimelineQuery(const TraceModel& Model)
{
	return std::make_unique<InMemoryTimelineQuery>(Model);
}

}  // namespace zen::trace_detail