#include "stdafx.h"
#include "ShaderPackageRegistry.h"

#include "RuntimeJson.h"

#include <algorithm>
#include <cmath>
#include <cctype>
#include <fstream>
#include <sstream>

namespace
{
std::string Trim(const std::string& text)
{
	std::size_t start = 0;
	while (start < text.size() && std::isspace(static_cast<unsigned char>(text[start])))
		++start;

	std::size_t end = text.size();
	while (end > start && std::isspace(static_cast<unsigned char>(text[end - 1])))
		--end;

	return text.substr(start, end - start);
}

bool IsFiniteNumber(double value)
{
	return std::isfinite(value) != 0;
}

std::vector<double> JsonArrayToNumbers(const JsonValue& value)
{
	std::vector<double> numbers;
	for (const JsonValue& item : value.asArray())
	{
		if (item.isNumber())
			numbers.push_back(item.asNumber());
	}
	return numbers;
}

bool ParseShaderParameterType(const std::string& typeName, ShaderParameterType& type)
{
	if (typeName == "float")
	{
		type = ShaderParameterType::Float;
		return true;
	}
	if (typeName == "vec2")
	{
		type = ShaderParameterType::Vec2;
		return true;
	}
	if (typeName == "color")
	{
		type = ShaderParameterType::Color;
		return true;
	}
	if (typeName == "bool")
	{
		type = ShaderParameterType::Boolean;
		return true;
	}
	if (typeName == "enum")
	{
		type = ShaderParameterType::Enum;
		return true;
	}
	return false;
}

bool ParseTemporalHistorySource(const std::string& sourceName, TemporalHistorySource& source)
{
	if (sourceName == "source")
	{
		source = TemporalHistorySource::Source;
		return true;
	}
	if (sourceName == "preLayerInput")
	{
		source = TemporalHistorySource::PreLayerInput;
		return true;
	}
	if (sourceName == "none")
	{
		source = TemporalHistorySource::None;
		return true;
	}
	return false;
}

std::string ReadTextFile(const std::filesystem::path& path, std::string& error)
{
	std::ifstream input(path, std::ios::binary);
	if (!input)
	{
		error = "Could not open file: " + path.string();
		return std::string();
	}

	std::ostringstream buffer;
	buffer << input.rdbuf();
	return buffer.str();
}

std::string ManifestPathMessage(const std::filesystem::path& manifestPath)
{
	return manifestPath.string();
}

bool RequireStringField(const JsonValue& object, const char* fieldName, std::string& value, const std::filesystem::path& manifestPath, std::string& error)
{
	const JsonValue* fieldValue = object.find(fieldName);
	if (!fieldValue || !fieldValue->isString())
	{
		error = "Shader manifest is missing required string '" + std::string(fieldName) + "' field: " + ManifestPathMessage(manifestPath);
		return false;
	}

	value = fieldValue->asString();
	return true;
}

bool RequireNonEmptyStringField(const JsonValue& object, const char* fieldName, std::string& value, const std::filesystem::path& manifestPath, std::string& error)
{
	if (!RequireStringField(object, fieldName, value, manifestPath, error))
		return false;
	if (Trim(value).empty())
	{
		error = "Shader manifest string '" + std::string(fieldName) + "' must not be empty: " + ManifestPathMessage(manifestPath);
		return false;
	}
	return true;
}

bool OptionalStringField(const JsonValue& object, const char* fieldName, std::string& value, const std::string& fallback, const std::filesystem::path& manifestPath, std::string& error)
{
	const JsonValue* fieldValue = object.find(fieldName);
	if (!fieldValue)
	{
		value = fallback;
		return true;
	}
	if (!fieldValue->isString())
	{
		error = "Shader manifest field '" + std::string(fieldName) + "' must be a string in: " + ManifestPathMessage(manifestPath);
		return false;
	}
	value = fieldValue->asString();
	return true;
}

bool OptionalArrayField(const JsonValue& object, const char* fieldName, const JsonValue*& value, const std::filesystem::path& manifestPath, std::string& error)
{
	value = object.find(fieldName);
	if (!value)
		return true;
	if (!value->isArray())
	{
		error = "Shader manifest '" + std::string(fieldName) + "' field must be an array in: " + ManifestPathMessage(manifestPath);
		return false;
	}
	return true;
}

bool OptionalObjectField(const JsonValue& object, const char* fieldName, const JsonValue*& value, const std::filesystem::path& manifestPath, std::string& error)
{
	value = object.find(fieldName);
	if (!value)
		return true;
	if (!value->isObject())
	{
		error = "Shader manifest '" + std::string(fieldName) + "' field must be an object in: " + ManifestPathMessage(manifestPath);
		return false;
	}
	return true;
}

bool NumberListFromJsonValue(const JsonValue& value, std::vector<double>& numbers, const std::string& fieldName, const std::filesystem::path& manifestPath, std::string& error)
{
	if (value.isNumber())
	{
		numbers.push_back(value.asNumber());
		return true;
	}
	if (value.isArray())
	{
		numbers = JsonArrayToNumbers(value);
		if (numbers.size() != value.asArray().size())
		{
			error = "Shader parameter field '" + fieldName + "' must contain only numbers in: " + ManifestPathMessage(manifestPath);
			return false;
		}
		return true;
	}

	error = "Shader parameter field '" + fieldName + "' must be a number or array of numbers in: " + ManifestPathMessage(manifestPath);
	return false;
}

bool ValidateShaderIdentifier(const std::string& identifier, const std::string& fieldName, const std::filesystem::path& manifestPath, std::string& error)
{
	if (identifier.empty() || !(std::isalpha(static_cast<unsigned char>(identifier.front())) || identifier.front() == '_'))
	{
		error = "Shader manifest field '" + fieldName + "' must be a valid shader identifier in: " + ManifestPathMessage(manifestPath);
		return false;
	}

	for (char ch : identifier)
	{
		const unsigned char unsignedCh = static_cast<unsigned char>(ch);
		if (!(std::isalnum(unsignedCh) || ch == '_'))
		{
			error = "Shader manifest field '" + fieldName + "' must be a valid shader identifier in: " + ManifestPathMessage(manifestPath);
			return false;
		}
	}

	return true;
}

bool ParseShaderMetadata(const JsonValue& manifestJson, ShaderPackage& shaderPackage, const std::filesystem::path& manifestPath, std::string& error)
{
	if (!RequireStringField(manifestJson, "id", shaderPackage.id, manifestPath, error) ||
		!RequireStringField(manifestJson, "name", shaderPackage.displayName, manifestPath, error) ||
		!OptionalStringField(manifestJson, "description", shaderPackage.description, "", manifestPath, error) ||
		!OptionalStringField(manifestJson, "category", shaderPackage.category, "", manifestPath, error) ||
		!OptionalStringField(manifestJson, "entryPoint", shaderPackage.entryPoint, "shadeVideo", manifestPath, error))
	{
		return false;
	}

	if (!ValidateShaderIdentifier(shaderPackage.entryPoint, "entryPoint", manifestPath, error))
		return false;

	shaderPackage.directoryPath = manifestPath.parent_path();
	shaderPackage.shaderPath = shaderPackage.directoryPath / "shader.slang";
	shaderPackage.manifestPath = manifestPath;
	return true;
}

bool ParseTextureAssets(const JsonValue& manifestJson, ShaderPackage& shaderPackage, const std::filesystem::path& manifestPath, std::string& error)
{
	const JsonValue* texturesValue = nullptr;
	if (!OptionalArrayField(manifestJson, "textures", texturesValue, manifestPath, error))
		return false;
	if (!texturesValue)
		return true;

	for (const JsonValue& textureJson : texturesValue->asArray())
	{
		if (!textureJson.isObject())
		{
			error = "Shader texture entry must be an object in: " + ManifestPathMessage(manifestPath);
			return false;
		}

		std::string textureId;
		std::string texturePath;
		if (!RequireNonEmptyStringField(textureJson, "id", textureId, manifestPath, error) ||
			!RequireNonEmptyStringField(textureJson, "path", texturePath, manifestPath, error))
		{
			error = "Shader texture is missing required 'id' or 'path' in: " + ManifestPathMessage(manifestPath);
			return false;
		}
		if (!ValidateShaderIdentifier(textureId, "textures[].id", manifestPath, error))
			return false;

		ShaderTextureAsset textureAsset;
		textureAsset.id = textureId;
		textureAsset.path = shaderPackage.directoryPath / texturePath;
		if (!std::filesystem::exists(textureAsset.path))
		{
			error = "Shader texture asset not found for package " + shaderPackage.id + ": " + textureAsset.path.string();
			return false;
		}

		textureAsset.writeTime = std::filesystem::last_write_time(textureAsset.path);
		shaderPackage.textureAssets.push_back(textureAsset);
	}

	return true;
}

bool ParseTemporalSettings(const JsonValue& manifestJson, ShaderPackage& shaderPackage, unsigned maxTemporalHistoryFrames, const std::filesystem::path& manifestPath, std::string& error)
{
	const JsonValue* temporalValue = nullptr;
	if (!OptionalObjectField(manifestJson, "temporal", temporalValue, manifestPath, error))
		return false;
	if (!temporalValue)
		return true;

	const JsonValue* enabledValue = temporalValue->find("enabled");
	if (!enabledValue || !enabledValue->asBoolean(false))
		return true;

	std::string historySourceName;
	if (!RequireNonEmptyStringField(*temporalValue, "historySource", historySourceName, manifestPath, error))
	{
		error = "Temporal shader is missing required 'historySource' in: " + ManifestPathMessage(manifestPath);
		return false;
	}

	const JsonValue* historyLengthValue = temporalValue->find("historyLength");
	if (!historyLengthValue || !historyLengthValue->isNumber())
	{
		error = "Temporal shader is missing required numeric 'historyLength' in: " + ManifestPathMessage(manifestPath);
		return false;
	}

	TemporalHistorySource historySource = TemporalHistorySource::None;
	if (!ParseTemporalHistorySource(historySourceName, historySource))
	{
		error = "Unsupported temporal historySource '" + historySourceName + "' in: " + ManifestPathMessage(manifestPath);
		return false;
	}

	const double requestedHistoryLength = historyLengthValue->asNumber();
	if (!IsFiniteNumber(requestedHistoryLength) || requestedHistoryLength <= 0.0 || std::floor(requestedHistoryLength) != requestedHistoryLength)
	{
		error = "Temporal shader 'historyLength' must be a positive integer in: " + ManifestPathMessage(manifestPath);
		return false;
	}

	shaderPackage.temporal.enabled = true;
	shaderPackage.temporal.historySource = historySource;
	shaderPackage.temporal.requestedHistoryLength = static_cast<unsigned>(requestedHistoryLength);
	shaderPackage.temporal.effectiveHistoryLength = std::min(shaderPackage.temporal.requestedHistoryLength, maxTemporalHistoryFrames);
	return true;
}

bool ParseParameterNumberField(const JsonValue& parameterJson, const char* fieldName, std::vector<double>& values, const std::filesystem::path& manifestPath, std::string& error)
{
	if (const JsonValue* fieldValue = parameterJson.find(fieldName))
		return NumberListFromJsonValue(*fieldValue, values, fieldName, manifestPath, error);
	return true;
}

bool ParseParameterDefault(const JsonValue& parameterJson, ShaderParameterDefinition& definition, const std::filesystem::path& manifestPath, std::string& error)
{
	const JsonValue* defaultValue = parameterJson.find("default");
	if (!defaultValue)
		return true;

	if (definition.type == ShaderParameterType::Boolean)
	{
		if (!defaultValue->isBoolean())
		{
			error = "Boolean parameter default must be a boolean for: " + definition.id;
			return false;
		}
		definition.defaultBoolean = defaultValue->asBoolean(false);
		return true;
	}

	if (definition.type == ShaderParameterType::Enum)
	{
		if (!defaultValue->isString())
		{
			error = "Enum parameter default must be a string for: " + definition.id;
			return false;
		}
		definition.defaultEnumValue = defaultValue->asString();
		return true;
	}

	return NumberListFromJsonValue(*defaultValue, definition.defaultNumbers, "default", manifestPath, error);
}

bool ParseParameterOptions(const JsonValue& parameterJson, ShaderParameterDefinition& definition, const std::filesystem::path& manifestPath, std::string& error)
{
	const JsonValue* optionsValue = nullptr;
	if (!OptionalArrayField(parameterJson, "options", optionsValue, manifestPath, error) || !optionsValue)
	{
		error = "Enum parameter is missing 'options' in: " + ManifestPathMessage(manifestPath);
		return false;
	}

	for (const JsonValue& optionJson : optionsValue->asArray())
	{
		if (!optionJson.isObject())
		{
			error = "Enum parameter option must be an object in: " + ManifestPathMessage(manifestPath);
			return false;
		}

		ShaderParameterOption option;
		if (!RequireStringField(optionJson, "value", option.value, manifestPath, error) ||
			!RequireStringField(optionJson, "label", option.label, manifestPath, error))
		{
			error = "Enum parameter option is missing 'value' or 'label' in: " + ManifestPathMessage(manifestPath);
			return false;
		}
		definition.enumOptions.push_back(option);
	}

	bool defaultFound = definition.defaultEnumValue.empty();
	for (const ShaderParameterOption& option : definition.enumOptions)
	{
		if (option.value == definition.defaultEnumValue)
		{
			defaultFound = true;
			break;
		}
	}

	if (!defaultFound)
	{
		error = "Enum parameter default is not present in its option list for: " + definition.id;
		return false;
	}

	return true;
}

bool ParseParameterDefinition(const JsonValue& parameterJson, ShaderParameterDefinition& definition, const std::filesystem::path& manifestPath, std::string& error)
{
	if (!parameterJson.isObject())
	{
		error = "Shader parameter entry must be an object in: " + ManifestPathMessage(manifestPath);
		return false;
	}

	std::string typeName;
	if (!RequireStringField(parameterJson, "id", definition.id, manifestPath, error) ||
		!RequireStringField(parameterJson, "label", definition.label, manifestPath, error) ||
		!RequireStringField(parameterJson, "type", typeName, manifestPath, error))
	{
		error = "Shader parameter is missing required fields in: " + ManifestPathMessage(manifestPath);
		return false;
	}

	if (!ParseShaderParameterType(typeName, definition.type))
	{
		error = "Unsupported parameter type '" + typeName + "' in: " + ManifestPathMessage(manifestPath);
		return false;
	}
	if (!ValidateShaderIdentifier(definition.id, "parameters[].id", manifestPath, error))
		return false;

	if (!ParseParameterDefault(parameterJson, definition, manifestPath, error) ||
		!ParseParameterNumberField(parameterJson, "min", definition.minNumbers, manifestPath, error) ||
		!ParseParameterNumberField(parameterJson, "max", definition.maxNumbers, manifestPath, error) ||
		!ParseParameterNumberField(parameterJson, "step", definition.stepNumbers, manifestPath, error))
	{
		return false;
	}

	if (definition.type == ShaderParameterType::Enum)
		return ParseParameterOptions(parameterJson, definition, manifestPath, error);

	return true;
}

bool ParseParameterDefinitions(const JsonValue& manifestJson, ShaderPackage& shaderPackage, const std::filesystem::path& manifestPath, std::string& error)
{
	const JsonValue* parametersValue = nullptr;
	if (!OptionalArrayField(manifestJson, "parameters", parametersValue, manifestPath, error))
		return false;
	if (!parametersValue)
		return true;

	for (const JsonValue& parameterJson : parametersValue->asArray())
	{
		ShaderParameterDefinition definition;
		if (!ParseParameterDefinition(parameterJson, definition, manifestPath, error))
			return false;
		shaderPackage.parameters.push_back(definition);
	}

	return true;
}
}

