#include "TemporalHistoryBuffers.h"

#include "GlRenderConstants.h"
#include "ShaderTypes.h"

#include <algorithm>
#include <sstream>
#include <set>

bool TemporalHistoryBuffers::ValidateTextureUnitBudget(const std::vector<RuntimeRenderState>& layerStates, unsigned historyCap, std::string& error) const
{
	unsigned requiredUnits = kSourceHistoryTextureUnitBase;
	for (const RuntimeRenderState& state : layerStates)
	{
		unsigned textTextureCount = 0;
		for (const ShaderParameterDefinition& definition : state.parameterDefinitions)
		{
			if (definition.type == ShaderParameterType::Text)
				++textTextureCount;
		}
		const unsigned totalShaderTextures = static_cast<unsigned>(state.textureAssets.size()) + textTextureCount;
		const unsigned feedbackTextureCount = state.feedback.enabled ? 1u : 0u;
		const unsigned layerRequiredUnits = kSourceHistoryTextureUnitBase + (state.isTemporal ? historyCap + historyCap : 0u) + feedbackTextureCount + totalShaderTextures;
		if (layerRequiredUnits > requiredUnits)
			requiredUnits = layerRequiredUnits;
	}

	GLint maxTextureUnits = 0;
	glGetIntegerv(GL_MAX_TEXTURE_IMAGE_UNITS, &maxTextureUnits);
	const unsigned availableUnits = maxTextureUnits > 0 ? static_cast<unsigned>(maxTextureUnits) : 0u;
	if (requiredUnits > availableUnits)
	{
		std::ostringstream message;
		message << "The current history and shader texture asset configuration requires " << requiredUnits
			<< " fragment texture units, but only " << maxTextureUnits << " are available.";
		error = message.str();
		return false;
	}
	return true;
}

bool TemporalHistoryBuffers::EnsureResources(const std::vector<RuntimeRenderState>& layerStates, unsigned historyCap, unsigned frameWidth, unsigned frameHeight, std::string& error)
{
	const bool sourceHistoryNeeded = std::any_of(layerStates.begin(), layerStates.end(),
		[](const RuntimeRenderState& state) { return state.isTemporal && state.effectiveTemporalHistoryLength > 0; });
	const unsigned sourceHistoryLength = sourceHistoryNeeded ? historyCap : 0;

	if (sourceHistoryRing.effectiveLength != sourceHistoryLength)
	{
		if (!CreateRing(sourceHistoryRing, sourceHistoryLength, TemporalHistorySource::Source, frameWidth, frameHeight, error))
			return false;
		mNeedsReset = true;
	}

	std::set<std::string> requiredPreLayerIds;
	for (const RuntimeRenderState& state : layerStates)
	{
		if (!state.isTemporal || state.temporalHistorySource != TemporalHistorySource::PreLayerInput)
			continue;
		requiredPreLayerIds.insert(state.layerId);
		auto historyIt = preLayerHistoryByLayerId.find(state.layerId);
		if (historyIt == preLayerHistoryByLayerId.end() || historyIt->second.effectiveLength != state.effectiveTemporalHistoryLength)
		{
			Ring replacement;
			if (!CreateRing(replacement, state.effectiveTemporalHistoryLength, TemporalHistorySource::PreLayerInput, frameWidth, frameHeight, error))
				return false;
			preLayerHistoryByLayerId[state.layerId] = std::move(replacement);
			mNeedsReset = true;
		}
	}

	for (auto it = preLayerHistoryByLayerId.begin(); it != preLayerHistoryByLayerId.end();)
	{
		if (requiredPreLayerIds.find(it->first) == requiredPreLayerIds.end())
		{
			DestroyRing(it->second);
			it = preLayerHistoryByLayerId.erase(it);
			mNeedsReset = true;
		}
		else
		{
			++it;
		}
	}

	if (mNeedsReset)
		ResetState();

	return true;
}

bool TemporalHistoryBuffers::CreateRing(Ring& ring, unsigned effectiveLength, TemporalHistorySource historySource, unsigned frameWidth, unsigned frameHeight, std::string& error)
{
	DestroyRing(ring);
	ring.effectiveLength = effectiveLength;
	ring.historySource = historySource;
	if (effectiveLength == 0)
		return true;

	ring.slots.resize(effectiveLength);
	for (Slot& slot : ring.slots)
	{
		glGenTextures(1, &slot.texture);
		glBindTexture(GL_TEXTURE_2D, slot.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);
		glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA16F, frameWidth, frameHeight, 0, GL_RGBA, GL_FLOAT, NULL);

		glGenFramebuffers(1, &slot.framebuffer);
		glBindFramebuffer(GL_FRAMEBUFFER, slot.framebuffer);
		glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, slot.texture, 0);
		const GLenum framebufferStatus = glCheckFramebufferStatus(GL_FRAMEBUFFER);
		if (framebufferStatus != GL_FRAMEBUFFER_COMPLETE)
		{
			error = "Failed to initialize a temporal history framebuffer.";
			glBindFramebuffer(GL_FRAMEBUFFER, 0);
			glBindTexture(GL_TEXTURE_2D, 0);
			DestroyRing(ring);
			return false;
		}
	}

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

void TemporalHistoryBuffers::DestroyRing(Ring& ring)
{
	for (Slot& slot : ring.slots)
	{
		if (slot.framebuffer != 0)
			glDeleteFramebuffers(1, &slot.framebuffer);
		if (slot.texture != 0)
			glDeleteTextures(1, &slot.texture);
		slot.framebuffer = 0;
		slot.texture = 0;
	}
	ring.slots.clear();
	ring.nextWriteIndex = 0;
	ring.filledCount = 0;
	ring.effectiveLength = 0;
	ring.historySource = TemporalHistorySource::None;
}

