#include "OpenGLShaderPrograms.h"

#include "GlRenderConstants.h"
#include "GlScopedObjects.h"
#include "GlShaderSources.h"
#include "Std140Buffer.h"
#include "TextRasterizer.h"
#include "TextureAssetLoader.h"

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <filesystem>
#include <vector>

namespace
{
void CopyErrorMessage(const std::string& message, int errorMessageSize, char* errorMessage)
{
	if (!errorMessage || errorMessageSize <= 0)
		return;

	strncpy_s(errorMessage, errorMessageSize, message.c_str(), _TRUNCATE);
}

std::string TextValueForBinding(const RuntimeRenderState& state, const std::string& parameterId)
{
	auto valueIt = state.parameterValues.find(parameterId);
	return valueIt == state.parameterValues.end() ? std::string() : valueIt->second.textValue;
}

const ShaderFontAsset* FindFontAssetForParameter(const RuntimeRenderState& state, const ShaderParameterDefinition& definition)
{
	if (!definition.fontId.empty())
	{
		for (const ShaderFontAsset& fontAsset : state.fontAssets)
		{
			if (fontAsset.id == definition.fontId)
				return &fontAsset;
		}
	}
	return state.fontAssets.empty() ? nullptr : &state.fontAssets.front();
}

GLint FindSamplerUniformLocation(GLuint program, const std::string& samplerName)
{
	GLint location = glGetUniformLocation(program, samplerName.c_str());
	if (location >= 0)
		return location;
	return glGetUniformLocation(program, (samplerName + "_0").c_str());
}
}

OpenGLShaderPrograms::OpenGLShaderPrograms(OpenGLRenderer& renderer, RuntimeHost& runtimeHost) :
	mRenderer(renderer),
	mRuntimeHost(runtimeHost)
{
}

bool OpenGLShaderPrograms::CompileLayerPrograms(unsigned inputFrameWidth, unsigned inputFrameHeight, int errorMessageSize, char* errorMessage)
{
	const std::vector<RuntimeRenderState> layerStates = mRuntimeHost.GetLayerRenderStates(inputFrameWidth, inputFrameHeight);
	std::string temporalError;
	const unsigned historyCap = mRuntimeHost.GetMaxTemporalHistoryFrames();
	if (!mRenderer.TemporalHistory().ValidateTextureUnitBudget(layerStates, historyCap, temporalError))
	{
		CopyErrorMessage(temporalError, errorMessageSize, errorMessage);
		return false;
	}
	if (!mRenderer.TemporalHistory().EnsureResources(layerStates, historyCap, inputFrameWidth, inputFrameHeight, temporalError))
	{
		CopyErrorMessage(temporalError, errorMessageSize, errorMessage);
		return false;
	}

	std::vector<LayerProgram> newPrograms;
	newPrograms.reserve(layerStates.size());

	for (const RuntimeRenderState& state : layerStates)
	{
		LayerProgram layerProgram;
		if (!CompileSingleLayerProgram(state, layerProgram, errorMessageSize, errorMessage))
		{
			for (LayerProgram& program : newPrograms)
				DestroySingleLayerProgram(program);
			return false;
		}
		newPrograms.push_back(layerProgram);
	}

	DestroyLayerPrograms();
	mRenderer.ReplaceLayerPrograms(newPrograms);

	mRuntimeHost.SetCompileStatus(true, "Shader layers compiled successfully.");
	mRuntimeHost.ClearReloadRequest();

	return true;
}

