#include "ShaderFeedbackBuffers.h"

#include <set>

namespace
{
void ConfigureFeedbackTexture(unsigned frameWidth, unsigned frameHeight)
{
	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);
}
}

bool ShaderFeedbackBuffers::EnsureResources(const std::vector<RuntimeRenderState>& layerStates, unsigned frameWidth, unsigned frameHeight, std::string& error)
{
	if (!EnsureZeroTexture())
	{
		error = "Failed to initialize shader feedback fallback texture.";
		return false;
	}

	std::set<std::string> requiredLayerIds;
	for (const RuntimeRenderState& state : layerStates)
	{
		if (!state.feedback.enabled)
			continue;

		requiredLayerIds.insert(state.layerId);
		auto surfaceIt = mSurfacesByLayerId.find(state.layerId);
		if (surfaceIt == mSurfacesByLayerId.end() ||
			surfaceIt->second.width != frameWidth ||
			surfaceIt->second.height != frameHeight)
		{
			Surface replacement;
			if (!CreateSurface(replacement, frameWidth, frameHeight, error))
				return false;
			mSurfacesByLayerId[state.layerId] = std::move(replacement);
		}
	}

	for (auto it = mSurfacesByLayerId.begin(); it != mSurfacesByLayerId.end();)
	{
		if (requiredLayerIds.find(it->first) == requiredLayerIds.end())
		{
			DestroySurface(it->second);
			it = mSurfacesByLayerId.erase(it);
		}
		else
		{
			++it;
		}
	}

	return true;
}

void ShaderFeedbackBuffers::DestroyResources()
{
	for (auto& entry : mSurfacesByLayerId)
		DestroySurface(entry.second);
	mSurfacesByLayerId.clear();

	if (mZeroTexture != 0)
	{
		glDeleteTextures(1, &mZeroTexture);
		mZeroTexture = 0;
	}
}

void ShaderFeedbackBuffers::ResetState()
{
	for (auto& entry : mSurfacesByLayerId)
		ClearSurfaceState(entry.second);
}

GLuint ShaderFeedbackBuffers::ResolveReadTexture(const RuntimeRenderState& state) const
{
	if (!state.feedback.enabled)
		return mZeroTexture;

	auto surfaceIt = mSurfacesByLayerId.find(state.layerId);
	if (surfaceIt == mSurfacesByLayerId.end() || !surfaceIt->second.hasData)
		return mZeroTexture;

	return surfaceIt->second.slots[surfaceIt->second.readIndex].texture != 0
		? surfaceIt->second.slots[surfaceIt->second.readIndex].texture
		: mZeroTexture;
}

bool ShaderFeedbackBuffers::FeedbackAvailable(const RuntimeRenderState& state) const
{
	if (!state.feedback.enabled)
		return false;

	auto surfaceIt = mSurfacesByLayerId.find(state.layerId);
	return surfaceIt != mSurfacesByLayerId.end() && surfaceIt->second.hasData;
}

void ShaderFeedbackBuffers::CaptureFeedbackFramebuffer(const std::string& layerId, GLuint sourceFramebuffer, unsigned frameWidth, unsigned frameHeight)
{
	auto surfaceIt = mSurfacesByLayerId.find(layerId);
	if (surfaceIt == mSurfacesByLayerId.end())
		return;

	Surface& surface = surfaceIt->second;
	const unsigned writeIndex = 1u - surface.readIndex;
	glBindFramebuffer(GL_READ_FRAMEBUFFER, sourceFramebuffer);
	glBindFramebuffer(GL_DRAW_FRAMEBUFFER, surface.slots[writeIndex].framebuffer);
	glBlitFramebuffer(0, 0, frameWidth, frameHeight, 0, 0, frameWidth, frameHeight, GL_COLOR_BUFFER_BIT, GL_LINEAR);
	surface.pendingWrite = true;
}

void ShaderFeedbackBuffers::FinalizeFrame()
{
	for (auto& entry : mSurfacesByLayerId)
	{
		Surface& surface = entry.second;
		if (!surface.pendingWrite)
			continue;

		surface.readIndex = 1u - surface.readIndex;
		surface.hasData = true;
		surface.pendingWrite = false;
	}
}

bool ShaderFeedbackBuffers::EnsureZeroTexture()
{
	if (mZeroTexture != 0)
		return true;

	glGenTextures(1, &mZeroTexture);
	glBindTexture(GL_TEXTURE_2D, mZeroTexture);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
	const float zeroPixel[4] = { 0.0f, 0.0f, 0.0f, 0.0f };
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA16F, 1, 1, 0, GL_RGBA, GL_FLOAT, zeroPixel);
	glBindTexture(GL_TEXTURE_2D, 0);
	return mZeroTexture != 0;
}

bool ShaderFeedbackBuffers::CreateSurface(Surface& surface, unsigned frameWidth, unsigned frameHeight, std::string& error)
{
	DestroySurface(surface);

	surface.width = frameWidth;
	surface.height = frameHeight;
	for (Slot& slot : surface.slots)
	{
		glGenTextures(1, &slot.texture);
		glBindTexture(GL_TEXTURE_2D, slot.texture);
		ConfigureFeedbackTexture(frameWidth, frameHeight);

		glGenFramebuffers(1, &slot.framebuffer);
		glBindFramebuffer(GL_FRAMEBUFFER, slot.framebuffer);
		glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, slot.texture, 0);
		if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
		{
			error = "Failed to initialize a shader feedback framebuffer.";
			glBindFramebuffer(GL_FRAMEBUFFER, 0);
			glBindTexture(GL_TEXTURE_2D, 0);
			DestroySurface(surface);
			return false;
		}
	}

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

void ShaderFeedbackBuffers::DestroySurface(Surface& surface)
{
	for (Slot& slot : surface.slots)
	{
		if (slot.framebuffer != 0)
			glDeleteFramebuffers(1, &slot.framebuffer);
		if (slot.texture != 0)
			glDeleteTextures(1, &slot.texture);
		slot.framebuffer = 0;
		slot.texture = 0;
	}

	surface.width = 0;
	surface.height = 0;
	surface.readIndex = 0;
	surface.hasData = false;
	surface.pendingWrite = false;
}

void ShaderFeedbackBuffers::ClearSurfaceState(Surface& surface)
{
	surface.readIndex = 0;
	surface.hasData = false;
	surface.pendingWrite = false;
}