#include "RuntimeLayerModel.h"

#include "RuntimeParameterUtils.h"
#include "RuntimeTextTextureComposer.h"

#include <algorithm>
#include <chrono>
#include <cctype>
#include <set>
#include <utility>

namespace RenderCadenceCompositor
{
namespace
{
JsonValue ParameterValueToJson(const ShaderParameterDefinition& definition, const ShaderParameterValue& value)
{
	switch (definition.type)
	{
	case ShaderParameterType::Float:
		return JsonValue(value.numberValues.empty() ? 0.0 : value.numberValues.front());
	case ShaderParameterType::Vec2:
	case ShaderParameterType::Color:
	{
		JsonValue array = JsonValue::MakeArray();
		for (double number : value.numberValues)
			array.pushBack(JsonValue(number));
		return array;
	}
	case ShaderParameterType::Boolean:
		return JsonValue(value.booleanValue);
	case ShaderParameterType::Enum:
		return JsonValue(value.enumValue);
	case ShaderParameterType::Text:
		return JsonValue(value.textValue);
	case ShaderParameterType::Trigger:
		return JsonValue(value.numberValues.empty() ? 0.0 : value.numberValues.front());
	}
	return JsonValue();
}

bool ParseRuntimeLayerNumber(const std::string& layerId, uint64_t& number)
{
	const std::string prefix = "runtime-layer-";
	if (layerId.compare(0, prefix.size(), prefix) != 0)
		return false;

	const std::string suffix = layerId.substr(prefix.size());
	if (suffix.empty())
		return false;

	uint64_t parsed = 0;
	for (char character : suffix)
	{
		if (!std::isdigit(static_cast<unsigned char>(character)))
			return false;
		parsed = parsed * 10 + static_cast<uint64_t>(character - '0');
	}

	number = parsed;
	return true;
}

std::string AllocateRestoredLayerId(std::set<std::string>& usedLayerIds, uint64_t& nextLayerNumber)
{
	for (;;)
	{
		std::string candidate = "runtime-layer-" + std::to_string(nextLayerNumber++);
		if (usedLayerIds.insert(candidate).second)
			return candidate;
	}
}
}

bool RuntimeLayerModel::InitializeSingleLayer(const SupportedShaderCatalog& shaderCatalog, const std::string& shaderId, std::string& error)
{
	Clear();
	if (shaderId.empty())
	{
		error.clear();
		return true;
	}

	const ShaderPackage* shaderPackage = shaderCatalog.FindPackage(shaderId);
	if (!shaderPackage)
	{
		error = "Shader '" + shaderId + "' is not in the supported shader catalog.";
		return false;
	}

	Layer layer;
	layer.id = AllocateLayerId();
	layer.shaderId = shaderPackage->id;
	layer.packageFingerprint = ShaderPackageFingerprint(*shaderPackage);
	layer.shaderName = shaderPackage->displayName.empty() ? shaderPackage->id : shaderPackage->displayName;
	layer.buildState = RuntimeLayerBuildState::Pending;
	layer.message = "Runtime Slang build is waiting to start.";
	InitializeDefaultParameterValues(layer, *shaderPackage);
	mLayers.push_back(std::move(layer));
	error.clear();
	return true;
}

bool RuntimeLayerModel::AddLayer(const SupportedShaderCatalog& shaderCatalog, const std::string& shaderId, std::string& layerId, std::string& error)
{
	const ShaderPackage* shaderPackage = shaderCatalog.FindPackage(shaderId);
	if (!shaderPackage)
	{
		error = "Shader '" + shaderId + "' is not in the supported shader catalog.";
		return false;
	}

	Layer layer;
	layer.id = AllocateLayerId();
	layer.shaderId = shaderPackage->id;
	layer.packageFingerprint = ShaderPackageFingerprint(*shaderPackage);
	layer.shaderName = shaderPackage->displayName.empty() ? shaderPackage->id : shaderPackage->displayName;
	layer.buildState = RuntimeLayerBuildState::Pending;
	layer.message = "Runtime Slang build is waiting to start.";
	InitializeDefaultParameterValues(layer, *shaderPackage);
	layerId = layer.id;
	mLayers.push_back(std::move(layer));
	error.clear();
	return true;
}

bool RuntimeLayerModel::RemoveLayer(const std::string& layerId, std::string& error)
{
	for (auto layerIt = mLayers.begin(); layerIt != mLayers.end(); ++layerIt)
	{
		if (layerIt->id != layerId)
			continue;

		mLayers.erase(layerIt);
		error.clear();
		return true;
	}

	error = "Unknown runtime layer id: " + layerId;
	return false;
}

bool RuntimeLayerModel::ReorderLayer(const std::string& layerId, int targetIndex, std::string& error)
{
	auto layerIt = std::find_if(mLayers.begin(), mLayers.end(), [&layerId](const Layer& layer) {
		return layer.id == layerId;
	});
	if (layerIt == mLayers.end())
	{
		error = "Unknown runtime layer id: " + layerId;
		return false;
	}

	if (targetIndex < 0)
		targetIndex = 0;
	if (targetIndex >= static_cast<int>(mLayers.size()))
		targetIndex = static_cast<int>(mLayers.size()) - 1;

	Layer layer = std::move(*layerIt);
	mLayers.erase(layerIt);
	std::size_t destinationIndex = static_cast<std::size_t>(targetIndex);
	if (destinationIndex > mLayers.size())
		destinationIndex = mLayers.size();
	mLayers.insert(mLayers.begin() + static_cast<std::ptrdiff_t>(destinationIndex), std::move(layer));
	error.clear();
	return true;
}

bool RuntimeLayerModel::SetLayerBypass(const std::string& layerId, bool bypass, std::string& error)
{
	Layer* layer = FindLayer(layerId);
	if (!layer)
	{
		error = "Unknown runtime layer id: " + layerId;
		return false;
	}
	layer->bypass = bypass;
	error.clear();
	return true;
}

bool RuntimeLayerModel::SetLayerShader(const SupportedShaderCatalog& shaderCatalog, const std::string& layerId, const std::string& shaderId, std::string& error)
{
	Layer* layer = FindLayer(layerId);
	if (!layer)
	{
		error = "Unknown runtime layer id: " + layerId;
		return false;
	}

	const ShaderPackage* shaderPackage = shaderCatalog.FindPackage(shaderId);
	if (!shaderPackage)
	{
		error = "Shader '" + shaderId + "' is not in the supported shader catalog.";
		return false;
	}

	layer->shaderId = shaderPackage->id;
	layer->packageFingerprint = ShaderPackageFingerprint(*shaderPackage);
	layer->shaderName = shaderPackage->displayName.empty() ? shaderPackage->id : shaderPackage->displayName;
	layer->buildState = RuntimeLayerBuildState::Pending;
	layer->message = "Runtime Slang build is waiting to start.";
	layer->renderReady = false;
	layer->artifact = RuntimeShaderArtifact();
	InitializeDefaultParameterValues(*layer, *shaderPackage);
	error.clear();
	return true;
}

bool RuntimeLayerModel::UpdateParameter(const std::string& layerId, const std::string& parameterId, const JsonValue& value, std::string& error)
{
	Layer* layer = FindLayer(layerId);
	if (!layer)
	{
		error = "Unknown runtime layer id: " + layerId;
		return false;
	}

	const ShaderParameterDefinition* definition = FindParameterDefinition(*layer, parameterId);
	if (!definition)
	{
		error = "Unknown parameter id '" + parameterId + "' for layer " + layerId + ".";
		return false;
	}

	ShaderParameterValue normalizedValue;
	if (definition->type == ShaderParameterType::Trigger)
	{
		const auto currentIt = layer->parameterValues.find(parameterId);
		const double previousCount = currentIt == layer->parameterValues.end() || currentIt->second.numberValues.empty()
			? 0.0
			: currentIt->second.numberValues.front();
		normalizedValue.numberValues = { previousCount + 1.0, RuntimeElapsedSeconds() };
	}
	else if (!NormalizeAndValidateParameterValue(*definition, value, normalizedValue, error))
	{
		return false;
	}

	layer->parameterValues[parameterId] = normalizedValue;
	if (layer->renderReady)
	{
		layer->artifact.parameterValues = layer->parameterValues;
		if (definition->type == ShaderParameterType::Text && !PrepareRuntimeTextTextures(layer->artifact, error))
			return false;
	}
	error.clear();
	return true;
}

bool RuntimeLayerModel::ResetParameters(const std::string& layerId, std::string& error)
{
	Layer* layer = FindLayer(layerId);
	if (!layer)
	{
		error = "Unknown runtime layer id: " + layerId;
		return false;
	}

	layer->parameterValues.clear();
	for (const ShaderParameterDefinition& definition : layer->parameterDefinitions)
		layer->parameterValues[definition.id] = DefaultValueForDefinition(definition);
	if (layer->renderReady)
	{
		layer->artifact.parameterValues = layer->parameterValues;
		if (!PrepareRuntimeTextTextures(layer->artifact, error))
			return false;
	}
	error.clear();
	return true;
}

bool RuntimeLayerModel::InitializeFromRuntimeState(const SupportedShaderCatalog& shaderCatalog, const JsonValue& runtimeState, std::string& error)
{
	if (!runtimeState.isObject())
	{
		error = "Runtime state root must be a JSON object.";
		return false;
	}

	const JsonValue* layersValue = runtimeState.find("layers");
	if (!layersValue || !layersValue->isArray())
	{
		error = "Runtime state must contain a layers array.";
		return false;
	}

	std::vector<Layer> restoredLayers;
	std::set<std::string> usedLayerIds;
	uint64_t nextLayerNumber = 1;

	for (const JsonValue& layerValue : layersValue->asArray())
	{
		if (!layerValue.isObject())
			continue;

		const JsonValue* shaderIdValue = layerValue.find("shaderId");
		if (!shaderIdValue || !shaderIdValue->isString() || shaderIdValue->asString().empty())
			continue;

		const ShaderPackage* shaderPackage = shaderCatalog.FindPackage(shaderIdValue->asString());
		if (!shaderPackage)
			continue;

		Layer layer;
		const JsonValue* layerIdValue = layerValue.find("id");
		if (layerIdValue && layerIdValue->isString() && !layerIdValue->asString().empty() && usedLayerIds.insert(layerIdValue->asString()).second)
			layer.id = layerIdValue->asString();
		else
			layer.id = AllocateRestoredLayerId(usedLayerIds, nextLayerNumber);

		uint64_t restoredLayerNumber = 0;
		if (ParseRuntimeLayerNumber(layer.id, restoredLayerNumber) && restoredLayerNumber >= nextLayerNumber)
			nextLayerNumber = restoredLayerNumber + 1;

		layer.shaderId = shaderPackage->id;
		layer.packageFingerprint = ShaderPackageFingerprint(*shaderPackage);
		layer.shaderName = shaderPackage->displayName.empty() ? shaderPackage->id : shaderPackage->displayName;
		const JsonValue* bypassValue = layerValue.find("bypass");
		layer.bypass = bypassValue && bypassValue->isBoolean() ? bypassValue->asBoolean() : false;
		layer.buildState = RuntimeLayerBuildState::Pending;
		layer.message = "Runtime Slang build is waiting to start.";
		InitializeDefaultParameterValues(layer, *shaderPackage);

		const JsonValue* parameterValues = layerValue.find("parameterValues");
		if (parameterValues && parameterValues->isObject())
		{
			for (const ShaderParameterDefinition& definition : layer.parameterDefinitions)
			{
				const JsonValue* value = parameterValues->find(definition.id);
				if (!value)
					continue;

				ShaderParameterValue normalizedValue;
				std::string normalizeError;
				if (NormalizeAndValidateParameterValue(definition, *value, normalizedValue, normalizeError))
					layer.parameterValues[definition.id] = normalizedValue;
			}
		}

		restoredLayers.push_back(std::move(layer));
	}

	if (restoredLayers.empty())
	{
		error = "Runtime state did not contain any supported layers.";
		return false;
	}

	mLayers = std::move(restoredLayers);
	mNextLayerNumber = nextLayerNumber;
	error.clear();
	return true;
}

bool RuntimeLayerModel::ReloadFromCatalog(const SupportedShaderCatalog& shaderCatalog, std::vector<std::pair<std::string, std::string>>& buildsToStart, std::string& error)
{
	buildsToStart.clear();
	for (Layer& layer : mLayers)
	{
		const ShaderPackage* shaderPackage = shaderCatalog.FindPackage(layer.shaderId);
		if (!shaderPackage)
		{
			layer.buildState = RuntimeLayerBuildState::Failed;
			layer.message = "Shader '" + layer.shaderId + "' is no longer available after reload.";
			continue;
		}

		const std::string nextFingerprint = ShaderPackageFingerprint(*shaderPackage);
		if (layer.packageFingerprint == nextFingerprint)
		{
			buildsToStart.push_back({ layer.id, layer.shaderId });
			continue;
		}

		std::map<std::string, ShaderParameterDefinition> previousDefinitions;
		for (const ShaderParameterDefinition& definition : layer.parameterDefinitions)
			previousDefinitions[definition.id] = definition;

		std::map<std::string, ShaderParameterValue> nextValues;
		for (const ShaderParameterDefinition& nextDefinition : shaderPackage->parameters)
		{
			const auto previousDefinitionIt = previousDefinitions.find(nextDefinition.id);
			const auto previousValueIt = layer.parameterValues.find(nextDefinition.id);
			if (previousDefinitionIt != previousDefinitions.end()
				&& previousValueIt != layer.parameterValues.end()
				&& previousDefinitionIt->second.type == nextDefinition.type)
			{
				ShaderParameterValue preservedValue;
				JsonValue valueJson = ParameterValueToJson(previousDefinitionIt->second, previousValueIt->second);
				std::string normalizeError;
				if (NormalizeAndValidateParameterValue(nextDefinition, valueJson, preservedValue, normalizeError))
				{
					nextValues[nextDefinition.id] = preservedValue;
					continue;
				}
			}
			nextValues[nextDefinition.id] = DefaultValueForDefinition(nextDefinition);
		}

		layer.shaderName = shaderPackage->displayName.empty() ? shaderPackage->id : shaderPackage->displayName;
		layer.packageFingerprint = nextFingerprint;
		layer.parameterDefinitions = shaderPackage->parameters;
		layer.parameterValues = std::move(nextValues);
		buildsToStart.push_back({ layer.id, layer.shaderId });
	}

	error.clear();
	return true;
}

void RuntimeLayerModel::Clear()
{
	mLayers.clear();
}

bool RuntimeLayerModel::MarkBuildStarted(const std::string& layerId, const std::string& message, std::string& error, bool preserveExistingRenderArtifact)
{
	Layer* layer = FindLayer(layerId);
	if (!layer)
	{
		error = "Unknown runtime layer id: " + layerId;
		return false;
	}

	layer->buildState = RuntimeLayerBuildState::Pending;
	layer->message = message;
	layer->preserveRenderDuringBuild = preserveExistingRenderArtifact && layer->renderReady;
	if (!layer->preserveRenderDuringBuild)
	{
		layer->renderReady = false;
		layer->artifact = RuntimeShaderArtifact();
	}
	error.clear();
	return true;
}

bool RuntimeLayerModel::MarkBuildReady(const RuntimeShaderArtifact& artifact, std::string& error)
{
	Layer* layer = artifact.layerId.empty() ? FindFirstLayerForShader(artifact.shaderId) : FindLayer(artifact.layerId);
	if (!layer)
	{
		error = artifact.layerId.empty()
			? "No runtime layer is waiting for shader artifact: " + artifact.shaderId
			: "No runtime layer is waiting for shader artifact on layer: " + artifact.layerId;
		return false;
	}

	RuntimeShaderArtifact nextArtifact = artifact;
	nextArtifact.parameterValues = layer->parameterValues;
	if (!PrepareRuntimeTextTextures(nextArtifact, error))
	{
		layer->buildState = RuntimeLayerBuildState::Failed;
		layer->message = error;
		if (!layer->preserveRenderDuringBuild)
		{
			layer->renderReady = false;
			layer->artifact = RuntimeShaderArtifact();
		}
		layer->preserveRenderDuringBuild = false;
		return false;
	}

	layer->shaderName = artifact.displayName.empty() ? artifact.shaderId : artifact.displayName;
	layer->packageFingerprint = artifact.packageFingerprint;
	layer->buildState = RuntimeLayerBuildState::Ready;
	layer->message = artifact.message;
	layer->renderReady = true;
	layer->preserveRenderDuringBuild = false;
	layer->artifact = std::move(nextArtifact);
	error.clear();
	return true;
}

bool RuntimeLayerModel::MarkBuildFailedForShader(const std::string& shaderId, const std::string& message)
{
	Layer* layer = FindFirstLayerForShader(shaderId);
	if (!layer)
		return false;

	std::string error;
	return MarkBuildFailed(layer->id, message, error);
}

bool RuntimeLayerModel::MarkBuildFailed(const std::string& layerId, const std::string& message, std::string& error)
{
	Layer* layer = FindLayer(layerId);
	if (!layer)
	{
		error = "Unknown runtime layer id: " + layerId;
		return false;
	}

	layer->buildState = RuntimeLayerBuildState::Failed;
	layer->message = message;
	if (!layer->preserveRenderDuringBuild)
	{
		layer->renderReady = false;
		layer->artifact = RuntimeShaderArtifact();
	}
	layer->preserveRenderDuringBuild = false;
	error.clear();
	return true;
}

bool RuntimeLayerModel::MarkRenderCommitFailed(const std::string& layerId, const std::string& message, std::string& error)
{
	return MarkBuildFailed(layerId, message, error);
}

RuntimeLayerModelSnapshot RuntimeLayerModel::Snapshot() const
{
	RuntimeLayerModelSnapshot snapshot;
	snapshot.compileSucceeded = true;

	for (const Layer& layer : mLayers)
	{
		snapshot.displayLayers.push_back(ToReadModel(layer));
		if (!layer.message.empty() && snapshot.compileMessage.empty())
			snapshot.compileMessage = layer.message;
		if (layer.buildState == RuntimeLayerBuildState::Failed)
			snapshot.compileSucceeded = false;
		if (layer.renderReady)
		{
			RuntimeRenderLayerModel renderLayer;
			renderLayer.id = layer.id;
			renderLayer.bypass = layer.bypass;
			renderLayer.artifact = layer.artifact;
			renderLayer.shaderId = renderLayer.artifact.shaderId.empty() ? layer.shaderId : renderLayer.artifact.shaderId;
			if (layer.buildState == RuntimeLayerBuildState::Ready)
				renderLayer.artifact.parameterValues = layer.parameterValues;
			renderLayer.artifact.fontAtlases.clear();
			snapshot.renderLayers.push_back(std::move(renderLayer));
		}
	}

	if (snapshot.compileMessage.empty())
		snapshot.compileMessage = mLayers.empty() ? "Runtime shader build disabled." : "Runtime shader build has not completed yet.";
	return snapshot;
}

std::string RuntimeLayerModel::FirstLayerId() const
{
	return mLayers.empty() ? std::string() : mLayers.front().id;
}

std::vector<std::pair<std::string, std::string>> RuntimeLayerModel::PendingLayerBuilds() const
{
	std::vector<std::pair<std::string, std::string>> builds;
	for (const Layer& layer : mLayers)
	{
		if (layer.buildState == RuntimeLayerBuildState::Pending)
			builds.push_back({ layer.id, layer.shaderId });
	}
	return builds;
}

RuntimeLayerModel::Layer* RuntimeLayerModel::FindLayer(const std::string& layerId)
{
	for (Layer& layer : mLayers)
	{
		if (layer.id == layerId)
			return &layer;
	}
	return nullptr;
}

const RuntimeLayerModel::Layer* RuntimeLayerModel::FindLayer(const std::string& layerId) const
{
	for (const Layer& layer : mLayers)
	{
		if (layer.id == layerId)
			return &layer;
	}
	return nullptr;
}

RuntimeLayerModel::Layer* RuntimeLayerModel::FindFirstLayerForShader(const std::string& shaderId)
{
	for (Layer& layer : mLayers)
	{
		if (layer.shaderId == shaderId)
			return &layer;
	}
	return nullptr;
}

void RuntimeLayerModel::InitializeDefaultParameterValues(Layer& layer, const ShaderPackage& shaderPackage)
{
	layer.parameterDefinitions = shaderPackage.parameters;
	layer.parameterValues.clear();
	for (const ShaderParameterDefinition& definition : layer.parameterDefinitions)
		layer.parameterValues[definition.id] = DefaultValueForDefinition(definition);
}

const ShaderParameterDefinition* RuntimeLayerModel::FindParameterDefinition(const Layer& layer, const std::string& parameterId)
{
	for (const ShaderParameterDefinition& definition : layer.parameterDefinitions)
	{
		if (definition.id == parameterId)
			return &definition;
	}
	return nullptr;
}

std::string RuntimeLayerModel::AllocateLayerId()
{
	return "runtime-layer-" + std::to_string(mNextLayerNumber++);
}

RuntimeLayerReadModel RuntimeLayerModel::ToReadModel(const Layer& layer)
{
	RuntimeLayerReadModel readModel;
	readModel.id = layer.id;
	readModel.shaderId = layer.shaderId;
	readModel.shaderName = layer.shaderName;
	readModel.bypass = layer.bypass;
	readModel.buildState = layer.buildState;
	readModel.message = layer.message;
	readModel.renderReady = layer.renderReady;
	readModel.parameterDefinitions = layer.parameterDefinitions;
	readModel.parameterValues = layer.parameterValues;
	return readModel;
}

double RuntimeLayerModel::RuntimeElapsedSeconds() const
{
	return std::chrono::duration_cast<std::chrono::duration<double>>(std::chrono::steady_clock::now() - mStartTime).count();
}
}