bool OpenGLShaderPrograms::CompileSingleLayerProgram(const RuntimeRenderState& state, LayerProgram& layerProgram, int errorMessageSize, char* errorMessage)
{
	GLsizei errorBufferSize = 0;
	GLint compileResult = GL_FALSE;
	GLint linkResult = GL_FALSE;
	std::string fragmentShaderSource;
	std::string loadError;
	std::vector<LayerProgram::TextureBinding> textureBindings;
	const char* vertexSource = kFullscreenTriangleVertexShaderSource;

	if (!mRuntimeHost.BuildLayerFragmentShaderSource(state.layerId, fragmentShaderSource, loadError))
	{
		CopyErrorMessage(loadError, errorMessageSize, errorMessage);
		return false;
	}

	const char* fragmentSource = fragmentShaderSource.c_str();

	ScopedGlShader newVertexShader(glCreateShader(GL_VERTEX_SHADER));
	glShaderSource(newVertexShader.get(), 1, (const GLchar**)&vertexSource, NULL);
	glCompileShader(newVertexShader.get());
	glGetShaderiv(newVertexShader.get(), GL_COMPILE_STATUS, &compileResult);
	if (compileResult == GL_FALSE)
	{
		glGetShaderInfoLog(newVertexShader.get(), errorMessageSize, &errorBufferSize, errorMessage);
		return false;
	}

	ScopedGlShader newFragmentShader(glCreateShader(GL_FRAGMENT_SHADER));
	glShaderSource(newFragmentShader.get(), 1, (const GLchar**)&fragmentSource, NULL);
	glCompileShader(newFragmentShader.get());
	glGetShaderiv(newFragmentShader.get(), GL_COMPILE_STATUS, &compileResult);
	if (compileResult == GL_FALSE)
	{
		glGetShaderInfoLog(newFragmentShader.get(), errorMessageSize, &errorBufferSize, errorMessage);
		return false;
	}

	ScopedGlProgram newProgram(glCreateProgram());
	glAttachShader(newProgram.get(), newVertexShader.get());
	glAttachShader(newProgram.get(), newFragmentShader.get());
	glLinkProgram(newProgram.get());
	glGetProgramiv(newProgram.get(), GL_LINK_STATUS, &linkResult);
	if (linkResult == GL_FALSE)
	{
		glGetProgramInfoLog(newProgram.get(), errorMessageSize, &errorBufferSize, errorMessage);
		return false;
	}

	for (const ShaderTextureAsset& textureAsset : state.textureAssets)
	{
		LayerProgram::TextureBinding textureBinding;
		textureBinding.samplerName = textureAsset.id;
		textureBinding.sourcePath = textureAsset.path;
		if (!LoadTextureAsset(textureAsset, textureBinding.texture, loadError))
		{
			for (LayerProgram::TextureBinding& loadedTexture : textureBindings)
			{
				if (loadedTexture.texture != 0)
					glDeleteTextures(1, &loadedTexture.texture);
			}
			CopyErrorMessage(loadError, errorMessageSize, errorMessage);
			return false;
		}
		textureBindings.push_back(textureBinding);
	}

	std::vector<LayerProgram::TextBinding> textBindings;
	for (const ShaderParameterDefinition& definition : state.parameterDefinitions)
	{
		if (definition.type != ShaderParameterType::Text)
			continue;
		LayerProgram::TextBinding textBinding;
		textBinding.parameterId = definition.id;
		textBinding.samplerName = definition.id + "Texture";
		textBinding.fontId = definition.fontId;
		glGenTextures(1, &textBinding.texture);
		glBindTexture(GL_TEXTURE_2D, textBinding.texture);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
		std::vector<unsigned char> empty(static_cast<std::size_t>(kTextTextureWidth) * kTextTextureHeight * 4, 0);
		glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, kTextTextureWidth, kTextTextureHeight, 0, GL_BGRA, GL_UNSIGNED_INT_8_8_8_8_REV, empty.data());
		glBindTexture(GL_TEXTURE_2D, 0);
		textBindings.push_back(textBinding);
	}

	const GLuint globalParamsIndex = glGetUniformBlockIndex(newProgram.get(), "GlobalParams");
	if (globalParamsIndex != GL_INVALID_INDEX)
		glUniformBlockBinding(newProgram.get(), globalParamsIndex, kGlobalParamsBindingPoint);

	const unsigned historyCap = mRuntimeHost.GetMaxTemporalHistoryFrames();
	const GLuint shaderTextureBase = state.isTemporal ? kSourceHistoryTextureUnitBase + historyCap + historyCap : kSourceHistoryTextureUnitBase;
	glUseProgram(newProgram.get());
	const GLint videoInputLocation = glGetUniformLocation(newProgram.get(), "gVideoInput");
	if (videoInputLocation >= 0)
		glUniform1i(videoInputLocation, static_cast<GLint>(kDecodedVideoTextureUnit));
	for (unsigned index = 0; index < historyCap; ++index)
	{
		const std::string sourceSamplerName = "gSourceHistory" + std::to_string(index);
		const GLint sourceSamplerLocation = glGetUniformLocation(newProgram.get(), sourceSamplerName.c_str());
		if (sourceSamplerLocation >= 0)
			glUniform1i(sourceSamplerLocation, static_cast<GLint>(kSourceHistoryTextureUnitBase + index));

		const std::string temporalSamplerName = "gTemporalHistory" + std::to_string(index);
		const GLint temporalSamplerLocation = glGetUniformLocation(newProgram.get(), temporalSamplerName.c_str());
		if (temporalSamplerLocation >= 0)
			glUniform1i(temporalSamplerLocation, static_cast<GLint>(kSourceHistoryTextureUnitBase + historyCap + index));
	}
	for (std::size_t index = 0; index < textureBindings.size(); ++index)
	{
		const GLint textureSamplerLocation = FindSamplerUniformLocation(newProgram.get(), textureBindings[index].samplerName);
		if (textureSamplerLocation >= 0)
			glUniform1i(textureSamplerLocation, static_cast<GLint>(shaderTextureBase + static_cast<GLuint>(index)));
	}
	const GLuint textTextureBase = shaderTextureBase + static_cast<GLuint>(textureBindings.size());
	for (std::size_t index = 0; index < textBindings.size(); ++index)
	{
		const GLint textSamplerLocation = FindSamplerUniformLocation(newProgram.get(), textBindings[index].samplerName);
		if (textSamplerLocation >= 0)
			glUniform1i(textSamplerLocation, static_cast<GLint>(textTextureBase + static_cast<GLuint>(index)));
	}
	glUseProgram(0);

	layerProgram.layerId = state.layerId;
	layerProgram.shaderId = state.shaderId;
	layerProgram.shaderTextureBase = shaderTextureBase;
	layerProgram.program = newProgram.release();
	layerProgram.vertexShader = newVertexShader.release();
	layerProgram.fragmentShader = newFragmentShader.release();
	layerProgram.textureBindings.swap(textureBindings);
	layerProgram.textBindings.swap(textBindings);
	return true;
}

