#include "RuntimeUpdateController.h"

#include "RenderEngine.h"
#include "RuntimeEventDispatcher.h"
#include "RuntimeServices.h"
#include "RuntimeStore.h"
#include "ShaderBuildQueue.h"
#include "VideoBackend.h"

#include <variant>

namespace
{
RuntimeCoordinatorRenderResetScope ToRuntimeCoordinatorRenderResetScope(RuntimeEventRenderResetScope scope)
{
	switch (scope)
	{
	case RuntimeEventRenderResetScope::TemporalHistoryOnly:
		return RuntimeCoordinatorRenderResetScope::TemporalHistoryOnly;
	case RuntimeEventRenderResetScope::TemporalHistoryAndFeedback:
		return RuntimeCoordinatorRenderResetScope::TemporalHistoryAndFeedback;
	case RuntimeEventRenderResetScope::None:
	default:
		return RuntimeCoordinatorRenderResetScope::None;
	}
}
}

RuntimeUpdateController::RuntimeUpdateController(
	RuntimeStore& runtimeStore,
	RuntimeCoordinator& runtimeCoordinator,
	RuntimeEventDispatcher& runtimeEventDispatcher,
	RuntimeServices& runtimeServices,
	RenderEngine& renderEngine,
	ShaderBuildQueue& shaderBuildQueue,
	VideoBackend& videoBackend) :
	mRuntimeStore(runtimeStore),
	mRuntimeCoordinator(runtimeCoordinator),
	mRuntimeEventDispatcher(runtimeEventDispatcher),
	mRuntimeServices(runtimeServices),
	mRenderEngine(renderEngine),
	mShaderBuildQueue(shaderBuildQueue),
	mVideoBackend(videoBackend)
{
	mRuntimeEventDispatcher.Subscribe(
		RuntimeEventType::RuntimeStateBroadcastRequested,
		[this](const RuntimeEvent& event) { HandleRuntimeStateBroadcastRequested(event); });
	mRuntimeEventDispatcher.Subscribe(
		RuntimeEventType::RuntimeReloadRequested,
		[this](const RuntimeEvent& event) { HandleRuntimeReloadRequested(event); });
	mRuntimeEventDispatcher.Subscribe(
		RuntimeEventType::RuntimePersistenceRequested,
		[this](const RuntimeEvent& event) { HandleRuntimePersistenceRequested(event); });
	mRuntimeEventDispatcher.Subscribe(
		RuntimeEventType::ShaderBuildRequested,
		[this](const RuntimeEvent& event) { HandleShaderBuildRequested(event); });
	mRuntimeEventDispatcher.Subscribe(
		RuntimeEventType::ShaderBuildPrepared,
		[this](const RuntimeEvent& event) { HandleShaderBuildPrepared(event); });
	mRuntimeEventDispatcher.Subscribe(
		RuntimeEventType::ShaderBuildFailed,
		[this](const RuntimeEvent& event) { HandleShaderBuildFailed(event); });
	mRuntimeEventDispatcher.Subscribe(
		RuntimeEventType::CompileStatusChanged,
		[this](const RuntimeEvent& event) { HandleCompileStatusChanged(event); });
	mRuntimeEventDispatcher.Subscribe(
		RuntimeEventType::RenderResetRequested,
		[this](const RuntimeEvent& event) { HandleRenderResetRequested(event); });
}

bool RuntimeUpdateController::ApplyRuntimeCoordinatorResult(const RuntimeCoordinatorResult& result, std::string* error)
{
	if (!result.accepted)
	{
		if (error)
			*error = result.errorMessage;
		return false;
	}

	if (result.compileStatusChanged)
	{
		mRuntimeStore.SetCompileStatus(result.compileStatusSucceeded, result.compileStatusMessage);
		++mPendingCoordinatorCompileStatusEvents;
	}

	if (result.clearReloadRequest)
		mRuntimeStore.ClearReloadRequest();

	mRuntimeCoordinator.ApplyCommittedStateMode(result.committedStateMode);

	switch (result.transientOscInvalidation)
	{
	case RuntimeCoordinatorTransientOscInvalidation::All:
		mRenderEngine.ClearOscOverlayState();
		mRuntimeServices.ClearOscState();
		break;
	case RuntimeCoordinatorTransientOscInvalidation::Layer:
		mRenderEngine.ClearOscOverlayStateForLayerKey(result.transientOscLayerKey);
		mRuntimeServices.ClearOscStateForLayerKey(result.transientOscLayerKey);
		break;
	case RuntimeCoordinatorTransientOscInvalidation::None:
	default:
		break;
	}

	mRenderEngine.ApplyRuntimeCoordinatorRenderReset(result.renderResetScope);
	if (result.renderResetScope != RuntimeCoordinatorRenderResetScope::None)
		++mPendingCoordinatorRenderResetEvents;

	if (result.shaderBuildRequested)
	{
		RequestShaderBuild();
		++mPendingCoordinatorShaderBuildEvents;
	}

	if (result.runtimeStateBroadcastRequired)
		BroadcastRuntimeState();

	return true;
}

