#include "RuntimeSnapshotProvider.h"

#include "RuntimeClock.h"
#include "ShaderCompiler.h"

#include <algorithm>
#include <mutex>
#include <utility>

RuntimeSnapshotProvider::RuntimeSnapshotProvider(RuntimeHost& runtimeHost) :
	mRuntimeHost(runtimeHost)
{
}

bool RuntimeSnapshotProvider::BuildLayerPassFragmentShaderSources(const std::string& layerId, std::vector<ShaderPassBuildSource>& passSources, std::string& error) const
{
	try
	{
		ShaderPackage shaderPackage;
		{
			std::lock_guard<std::mutex> lock(mRuntimeHost.mMutex);
			const RuntimeHost::LayerPersistentState* layer = mRuntimeHost.FindLayerById(layerId);
			if (!layer)
			{
				error = "Unknown layer id: " + layerId;
				return false;
			}

			auto it = mRuntimeHost.mPackagesById.find(layer->shaderId);
			if (it == mRuntimeHost.mPackagesById.end())
			{
				error = "Unknown shader id: " + layer->shaderId;
				return false;
			}
			shaderPackage = it->second;
		}

		ShaderCompiler compiler(
			mRuntimeHost.mRepoRoot,
			mRuntimeHost.mWrapperPath,
			mRuntimeHost.mGeneratedGlslPath,
			mRuntimeHost.mPatchedGlslPath,
			mRuntimeHost.mConfig.maxTemporalHistoryFrames);
		passSources.clear();
		passSources.reserve(shaderPackage.passes.size());
		for (const ShaderPassDefinition& pass : shaderPackage.passes)
		{
			ShaderPassBuildSource passSource;
			passSource.passId = pass.id;
			passSource.inputNames = pass.inputNames;
			passSource.outputName = pass.outputName;
			if (!compiler.BuildPassFragmentShaderSource(shaderPackage, pass, passSource.fragmentShaderSource, error))
				return false;
			passSources.push_back(std::move(passSource));
		}
		return true;
	}
	catch (const std::exception& exception)
	{
		error = std::string("RuntimeSnapshotProvider::BuildLayerPassFragmentShaderSources exception: ") + exception.what();
		return false;
	}
	catch (...)
	{
		error = "RuntimeSnapshotProvider::BuildLayerPassFragmentShaderSources threw a non-standard exception.";
		return false;
	}
}

unsigned RuntimeSnapshotProvider::GetMaxTemporalHistoryFrames() const
{
	return mRuntimeHost.mConfig.maxTemporalHistoryFrames;
}

RuntimeSnapshotVersions RuntimeSnapshotProvider::GetVersions() const
{
	RuntimeSnapshotVersions versions;
	versions.renderStateVersion = mRuntimeHost.mRenderStateVersion.load(std::memory_order_relaxed);
	versions.parameterStateVersion = mRuntimeHost.mParameterStateVersion.load(std::memory_order_relaxed);
	return versions;
}

void RuntimeSnapshotProvider::AdvanceFrame()
{
	++mRuntimeHost.mFrameCounter;
}

RuntimeRenderStateSnapshot RuntimeSnapshotProvider::GetRenderStateSnapshot(unsigned outputWidth, unsigned outputHeight) const
{
	for (;;)
	{
		const RuntimeSnapshotVersions versionsBefore = GetVersions();

		RuntimeRenderStateSnapshot snapshot;
		snapshot.outputWidth = outputWidth;
		snapshot.outputHeight = outputHeight;
		{
			std::lock_guard<std::mutex> lock(mRuntimeHost.mMutex);
			BuildLayerRenderStatesLocked(outputWidth, outputHeight, snapshot.states);
		}

		const RuntimeSnapshotVersions versionsAfter = GetVersions();
		if (versionsBefore.renderStateVersion == versionsAfter.renderStateVersion &&
			versionsBefore.parameterStateVersion == versionsAfter.parameterStateVersion)
		{
			snapshot.versions = versionsAfter;
			return snapshot;
		}
	}
}

bool RuntimeSnapshotProvider::TryGetRenderStateSnapshot(unsigned outputWidth, unsigned outputHeight, RuntimeRenderStateSnapshot& snapshot) const
{
	const RuntimeSnapshotVersions versionsBefore = GetVersions();

	std::vector<RuntimeRenderState> states;
	{
		std::unique_lock<std::mutex> lock(mRuntimeHost.mMutex, std::try_to_lock);
		if (!lock.owns_lock())
			return false;

		BuildLayerRenderStatesLocked(outputWidth, outputHeight, states);
	}

	const RuntimeSnapshotVersions versionsAfter = GetVersions();
	if (versionsBefore.renderStateVersion != versionsAfter.renderStateVersion ||
		versionsBefore.parameterStateVersion != versionsAfter.parameterStateVersion)
	{
		return false;
	}

	snapshot.outputWidth = outputWidth;
	snapshot.outputHeight = outputHeight;
	snapshot.versions = versionsAfter;
	snapshot.states = std::move(states);
	return true;
}