bool OpenGLShaderPrograms::CompileDecodeShader(int errorMessageSize, char* errorMessage)
{
	GLsizei errorBufferSize = 0;
	GLint compileResult = GL_FALSE;
	GLint linkResult = GL_FALSE;
	const char* vertexSource = kFullscreenTriangleVertexShaderSource;
	const char* fragmentSource = kDecodeFragmentShaderSource;

	ScopedGlShader newVertexShader(glCreateShader(GL_VERTEX_SHADER));
	glShaderSource(newVertexShader.get(), 1, (const GLchar**)&vertexSource, NULL);
	glCompileShader(newVertexShader.get());
	glGetShaderiv(newVertexShader.get(), GL_COMPILE_STATUS, &compileResult);
	if (compileResult == GL_FALSE)
	{
		glGetShaderInfoLog(newVertexShader.get(), errorMessageSize, &errorBufferSize, errorMessage);
		return false;
	}

	ScopedGlShader newFragmentShader(glCreateShader(GL_FRAGMENT_SHADER));
	glShaderSource(newFragmentShader.get(), 1, (const GLchar**)&fragmentSource, NULL);
	glCompileShader(newFragmentShader.get());
	glGetShaderiv(newFragmentShader.get(), GL_COMPILE_STATUS, &compileResult);
	if (compileResult == GL_FALSE)
	{
		glGetShaderInfoLog(newFragmentShader.get(), errorMessageSize, &errorBufferSize, errorMessage);
		return false;
	}

	ScopedGlProgram newProgram(glCreateProgram());
	glAttachShader(newProgram.get(), newVertexShader.get());
	glAttachShader(newProgram.get(), newFragmentShader.get());
	glLinkProgram(newProgram.get());
	glGetProgramiv(newProgram.get(), GL_LINK_STATUS, &linkResult);
	if (linkResult == GL_FALSE)
	{
		glGetProgramInfoLog(newProgram.get(), errorMessageSize, &errorBufferSize, errorMessage);
		return false;
	}

	DestroyDecodeShaderProgram();
	mRenderer.SetDecodeShaderProgram(newProgram.release(), newVertexShader.release(), newFragmentShader.release());
	return true;
}

