#include "RuntimeRenderScene.h"

#include "../../platform/HiddenGlWindow.h"

#include <algorithm>
#include <functional>
#include <utility>

#ifndef GL_FRAMEBUFFER_BINDING
#define GL_FRAMEBUFFER_BINDING 0x8CA6
#endif

RuntimeRenderScene::~RuntimeRenderScene()
{
	ShutdownGl();
}

bool RuntimeRenderScene::StartPrepareWorker(std::unique_ptr<HiddenGlWindow> sharedWindow, std::string& error)
{
	return mPrepareWorker.Start(std::move(sharedWindow), error);
}

bool RuntimeRenderScene::CommitRenderLayers(const std::vector<RenderCadenceCompositor::RuntimeRenderLayerModel>& layers, std::string& error)
{
	ConsumePreparedPrograms();

	std::vector<std::string> nextOrder;
	std::vector<RenderCadenceCompositor::RuntimeRenderLayerModel> layersToPrepare;
	nextOrder.reserve(layers.size());
	for (const RenderCadenceCompositor::RuntimeRenderLayerModel& layer : layers)
		nextOrder.push_back(layer.id);

	for (auto layerIt = mLayers.begin(); layerIt != mLayers.end();)
	{
		const bool stillPresent = std::find(nextOrder.begin(), nextOrder.end(), layerIt->layerId) != nextOrder.end();
		if (stillPresent)
		{
			++layerIt;
			continue;
		}

		for (LayerProgram::PassProgram& pass : layerIt->passes)
		{
			if (pass.renderer)
				pass.renderer->ShutdownGl();
		}
		ReleasePendingPrograms(*layerIt);
		layerIt = mLayers.erase(layerIt);
	}

	for (const RenderCadenceCompositor::RuntimeRenderLayerModel& layer : layers)
	{
		if (layer.artifact.passes.empty() && layer.artifact.fragmentShaderSource.empty())
			continue;

		const std::string fingerprint = Fingerprint(layer.artifact);
		LayerProgram* program = FindLayer(layer.id);
		if (!program)
		{
			LayerProgram next;
			next.layerId = layer.id;
			mLayers.push_back(std::move(next));
			program = &mLayers.back();
		}

		bool hasReadyPass = false;
		for (const LayerProgram::PassProgram& pass : program->passes)
		{
			if (pass.renderer && pass.renderer->HasProgram())
			{
				hasReadyPass = true;
				break;
			}
		}
		if (program->shaderId == layer.shaderId && program->sourceFingerprint == fingerprint && hasReadyPass)
			continue;
		if (program->pendingFingerprint == fingerprint)
			continue;

		ReleasePendingPrograms(*program);
		program->shaderId = layer.shaderId;
		program->pendingFingerprint = fingerprint;
		layersToPrepare.push_back(layer);
	}

	mLayerOrder = std::move(nextOrder);
	if (!layersToPrepare.empty())
		mPrepareWorker.Submit(layersToPrepare);
	error.clear();
	return true;
}

bool RuntimeRenderScene::HasLayers()
{
	ConsumePreparedPrograms();

	for (const std::string& layerId : mLayerOrder)
	{
		const LayerProgram* layer = FindLayer(layerId);
		if (!layer)
			continue;
		for (const LayerProgram::PassProgram& pass : layer->passes)
		{
			if (pass.renderer && pass.renderer->HasProgram())
				return true;
		}
	}
	return false;
}