void TemporalHistoryBuffers::DestroyResources()
{
	DestroyRing(sourceHistoryRing);
	for (auto& historyEntry : preLayerHistoryByLayerId)
		DestroyRing(historyEntry.second);
	preLayerHistoryByLayerId.clear();
	mNeedsReset = true;
}

void TemporalHistoryBuffers::ResetState()
{
	sourceHistoryRing.nextWriteIndex = 0;
	sourceHistoryRing.filledCount = 0;
	for (auto& historyEntry : preLayerHistoryByLayerId)
	{
		historyEntry.second.nextWriteIndex = 0;
		historyEntry.second.filledCount = 0;
	}
	mNeedsReset = false;
}

void TemporalHistoryBuffers::PushFramebuffer(GLuint sourceFramebuffer, Ring& ring, unsigned frameWidth, unsigned frameHeight)
{
	if (ring.effectiveLength == 0 || ring.slots.empty())
		return;

	Slot& targetSlot = ring.slots[ring.nextWriteIndex];
	glBindFramebuffer(GL_READ_FRAMEBUFFER, sourceFramebuffer);
	glBindFramebuffer(GL_DRAW_FRAMEBUFFER, targetSlot.framebuffer);
	glBlitFramebuffer(0, 0, frameWidth, frameHeight, 0, 0, frameWidth, frameHeight, GL_COLOR_BUFFER_BIT, GL_LINEAR);
	ring.nextWriteIndex = (ring.nextWriteIndex + 1) % ring.slots.size();
	ring.filledCount = std::min<std::size_t>(ring.filledCount + 1, ring.slots.size());
}

void TemporalHistoryBuffers::PushSourceFramebuffer(GLuint sourceFramebuffer, unsigned frameWidth, unsigned frameHeight)
{
	PushFramebuffer(sourceFramebuffer, sourceHistoryRing, frameWidth, frameHeight);
}

void TemporalHistoryBuffers::PushPreLayerFramebuffer(const std::string& layerId, GLuint sourceFramebuffer, unsigned frameWidth, unsigned frameHeight)
{
	auto historyIt = preLayerHistoryByLayerId.find(layerId);
	if (historyIt != preLayerHistoryByLayerId.end())
		PushFramebuffer(sourceFramebuffer, historyIt->second, frameWidth, frameHeight);
}

void TemporalHistoryBuffers::BindSamplers(const RuntimeRenderState& state, GLuint currentSourceTexture, unsigned historyCap)
{
	for (unsigned index = 0; index < historyCap; ++index)
	{
		glActiveTexture(GL_TEXTURE0 + kSourceHistoryTextureUnitBase + index);
		glBindTexture(GL_TEXTURE_2D, ResolveTexture(sourceHistoryRing, currentSourceTexture, index));
	}

	const GLuint temporalBase = kSourceHistoryTextureUnitBase + historyCap;
	const Ring* temporalRing = nullptr;
	auto it = preLayerHistoryByLayerId.find(state.layerId);
	if (it != preLayerHistoryByLayerId.end())
		temporalRing = &it->second;

	for (unsigned index = 0; index < historyCap; ++index)
	{
		glActiveTexture(GL_TEXTURE0 + temporalBase + index);
		glBindTexture(GL_TEXTURE_2D, temporalRing ? ResolveTexture(*temporalRing, currentSourceTexture, index) : currentSourceTexture);
	}
	glActiveTexture(GL_TEXTURE0);
}

std::vector<GLuint> TemporalHistoryBuffers::ResolveSourceHistoryTextures(GLuint fallbackTexture, unsigned historyCap) const
{
	std::vector<GLuint> textures;
	textures.reserve(historyCap);
	for (unsigned index = 0; index < historyCap; ++index)
		textures.push_back(ResolveTexture(sourceHistoryRing, fallbackTexture, index));
	return textures;
}

std::vector<GLuint> TemporalHistoryBuffers::ResolveTemporalHistoryTextures(const RuntimeRenderState& state, GLuint fallbackTexture, unsigned historyCap) const
{
	const Ring* temporalRing = nullptr;
	auto it = preLayerHistoryByLayerId.find(state.layerId);
	if (it != preLayerHistoryByLayerId.end())
		temporalRing = &it->second;

	std::vector<GLuint> textures;
	textures.reserve(historyCap);
	for (unsigned index = 0; index < historyCap; ++index)
		textures.push_back(temporalRing ? ResolveTexture(*temporalRing, fallbackTexture, index) : fallbackTexture);
	return textures;
}

GLuint TemporalHistoryBuffers::ResolveTexture(const Ring& ring, GLuint fallbackTexture, std::size_t framesAgo) const
{
	if (ring.filledCount == 0 || ring.slots.empty())
		return fallbackTexture;

	const std::size_t clampedOffset = std::min<std::size_t>(framesAgo, ring.filledCount - 1);
	const std::size_t newestIndex = (ring.nextWriteIndex + ring.slots.size() - 1) % ring.slots.size();
	const std::size_t slotIndex = (newestIndex + ring.slots.size() - clampedOffset) % ring.slots.size();
	return ring.slots[slotIndex].texture != 0 ? ring.slots[slotIndex].texture : fallbackTexture;
}

unsigned TemporalHistoryBuffers::SourceAvailableCount() const
{
	return static_cast<unsigned>(sourceHistoryRing.filledCount);
}

unsigned TemporalHistoryBuffers::AvailableCountForLayer(const std::string& layerId) const
{
	auto it = preLayerHistoryByLayerId.find(layerId);
	if (it == preLayerHistoryByLayerId.end())
		return 0;
	return static_cast<unsigned>(it->second.filledCount);
}