ShaderPackageRegistry::ShaderPackageRegistry(unsigned maxTemporalHistoryFrames)
	: mMaxTemporalHistoryFrames(maxTemporalHistoryFrames)
{
}

bool ShaderPackageRegistry::Scan(const std::filesystem::path& shaderRoot, std::map<std::string, ShaderPackage>& packagesById, std::vector<std::string>& packageOrder, std::string& error) const
{
	packagesById.clear();
	packageOrder.clear();

	if (!std::filesystem::exists(shaderRoot))
	{
		error = "Shader library directory does not exist: " + shaderRoot.string();
		return false;
	}

	for (const auto& entry : std::filesystem::directory_iterator(shaderRoot))
	{
		if (!entry.is_directory())
			continue;

		std::filesystem::path manifestPath = entry.path() / "shader.json";
		if (!std::filesystem::exists(manifestPath))
			continue;

		ShaderPackage shaderPackage;
		if (!ParseManifest(manifestPath, shaderPackage, error))
			return false;

		if (packagesById.find(shaderPackage.id) != packagesById.end())
		{
			error = "Duplicate shader id found: " + shaderPackage.id;
			return false;
		}

		packageOrder.push_back(shaderPackage.id);
		packagesById[shaderPackage.id] = shaderPackage;
	}

	std::sort(packageOrder.begin(), packageOrder.end());
	return true;
}