bool RuntimeUpdateController::ProcessRuntimeWork()
{
	DispatchRuntimeEvents();

	return ConsumeReadyShaderBuild(0, true, true);
}

void RuntimeUpdateController::RequestShaderBuild()
{
	mShaderBuildQueue.RequestBuild(mVideoBackend.InputFrameWidth(), mVideoBackend.InputFrameHeight());
}

void RuntimeUpdateController::BroadcastRuntimeState()
{
	RuntimeStateBroadcastRequestedEvent event;
	event.reason = "runtime-state-changed";
	if (!mRuntimeEventDispatcher.PublishPayload(event, "RuntimeUpdateController"))
	{
		mRuntimeServices.BroadcastState();
		return;
	}

	DispatchRuntimeEvents();
}

void RuntimeUpdateController::HandleRuntimeStateBroadcastRequested(const RuntimeEvent& event)
{
	if (event.source == "ControlServices")
		return;

	mRuntimeServices.BroadcastState();
}

void RuntimeUpdateController::HandleRuntimeReloadRequested(const RuntimeEvent& event)
{
	const RuntimeReloadRequestedEvent* payload = std::get_if<RuntimeReloadRequestedEvent>(&event.payload);
	if (!payload)
		return;

	mRuntimeStore.ClearReloadRequest();
}

void RuntimeUpdateController::HandleRuntimePersistenceRequested(const RuntimeEvent& event)
{
	const RuntimePersistenceRequestedEvent* payload = std::get_if<RuntimePersistenceRequestedEvent>(&event.payload);
	if (!payload)
		return;

	std::string error;
	mRuntimeStore.RequestPersistence(payload->request, error);
}

void RuntimeUpdateController::HandleShaderBuildRequested(const RuntimeEvent& event)
{
	const ShaderBuildEvent* payload = std::get_if<ShaderBuildEvent>(&event.payload);
	if (!payload || payload->phase != RuntimeEventShaderBuildPhase::Requested)
		return;
	if (ShouldSuppressCoordinatorFollowUp(event, mPendingCoordinatorShaderBuildEvents))
		return;

	RequestShaderBuild();
}

void RuntimeUpdateController::HandleShaderBuildPrepared(const RuntimeEvent& event)
{
	const ShaderBuildEvent* payload = std::get_if<ShaderBuildEvent>(&event.payload);
	if (!payload || payload->phase != RuntimeEventShaderBuildPhase::Prepared)
		return;

	ConsumeReadyShaderBuild(payload->generation, false, true);
}

void RuntimeUpdateController::HandleShaderBuildFailed(const RuntimeEvent& event)
{
	const ShaderBuildEvent* payload = std::get_if<ShaderBuildEvent>(&event.payload);
	if (!payload || payload->phase != RuntimeEventShaderBuildPhase::Failed)
		return;

	ConsumeReadyShaderBuild(payload->generation, false, false);
}

void RuntimeUpdateController::HandleCompileStatusChanged(const RuntimeEvent& event)
{
	const CompileStatusChangedEvent* payload = std::get_if<CompileStatusChangedEvent>(&event.payload);
	if (!payload)
		return;
	if (ShouldSuppressCoordinatorFollowUp(event, mPendingCoordinatorCompileStatusEvents))
		return;

	mRuntimeStore.SetCompileStatus(payload->succeeded, payload->message);
}

void RuntimeUpdateController::HandleRenderResetRequested(const RuntimeEvent& event)
{
	const RenderResetEvent* payload = std::get_if<RenderResetEvent>(&event.payload);
	if (!payload || payload->applied)
		return;
	if (ShouldSuppressCoordinatorFollowUp(event, mPendingCoordinatorRenderResetEvents))
		return;

	mRenderEngine.ApplyRuntimeCoordinatorRenderReset(ToRuntimeCoordinatorRenderResetScope(payload->scope));
}