void RuntimeRenderScene::RenderFrame(uint64_t frameIndex, unsigned width, unsigned height)
{
	ConsumePreparedPrograms();

	std::vector<LayerProgram*> readyLayers;
	for (const std::string& layerId : mLayerOrder)
	{
		LayerProgram* layer = FindLayer(layerId);
		if (!layer)
			continue;
		for (const LayerProgram::PassProgram& pass : layer->passes)
		{
			if (pass.renderer && pass.renderer->HasProgram())
			{
				readyLayers.push_back(layer);
				break;
			}
		}
	}

	if (readyLayers.empty())
		return;

	GLint outputFramebuffer = 0;
	glGetIntegerv(GL_FRAMEBUFFER_BINDING, &outputFramebuffer);

	if (readyLayers.size() == 1)
	{
		RenderLayer(*readyLayers.front(), frameIndex, width, height, 0, static_cast<GLuint>(outputFramebuffer), true);
		return;
	}

	if (!EnsureLayerTargets(width, height))
	{
		glBindFramebuffer(GL_FRAMEBUFFER, static_cast<GLuint>(outputFramebuffer));
		RenderLayer(*readyLayers.back(), frameIndex, width, height, 0, static_cast<GLuint>(outputFramebuffer), true);
		return;
	}

	GLuint layerInputTexture = 0;
	std::size_t nextTargetIndex = 0;
	for (std::size_t layerIndex = 0; layerIndex < readyLayers.size(); ++layerIndex)
	{
		const bool isFinalLayer = layerIndex == readyLayers.size() - 1;
		if (isFinalLayer)
		{
			glBindFramebuffer(GL_FRAMEBUFFER, static_cast<GLuint>(outputFramebuffer));
			RenderLayer(*readyLayers[layerIndex], frameIndex, width, height, layerInputTexture, static_cast<GLuint>(outputFramebuffer), true);
			continue;
		}

		RenderLayer(*readyLayers[layerIndex], frameIndex, width, height, layerInputTexture, mLayerFramebuffers[nextTargetIndex], true);
		layerInputTexture = mLayerTextures[nextTargetIndex];
		nextTargetIndex = 1 - nextTargetIndex;
	}
}

void RuntimeRenderScene::ShutdownGl()
{
	mPrepareWorker.Stop();
	for (LayerProgram& layer : mLayers)
	{
		for (LayerProgram::PassProgram& pass : layer.passes)
		{
			if (pass.renderer)
				pass.renderer->ShutdownGl();
		}
		ReleasePendingPrograms(layer);
	}
	mLayers.clear();
	mLayerOrder.clear();
	DestroyLayerTargets();
}

void RuntimeRenderScene::ConsumePreparedPrograms()
{
	RuntimePreparedShaderProgram preparedProgram;
	while (mPrepareWorker.TryConsume(preparedProgram))
	{
		if (!preparedProgram.succeeded)
		{
			preparedProgram.ReleaseGl();
			continue;
		}

		LayerProgram* layer = FindLayer(preparedProgram.layerId);
		if (!layer || layer->pendingFingerprint != preparedProgram.sourceFingerprint)
		{
			preparedProgram.ReleaseGl();
			continue;
		}

		bool replacesExistingPendingPass = false;
		for (RuntimePreparedShaderProgram& existing : layer->pendingPreparedPrograms)
		{
			if (existing.passId != preparedProgram.passId)
				continue;
			existing.ReleaseGl();
			existing = std::move(preparedProgram);
			replacesExistingPendingPass = true;
			break;
		}
		if (!replacesExistingPendingPass)
			layer->pendingPreparedPrograms.push_back(std::move(preparedProgram));
		TryCommitPendingPrograms(*layer);
	}
}

void RuntimeRenderScene::ReleasePendingPrograms(LayerProgram& layer)
{
	for (RuntimePreparedShaderProgram& program : layer.pendingPreparedPrograms)
		program.ReleaseGl();
	layer.pendingPreparedPrograms.clear();
}