bool RuntimeSnapshotProvider::TryRefreshSnapshotParameters(RuntimeRenderStateSnapshot& snapshot) const
{
	const uint64_t expectedRenderStateVersion = snapshot.versions.renderStateVersion;
	{
		std::unique_lock<std::mutex> lock(mRuntimeHost.mMutex, std::try_to_lock);
		if (!lock.owns_lock())
			return false;

		RefreshCachedLayerStatesLocked(snapshot.states);
	}

	const RuntimeSnapshotVersions versions = GetVersions();
	if (versions.renderStateVersion != expectedRenderStateVersion)
		return false;

	snapshot.versions = versions;
	return true;
}

void RuntimeSnapshotProvider::RefreshDynamicRenderStateFields(std::vector<RuntimeRenderState>& states) const
{
	std::lock_guard<std::mutex> lock(mRuntimeHost.mMutex);
	RefreshDynamicRenderStateFieldsLocked(states);
}

void RuntimeSnapshotProvider::BuildLayerRenderStatesLocked(unsigned outputWidth, unsigned outputHeight, std::vector<RuntimeRenderState>& states) const
{
	const HealthTelemetry::SignalStatusSnapshot signalStatus = mRuntimeHost.mHealthTelemetry.GetSignalStatusSnapshot();

	for (const RuntimeHost::LayerPersistentState& layer : mRuntimeHost.mPersistentState.layers)
	{
		auto shaderIt = mRuntimeHost.mPackagesById.find(layer.shaderId);
		if (shaderIt == mRuntimeHost.mPackagesById.end())
			continue;

		RuntimeRenderState state;
		state.layerId = layer.id;
		state.shaderId = layer.shaderId;
		state.shaderName = shaderIt->second.displayName;
		state.mixAmount = 1.0;
		state.bypass = layer.bypass ? 1.0 : 0.0;
		state.inputWidth = signalStatus.width;
		state.inputHeight = signalStatus.height;
		state.outputWidth = outputWidth;
		state.outputHeight = outputHeight;
		state.parameterDefinitions = shaderIt->second.parameters;
		state.textureAssets = shaderIt->second.textureAssets;
		state.fontAssets = shaderIt->second.fontAssets;
		state.isTemporal = shaderIt->second.temporal.enabled;
		state.temporalHistorySource = shaderIt->second.temporal.historySource;
		state.requestedTemporalHistoryLength = shaderIt->second.temporal.requestedHistoryLength;
		state.effectiveTemporalHistoryLength = shaderIt->second.temporal.effectiveHistoryLength;
		state.feedback = shaderIt->second.feedback;

		for (const ShaderParameterDefinition& definition : shaderIt->second.parameters)
		{
			ShaderParameterValue value = mRuntimeHost.DefaultValueForDefinition(definition);
			auto valueIt = layer.parameterValues.find(definition.id);
			if (valueIt != layer.parameterValues.end())
				value = valueIt->second;
			state.parameterValues[definition.id] = value;
		}

		states.push_back(state);
	}

	RefreshDynamicRenderStateFieldsLocked(states);
}

void RuntimeSnapshotProvider::RefreshCachedLayerStatesLocked(std::vector<RuntimeRenderState>& states) const
{
	for (RuntimeRenderState& state : states)
	{
		const auto layerIt = std::find_if(mRuntimeHost.mPersistentState.layers.begin(), mRuntimeHost.mPersistentState.layers.end(),
			[&state](const RuntimeHost::LayerPersistentState& layer) { return layer.id == state.layerId; });
		if (layerIt == mRuntimeHost.mPersistentState.layers.end())
			continue;

		state.bypass = layerIt->bypass ? 1.0 : 0.0;
		state.parameterValues.clear();
		for (const ShaderParameterDefinition& definition : state.parameterDefinitions)
		{
			ShaderParameterValue value = mRuntimeHost.DefaultValueForDefinition(definition);
			auto valueIt = layerIt->parameterValues.find(definition.id);
			if (valueIt != layerIt->parameterValues.end())
				value = valueIt->second;
			state.parameterValues[definition.id] = value;
		}
	}
}

void RuntimeSnapshotProvider::RefreshDynamicRenderStateFieldsLocked(std::vector<RuntimeRenderState>& states) const
{
	const RuntimeClockSnapshot clock = GetRuntimeClockSnapshot();
	const double timeSeconds = std::chrono::duration_cast<std::chrono::duration<double>>(std::chrono::steady_clock::now() - mRuntimeHost.mStartTime).count();
	const double frameCount = static_cast<double>(mRuntimeHost.mFrameCounter.load(std::memory_order_relaxed));

	for (RuntimeRenderState& state : states)
	{
		state.timeSeconds = timeSeconds;
		state.utcTimeSeconds = clock.utcTimeSeconds;
		state.utcOffsetSeconds = clock.utcOffsetSeconds;
		state.startupRandom = mRuntimeHost.mStartupRandom;
		state.frameCount = frameCount;
	}
}