bool RuntimeUpdateController::ConsumeReadyShaderBuild(uint64_t expectedGeneration, bool publishPreparedEvent, bool publishFailureEvent)
{
	PreparedShaderBuild readyBuild;
	const bool consumed = expectedGeneration == 0
		? mShaderBuildQueue.TryConsumeReadyBuild(readyBuild)
		: mShaderBuildQueue.TryConsumeReadyBuild(expectedGeneration, readyBuild);
	if (!consumed)
		return true;

	const unsigned inputWidth = mVideoBackend.InputFrameWidth();
	const unsigned inputHeight = mVideoBackend.InputFrameHeight();
	if (!readyBuild.succeeded)
	{
		if (publishFailureEvent)
		{
			PublishShaderBuildLifecycleEvent(
				RuntimeEventShaderBuildPhase::Failed,
				readyBuild.generation,
				inputWidth,
				inputHeight,
				false,
				readyBuild.message);
			DispatchRuntimeEvents();
		}
		ApplyRuntimeCoordinatorResult(mRuntimeCoordinator.HandlePreparedShaderBuildFailure(readyBuild.message));
		return false;
	}

	if (publishPreparedEvent)
	{
		PublishShaderBuildLifecycleEvent(
			RuntimeEventShaderBuildPhase::Prepared,
			readyBuild.generation,
			inputWidth,
			inputHeight,
			true,
			readyBuild.message);
		DispatchRuntimeEvents();
	}

	char compilerErrorMessage[1024] = {};
	if (!mRenderEngine.ApplyPreparedShaderBuild(
		readyBuild,
		inputWidth,
		inputHeight,
		mRuntimeCoordinator.PreserveFeedbackOnNextShaderBuild(),
		sizeof(compilerErrorMessage),
		compilerErrorMessage))
	{
		const std::string errorMessage = compilerErrorMessage;
		if (publishFailureEvent)
		{
			PublishShaderBuildLifecycleEvent(
				RuntimeEventShaderBuildPhase::Failed,
				readyBuild.generation,
				inputWidth,
				inputHeight,
				false,
				errorMessage);
			DispatchRuntimeEvents();
		}
		ApplyRuntimeCoordinatorResult(mRuntimeCoordinator.HandlePreparedShaderBuildFailure(errorMessage));
		return false;
	}

	PublishShaderBuildLifecycleEvent(
		RuntimeEventShaderBuildPhase::Applied,
		readyBuild.generation,
		inputWidth,
		inputHeight,
		true,
		"Shader layers applied successfully.");
	ApplyRuntimeCoordinatorResult(mRuntimeCoordinator.HandlePreparedShaderBuildSuccess());
	return true;
}

void RuntimeUpdateController::PublishShaderBuildLifecycleEvent(
	RuntimeEventShaderBuildPhase phase,
	uint64_t generation,
	unsigned inputWidth,
	unsigned inputHeight,
	bool succeeded,
	const std::string& message)
{
	ShaderBuildEvent event;
	event.phase = phase;
	event.generation = generation;
	event.inputWidth = inputWidth;
	event.inputHeight = inputHeight;
	event.preserveFeedbackState = mRuntimeCoordinator.PreserveFeedbackOnNextShaderBuild();
	event.succeeded = succeeded;
	event.message = message;
	mRuntimeEventDispatcher.PublishPayload(event, "RuntimeUpdateController");
}

bool RuntimeUpdateController::ShouldSuppressCoordinatorFollowUp(const RuntimeEvent& event, std::size_t& pendingSuppressions)
{
	if (event.source != "RuntimeCoordinator")
		return false;

	if (pendingSuppressions > 0)
		--pendingSuppressions;
	return true;
}

RuntimeEventDispatchResult RuntimeUpdateController::DispatchRuntimeEvents(std::size_t maxEvents)
{
	RuntimeEventDispatchResult result = mRuntimeEventDispatcher.DispatchPending(maxEvents);
	const RuntimeEventQueueMetrics queueMetrics = mRuntimeEventDispatcher.GetQueueMetrics();
	HealthTelemetry& telemetry = mRuntimeStore.GetHealthTelemetry();
	telemetry.TryRecordRuntimeEventDispatchStats(
		result.dispatchedEvents,
		result.handlerInvocations,
		result.handlerFailures,
		result.dispatchDurationMilliseconds);
	telemetry.TryRecordRuntimeEventQueueMetrics(
		"runtime-events",
		queueMetrics.depth,
		queueMetrics.capacity,
		static_cast<uint64_t>(queueMetrics.droppedCount),
		queueMetrics.oldestEventAgeMilliseconds);
	PublishRuntimeEventHealthObservations(result);
	return result;
}

void RuntimeUpdateController::PublishRuntimeEventHealthObservations(const RuntimeEventDispatchResult& result)
{
	const RuntimeEventQueueMetrics queueMetrics = mRuntimeEventDispatcher.GetQueueMetrics();
	if (queueMetrics.depth != mLastReportedRuntimeEventQueueDepth ||
		queueMetrics.droppedCount != mLastReportedRuntimeEventDroppedCount ||
		queueMetrics.coalescedCount != mLastReportedRuntimeEventCoalescedCount)
	{
		QueueDepthChangedEvent queueDepth;
		queueDepth.queueName = "runtime-events";
		queueDepth.depth = queueMetrics.depth;
		queueDepth.capacity = queueMetrics.capacity;
		queueDepth.droppedCount = queueMetrics.droppedCount;
		queueDepth.coalescedCount = queueMetrics.coalescedCount;
		mRuntimeEventDispatcher.PublishPayload(queueDepth, "HealthTelemetry");
		mLastReportedRuntimeEventQueueDepth = queueMetrics.depth;
		mLastReportedRuntimeEventDroppedCount = queueMetrics.droppedCount;
		mLastReportedRuntimeEventCoalescedCount = queueMetrics.coalescedCount;
	}

	if (result.handlerInvocations == 0 && result.handlerFailures == 0)
		return;

	TimingSampleRecordedEvent timing;
	timing.subsystem = "RuntimeEventDispatcher";
	timing.metric = "dispatchDuration";
	timing.value = result.dispatchDurationMilliseconds;
	timing.unit = "ms";
	mRuntimeEventDispatcher.PublishPayload(timing, "HealthTelemetry");
}