void RuntimeRenderScene::TryCommitPendingPrograms(LayerProgram& layer)
{
	if (layer.pendingPreparedPrograms.empty())
		return;

	const RuntimeShaderArtifact& artifact = layer.pendingPreparedPrograms.front().artifact;
	const std::size_t expectedPassCount = artifact.passes.empty() ? 1 : artifact.passes.size();
	if (layer.pendingPreparedPrograms.size() < expectedPassCount)
		return;

	std::vector<LayerProgram::PassProgram> nextPasses;
	nextPasses.reserve(expectedPassCount);
	for (const RuntimeShaderPassArtifact& passArtifact : artifact.passes)
	{
		auto preparedIt = std::find_if(
			layer.pendingPreparedPrograms.begin(),
			layer.pendingPreparedPrograms.end(),
			[&passArtifact](const RuntimePreparedShaderProgram& prepared) {
				return prepared.passId == passArtifact.passId;
			});
		if (preparedIt == layer.pendingPreparedPrograms.end())
			return;

		std::unique_ptr<RuntimeShaderRenderer> nextRenderer = std::make_unique<RuntimeShaderRenderer>();
		std::string error;
		if (!nextRenderer->CommitPreparedProgram(*preparedIt, error))
		{
			ReleasePendingPrograms(layer);
			return;
		}

		LayerProgram::PassProgram nextPass;
		nextPass.passId = preparedIt->passId;
		nextPass.inputNames = preparedIt->inputNames;
		nextPass.outputName = preparedIt->outputName.empty() ? preparedIt->passId : preparedIt->outputName;
		nextPass.renderer = std::move(nextRenderer);
		nextPasses.push_back(std::move(nextPass));
	}
	if (artifact.passes.empty())
	{
		RuntimePreparedShaderProgram& prepared = layer.pendingPreparedPrograms.front();
		std::unique_ptr<RuntimeShaderRenderer> nextRenderer = std::make_unique<RuntimeShaderRenderer>();
		std::string error;
		if (!nextRenderer->CommitPreparedProgram(prepared, error))
		{
			ReleasePendingPrograms(layer);
			return;
		}

		LayerProgram::PassProgram nextPass;
		nextPass.passId = prepared.passId;
		nextPass.inputNames = prepared.inputNames;
		nextPass.outputName = prepared.outputName.empty() ? prepared.passId : prepared.outputName;
		nextPass.renderer = std::move(nextRenderer);
		nextPasses.push_back(std::move(nextPass));
	}

	for (LayerProgram::PassProgram& pass : layer.passes)
	{
		if (pass.renderer)
			pass.renderer->ShutdownGl();
	}
	layer.passes = std::move(nextPasses);
	layer.shaderId = artifact.shaderId;
	layer.sourceFingerprint = layer.pendingPreparedPrograms.front().sourceFingerprint;
	layer.pendingFingerprint.clear();
	layer.pendingPreparedPrograms.clear();
}

GLuint RuntimeRenderScene::RenderLayer(
	LayerProgram& layer,
	uint64_t frameIndex,
	unsigned width,
	unsigned height,
	GLuint layerInputTexture,
	GLuint outputFramebuffer,
	bool renderToOutput)
{
	GLuint namedOutputs[2] = {};
	std::string namedOutputNames[2];
	std::size_t nextTargetIndex = 2;
	GLuint lastOutputTexture = layerInputTexture;

	for (LayerProgram::PassProgram& pass : layer.passes)
	{
		if (!pass.renderer || !pass.renderer->HasProgram())
			continue;

		GLuint sourceTexture = layerInputTexture;
		if (!pass.inputNames.empty())
		{
			const std::string& inputName = pass.inputNames.front();
			if (inputName != "layerInput" && inputName != "videoInput")
			{
				for (std::size_t index = 0; index < 2; ++index)
				{
					if (namedOutputNames[index] == inputName)
					{
						sourceTexture = namedOutputs[index];
						break;
					}
				}
			}
		}

		const bool writesLayerOutput = pass.outputName == "layerOutput";
		if (writesLayerOutput && renderToOutput)
		{
			glBindFramebuffer(GL_FRAMEBUFFER, outputFramebuffer);
			pass.renderer->RenderFrame(frameIndex, width, height, sourceTexture, sourceTexture);
			lastOutputTexture = 0;
			continue;
		}

		if (!EnsureLayerTargets(width, height))
			continue;

		const std::size_t targetIndex = nextTargetIndex;
		nextTargetIndex = nextTargetIndex == 2 ? 3 : 2;
		glBindFramebuffer(GL_FRAMEBUFFER, mLayerFramebuffers[targetIndex]);
		pass.renderer->RenderFrame(frameIndex, width, height, sourceTexture, sourceTexture);
		const std::size_t namedIndex = targetIndex - 2;
		namedOutputs[namedIndex] = mLayerTextures[targetIndex];
		namedOutputNames[namedIndex] = pass.outputName;
		lastOutputTexture = mLayerTextures[targetIndex];
	}

	return lastOutputTexture;
}