void OpenGLShaderPrograms::DestroySingleLayerProgram(LayerProgram& layerProgram)
{
	mRenderer.DestroySingleLayerProgram(layerProgram);
}

void OpenGLShaderPrograms::DestroyLayerPrograms()
{
	mRenderer.DestroyLayerPrograms();
}

void OpenGLShaderPrograms::DestroyDecodeShaderProgram()
{
	mRenderer.DestroyDecodeShaderProgram();
}

void OpenGLShaderPrograms::ResetTemporalHistoryState()
{
	mRenderer.TemporalHistory().ResetState();
}

bool OpenGLShaderPrograms::LoadTextureAsset(const ShaderTextureAsset& textureAsset, GLuint& textureId, std::string& error)
{
	return ::LoadTextureAsset(textureAsset, textureId, error);
}

bool OpenGLShaderPrograms::UpdateTextBindingTexture(const RuntimeRenderState& state, LayerProgram::TextBinding& textBinding, std::string& error)
{
	const std::string text = TextValueForBinding(state, textBinding.parameterId);
	if (text == textBinding.renderedText && textBinding.renderedWidth == kTextTextureWidth && textBinding.renderedHeight == kTextTextureHeight)
		return true;

	auto definitionIt = std::find_if(state.parameterDefinitions.begin(), state.parameterDefinitions.end(),
		[&textBinding](const ShaderParameterDefinition& definition) { return definition.id == textBinding.parameterId; });
	if (definitionIt == state.parameterDefinitions.end())
		return true;

	const ShaderFontAsset* fontAsset = FindFontAssetForParameter(state, *definitionIt);
	std::filesystem::path fontPath;
	if (fontAsset)
		fontPath = fontAsset->path;

	std::vector<unsigned char> sdf;
	if (!RasterizeTextSdf(text, fontPath, sdf, error))
		return false;

	GLint previousActiveTexture = 0;
	GLint previousUnpackBuffer = 0;
	glGetIntegerv(GL_ACTIVE_TEXTURE, &previousActiveTexture);
	glGetIntegerv(GL_PIXEL_UNPACK_BUFFER_BINDING, &previousUnpackBuffer);
	glActiveTexture(GL_TEXTURE0);
	glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);
	glBindTexture(GL_TEXTURE_2D, textBinding.texture);
	glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, kTextTextureWidth, kTextTextureHeight, GL_BGRA, GL_UNSIGNED_INT_8_8_8_8_REV, sdf.data());
	glBindTexture(GL_TEXTURE_2D, 0);
	glBindBuffer(GL_PIXEL_UNPACK_BUFFER, static_cast<GLuint>(previousUnpackBuffer));
	glActiveTexture(static_cast<GLenum>(previousActiveTexture));

	textBinding.renderedText = text;
	textBinding.renderedWidth = kTextTextureWidth;
	textBinding.renderedHeight = kTextTextureHeight;
	return true;
}