bool ShaderPackageRegistry::ParseManifest(const std::filesystem::path& manifestPath, ShaderPackage& shaderPackage, std::string& error) const
{
	const std::string manifestText = ReadTextFile(manifestPath, error);
	if (manifestText.empty())
		return false;

	JsonValue manifestJson;
	if (!ParseJson(manifestText, manifestJson, error))
		return false;
	if (!manifestJson.isObject())
	{
		error = "Shader manifest root must be an object: " + manifestPath.string();
		return false;
	}

	if (!ParseShaderMetadata(manifestJson, shaderPackage, manifestPath, error))
		return false;

	if (!std::filesystem::exists(shaderPackage.shaderPath))
	{
		error = "Shader source not found for package " + shaderPackage.id + ": " + shaderPackage.shaderPath.string();
		return false;
	}

	shaderPackage.shaderWriteTime = std::filesystem::last_write_time(shaderPackage.shaderPath);
	shaderPackage.manifestWriteTime = std::filesystem::last_write_time(shaderPackage.manifestPath);

	return ParseTextureAssets(manifestJson, shaderPackage, manifestPath, error) &&
		ParseTemporalSettings(manifestJson, shaderPackage, mMaxTemporalHistoryFrames, manifestPath, error) &&
		ParseParameterDefinitions(manifestJson, shaderPackage, manifestPath, error);
}