bool RuntimeRenderScene::EnsureLayerTargets(unsigned width, unsigned height)
{
	if (width == 0 || height == 0)
		return false;
	if (mLayerFramebuffers[0] != 0 && mLayerFramebuffers[1] != 0 && mLayerFramebuffers[2] != 0 && mLayerFramebuffers[3] != 0
		&& mLayerTextures[0] != 0 && mLayerTextures[1] != 0 && mLayerTextures[2] != 0 && mLayerTextures[3] != 0
		&& mLayerTargetWidth == width && mLayerTargetHeight == height)
		return true;

	DestroyLayerTargets();
	mLayerTargetWidth = width;
	mLayerTargetHeight = height;

	glGenFramebuffers(4, mLayerFramebuffers);
	glGenTextures(4, mLayerTextures);
	for (int index = 0; index < 4; ++index)
	{
		glBindTexture(GL_TEXTURE_2D, mLayerTextures[index]);
		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);
		glTexImage2D(
			GL_TEXTURE_2D,
			0,
			GL_RGBA8,
			static_cast<GLsizei>(width),
			static_cast<GLsizei>(height),
			0,
			GL_BGRA,
			GL_UNSIGNED_INT_8_8_8_8_REV,
			nullptr);

		glBindFramebuffer(GL_FRAMEBUFFER, mLayerFramebuffers[index]);
		glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, mLayerTextures[index], 0);
		if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
		{
			glBindFramebuffer(GL_FRAMEBUFFER, 0);
			glBindTexture(GL_TEXTURE_2D, 0);
			DestroyLayerTargets();
			return false;
		}
	}

	glBindFramebuffer(GL_FRAMEBUFFER, 0);
	glBindTexture(GL_TEXTURE_2D, 0);
	return true;
}

void RuntimeRenderScene::DestroyLayerTargets()
{
	if (mLayerFramebuffers[0] != 0 || mLayerFramebuffers[1] != 0 || mLayerFramebuffers[2] != 0 || mLayerFramebuffers[3] != 0)
		glDeleteFramebuffers(4, mLayerFramebuffers);
	if (mLayerTextures[0] != 0 || mLayerTextures[1] != 0 || mLayerTextures[2] != 0 || mLayerTextures[3] != 0)
		glDeleteTextures(4, mLayerTextures);
	for (int index = 0; index < 4; ++index)
	{
		mLayerFramebuffers[index] = 0;
		mLayerTextures[index] = 0;
	}
	mLayerTargetWidth = 0;
	mLayerTargetHeight = 0;
}

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

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

RuntimeRenderScene::LayerProgram::PassProgram* RuntimeRenderScene::FindPass(LayerProgram& layer, const std::string& passId)
{
	for (LayerProgram::PassProgram& pass : layer.passes)
	{
		if (pass.passId == passId)
			return &pass;
	}
	return nullptr;
}

std::string RuntimeRenderScene::Fingerprint(const RuntimeShaderArtifact& artifact)
{
	const std::hash<std::string> hasher;
	std::string source;
	for (const RuntimeShaderPassArtifact& pass : artifact.passes)
		source += pass.passId + ":" + pass.outputName + ":" + pass.fragmentShaderSource + "\n";
	if (source.empty())
		source = artifact.fragmentShaderSource;
	return artifact.shaderId + ":" + std::to_string(source.size()) + ":" + std::to_string(hasher(source));
}