bool OpenGLShaderPrograms::UpdateGlobalParamsBuffer(const RuntimeRenderState& state, unsigned availableSourceHistoryLength, unsigned availableTemporalHistoryLength)
{
	std::vector<unsigned char> buffer;
	buffer.reserve(512);

	AppendStd140Float(buffer, static_cast<float>(state.timeSeconds));
	AppendStd140Vec2(buffer, static_cast<float>(state.inputWidth), static_cast<float>(state.inputHeight));
	AppendStd140Vec2(buffer, static_cast<float>(state.outputWidth), static_cast<float>(state.outputHeight));
	AppendStd140Float(buffer, static_cast<float>(state.frameCount));
	AppendStd140Float(buffer, static_cast<float>(state.mixAmount));
	AppendStd140Float(buffer, static_cast<float>(state.bypass));
	const unsigned effectiveSourceHistoryLength = availableSourceHistoryLength < state.effectiveTemporalHistoryLength
		? availableSourceHistoryLength
		: state.effectiveTemporalHistoryLength;
	const unsigned effectiveTemporalHistoryLength = (state.temporalHistorySource == TemporalHistorySource::PreLayerInput)
		? (availableTemporalHistoryLength < state.effectiveTemporalHistoryLength ? availableTemporalHistoryLength : state.effectiveTemporalHistoryLength)
		: 0u;
	AppendStd140Int(buffer, static_cast<int>(effectiveSourceHistoryLength));
	AppendStd140Int(buffer, static_cast<int>(effectiveTemporalHistoryLength));

	for (const ShaderParameterDefinition& definition : state.parameterDefinitions)
	{
		auto valueIt = state.parameterValues.find(definition.id);
		const ShaderParameterValue value = valueIt != state.parameterValues.end()
			? valueIt->second
			: ShaderParameterValue();

		switch (definition.type)
		{
		case ShaderParameterType::Float:
			AppendStd140Float(buffer, value.numberValues.empty() ? 0.0f : static_cast<float>(value.numberValues[0]));
			break;
		case ShaderParameterType::Vec2:
			AppendStd140Vec2(buffer,
				value.numberValues.size() > 0 ? static_cast<float>(value.numberValues[0]) : 0.0f,
				value.numberValues.size() > 1 ? static_cast<float>(value.numberValues[1]) : 0.0f);
			break;
		case ShaderParameterType::Color:
			AppendStd140Vec4(buffer,
				value.numberValues.size() > 0 ? static_cast<float>(value.numberValues[0]) : 1.0f,
				value.numberValues.size() > 1 ? static_cast<float>(value.numberValues[1]) : 1.0f,
				value.numberValues.size() > 2 ? static_cast<float>(value.numberValues[2]) : 1.0f,
				value.numberValues.size() > 3 ? static_cast<float>(value.numberValues[3]) : 1.0f);
			break;
		case ShaderParameterType::Boolean:
			AppendStd140Int(buffer, value.booleanValue ? 1 : 0);
			break;
		case ShaderParameterType::Enum:
		{
			int selectedIndex = 0;
			for (std::size_t optionIndex = 0; optionIndex < definition.enumOptions.size(); ++optionIndex)
			{
				if (definition.enumOptions[optionIndex].value == value.enumValue)
				{
					selectedIndex = static_cast<int>(optionIndex);
					break;
				}
			}
			AppendStd140Int(buffer, selectedIndex);
			break;
		}
		case ShaderParameterType::Text:
			break;
		}
	}

	buffer.resize(AlignStd140(buffer.size(), 16), 0);

	glBindBuffer(GL_UNIFORM_BUFFER, mRenderer.GlobalParamsUBO());
	if (mRenderer.GlobalParamsUBOSize() != static_cast<GLsizeiptr>(buffer.size()))
	{
		glBufferData(GL_UNIFORM_BUFFER, static_cast<GLsizeiptr>(buffer.size()), buffer.data(), GL_DYNAMIC_DRAW);
		mRenderer.SetGlobalParamsUBOSize(static_cast<GLsizeiptr>(buffer.size()));
	}
	else
	{
		glBufferSubData(GL_UNIFORM_BUFFER, 0, static_cast<GLsizeiptr>(buffer.size()), buffer.data());
	}
	glBindBufferBase(GL_UNIFORM_BUFFER, kGlobalParamsBindingPoint, mRenderer.GlobalParamsUBO());
	glBindBuffer(GL_UNIFORM_BUFFER, 0);

	return true;
}