#include "InputFrameTexture.h"

#include <chrono>

#ifndef GL_FRAMEBUFFER_BINDING
#define GL_FRAMEBUFFER_BINDING 0x8CA6
#endif

namespace
{
constexpr GLuint kUyvyTextureUnit = 0;

const char* kDecodeVertexShader = R"GLSL(
#version 430 core
out vec2 vTexCoord;
void main()
{
	vec2 positions[3] = vec2[3](
		vec2(-1.0, -1.0),
		vec2( 3.0, -1.0),
		vec2(-1.0,  3.0));
	vec2 texCoords[3] = vec2[3](
		vec2(0.0, 0.0),
		vec2(2.0, 0.0),
		vec2(0.0, 2.0));
	gl_Position = vec4(positions[gl_VertexID], 0.0, 1.0);
	vTexCoord = texCoords[gl_VertexID];
}
)GLSL";

const char* kUyvyDecodeFragmentShader = R"GLSL(
#version 430 core
layout(binding = 0) uniform sampler2D uPackedUyvy;
uniform vec2 uDecodedSize;
in vec2 vTexCoord;
out vec4 fragColor;

vec4 rec709YCbCr2rgba(float yByte, float cbByte, float crByte)
{
	float y = (yByte - 16.0) / 219.0;
	float cb = (cbByte - 16.0) / 224.0 - 0.5;
	float cr = (crByte - 16.0) / 224.0 - 0.5;
	return vec4(
		y + 1.5748 * cr,
		y - 0.1873 * cb - 0.4681 * cr,
		y + 1.8556 * cb,
		1.0);
}

void main()
{
	ivec2 decodedSize = ivec2(uDecodedSize);
	ivec2 outputCoord = ivec2(clamp(gl_FragCoord.xy, vec2(0.0), vec2(decodedSize - ivec2(1))));
	int sourceY = decodedSize.y - 1 - outputCoord.y;
	ivec2 packedCoord = ivec2(clamp(outputCoord.x / 2, 0, max(decodedSize.x / 2 - 1, 0)), sourceY);
	vec4 macroPixel = texelFetch(uPackedUyvy, packedCoord, 0) * 255.0;
	float ySample = (outputCoord.x & 1) != 0 ? macroPixel.a : macroPixel.g;
	fragColor = clamp(rec709YCbCr2rgba(ySample, macroPixel.r, macroPixel.b), vec4(0.0), vec4(1.0));
}
)GLSL";
}

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

GLuint InputFrameTexture::PollAndUpload(InputFrameMailbox* mailbox)
{
	if (mailbox == nullptr)
		return mTexture;

	InputFrame frame;
	if (!mailbox->TryAcquireOldest(frame))
	{
		++mUploadMisses;
		mLastUploadMilliseconds = 0.0;
		return mTexture;
	}

	if (frame.bytes != nullptr && frame.pixelFormat == VideoIOPixelFormat::Bgra8 && EnsureTexture(frame))
	{
		mLastFrameFormatSupported = true;
		const auto uploadStart = std::chrono::steady_clock::now();
		UploadBgra8FrameFlippedVertically(frame);
		const auto uploadEnd = std::chrono::steady_clock::now();
		mLastUploadMilliseconds = std::chrono::duration_cast<std::chrono::duration<double, std::milli>>(uploadEnd - uploadStart).count();
		++mUploadedFrames;
	}
	else if (frame.bytes != nullptr && frame.pixelFormat == VideoIOPixelFormat::Uyvy8 && EnsureTexture(frame) && EnsureRawUyvyTexture(frame) && EnsureDecodeProgram())
	{
		mLastFrameFormatSupported = true;
		const auto uploadStart = std::chrono::steady_clock::now();
		UploadUyvy8Frame(frame);
		DecodeUyvy8Frame(frame);
		const auto uploadEnd = std::chrono::steady_clock::now();
		mLastUploadMilliseconds = std::chrono::duration_cast<std::chrono::duration<double, std::milli>>(uploadEnd - uploadStart).count();
		++mUploadedFrames;
	}
	else
	{
		mLastFrameFormatSupported = frame.pixelFormat == VideoIOPixelFormat::Bgra8 || frame.pixelFormat == VideoIOPixelFormat::Uyvy8;
		mLastUploadMilliseconds = 0.0;
	}

	mailbox->Release(frame);
	return mTexture;
}

void InputFrameTexture::ShutdownGl()
{
	if (mTexture != 0)
		glDeleteTextures(1, &mTexture);
	if (mRawTexture != 0)
		glDeleteTextures(1, &mRawTexture);
	mTexture = 0;
	mRawTexture = 0;
	mWidth = 0;
	mHeight = 0;
	mRawWidth = 0;
	mRawHeight = 0;
	DestroyDecodeResources();
}

bool InputFrameTexture::EnsureTexture(const InputFrame& frame)
{
	if (frame.width == 0 || frame.height == 0)
		return false;

	if (mTexture != 0 && mWidth == frame.width && mHeight == frame.height)
		return true;

	ShutdownGl();
	glGenTextures(1, &mTexture);
	glBindTexture(GL_TEXTURE_2D, mTexture);
	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>(frame.width),
		static_cast<GLsizei>(frame.height),
		0,
		GL_BGRA,
		GL_UNSIGNED_INT_8_8_8_8_REV,
		nullptr);
	glBindTexture(GL_TEXTURE_2D, 0);

	mWidth = frame.width;
	mHeight = frame.height;
	return mTexture != 0;
}

bool InputFrameTexture::EnsureRawUyvyTexture(const InputFrame& frame)
{
	if (frame.width == 0 || frame.height == 0)
		return false;

	const unsigned rawWidth = (frame.width + 1u) / 2u;
	if (mRawTexture != 0 && mRawWidth == rawWidth && mRawHeight == frame.height)
		return true;

	if (mRawTexture != 0)
		glDeleteTextures(1, &mRawTexture);
	mRawTexture = 0;
	glGenTextures(1, &mRawTexture);
	glBindTexture(GL_TEXTURE_2D, mRawTexture);
	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);
	glTexImage2D(
		GL_TEXTURE_2D,
		0,
		GL_RGBA8,
		static_cast<GLsizei>(rawWidth),
		static_cast<GLsizei>(frame.height),
		0,
		GL_RGBA,
		GL_UNSIGNED_BYTE,
		nullptr);
	glBindTexture(GL_TEXTURE_2D, 0);

	mRawWidth = rawWidth;
	mRawHeight = frame.height;
	return mRawTexture != 0;
}

void InputFrameTexture::UploadBgra8FrameFlippedVertically(const InputFrame& frame)
{
	glBindTexture(GL_TEXTURE_2D, mTexture);
	glPixelStorei(GL_UNPACK_ALIGNMENT, 4);
	glPixelStorei(GL_UNPACK_ROW_LENGTH, frame.rowBytes > 0 ? static_cast<GLint>(frame.rowBytes / 4) : 0);

	const unsigned char* sourceBytes = static_cast<const unsigned char*>(frame.bytes);
	for (unsigned destinationY = 0; destinationY < frame.height; ++destinationY)
	{
		const unsigned sourceY = frame.height - 1u - destinationY;
		const unsigned char* sourceRow = sourceBytes + static_cast<std::size_t>(sourceY) * static_cast<std::size_t>(frame.rowBytes);
		glTexSubImage2D(
			GL_TEXTURE_2D,
			0,
			0,
			static_cast<GLint>(destinationY),
			static_cast<GLsizei>(frame.width),
			1,
			GL_BGRA,
			GL_UNSIGNED_INT_8_8_8_8_REV,
			sourceRow);
	}

	glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
	glBindTexture(GL_TEXTURE_2D, 0);
}

void InputFrameTexture::UploadUyvy8Frame(const InputFrame& frame)
{
	glBindTexture(GL_TEXTURE_2D, mRawTexture);
	glPixelStorei(GL_UNPACK_ALIGNMENT, 4);
	glPixelStorei(GL_UNPACK_ROW_LENGTH, frame.rowBytes > 0 ? static_cast<GLint>(frame.rowBytes / 4) : 0);
	glTexSubImage2D(
		GL_TEXTURE_2D,
		0,
		0,
		0,
		static_cast<GLsizei>((frame.width + 1u) / 2u),
		static_cast<GLsizei>(frame.height),
		GL_RGBA,
		GL_UNSIGNED_BYTE,
		frame.bytes);
	glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
	glBindTexture(GL_TEXTURE_2D, 0);
}

void InputFrameTexture::DecodeUyvy8Frame(const InputFrame& frame)
{
	GLint previousFramebuffer = 0;
	glGetIntegerv(GL_FRAMEBUFFER_BINDING, &previousFramebuffer);

	if (mDecodeFramebuffer == 0)
		glGenFramebuffers(1, &mDecodeFramebuffer);
	glBindFramebuffer(GL_FRAMEBUFFER, mDecodeFramebuffer);
	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, mTexture, 0);
	if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
	{
		glBindFramebuffer(GL_FRAMEBUFFER, static_cast<GLuint>(previousFramebuffer));
		return;
	}

	glViewport(0, 0, static_cast<GLsizei>(frame.width), static_cast<GLsizei>(frame.height));
	glDisable(GL_SCISSOR_TEST);
	glDisable(GL_DEPTH_TEST);
	glDisable(GL_BLEND);
	glUseProgram(mDecodeProgram);
	const GLint decodedSizeLocation = glGetUniformLocation(mDecodeProgram, "uDecodedSize");
	if (decodedSizeLocation >= 0)
		glUniform2f(decodedSizeLocation, static_cast<GLfloat>(frame.width), static_cast<GLfloat>(frame.height));
	glActiveTexture(GL_TEXTURE0 + kUyvyTextureUnit);
	glBindTexture(GL_TEXTURE_2D, mRawTexture);
	glBindVertexArray(mDecodeVertexArray);
	glDrawArrays(GL_TRIANGLES, 0, 3);
	glBindVertexArray(0);
	glBindTexture(GL_TEXTURE_2D, 0);
	glActiveTexture(GL_TEXTURE0);
	glUseProgram(0);
	glBindFramebuffer(GL_FRAMEBUFFER, static_cast<GLuint>(previousFramebuffer));
}

bool InputFrameTexture::EnsureDecodeProgram()
{
	if (mDecodeProgram != 0)
		return true;

	if (!CompileShader(GL_VERTEX_SHADER, kDecodeVertexShader, mDecodeVertexShader))
		return false;
	if (!CompileShader(GL_FRAGMENT_SHADER, kUyvyDecodeFragmentShader, mDecodeFragmentShader))
		return false;
	if (!LinkProgram(mDecodeVertexShader, mDecodeFragmentShader, mDecodeProgram))
		return false;

	glUseProgram(mDecodeProgram);
	const GLint samplerLocation = glGetUniformLocation(mDecodeProgram, "uPackedUyvy");
	if (samplerLocation >= 0)
		glUniform1i(samplerLocation, static_cast<GLint>(kUyvyTextureUnit));
	glUseProgram(0);

	if (mDecodeVertexArray == 0)
		glGenVertexArrays(1, &mDecodeVertexArray);
	return mDecodeProgram != 0 && mDecodeVertexArray != 0;
}

void InputFrameTexture::DestroyDecodeResources()
{
	if (mDecodeFramebuffer != 0)
		glDeleteFramebuffers(1, &mDecodeFramebuffer);
	if (mDecodeVertexArray != 0)
		glDeleteVertexArrays(1, &mDecodeVertexArray);
	if (mDecodeProgram != 0)
		glDeleteProgram(mDecodeProgram);
	if (mDecodeVertexShader != 0)
		glDeleteShader(mDecodeVertexShader);
	if (mDecodeFragmentShader != 0)
		glDeleteShader(mDecodeFragmentShader);
	mDecodeFramebuffer = 0;
	mDecodeVertexArray = 0;
	mDecodeProgram = 0;
	mDecodeVertexShader = 0;
	mDecodeFragmentShader = 0;
}

bool InputFrameTexture::CompileShader(GLenum shaderType, const char* source, GLuint& shader)
{
	shader = glCreateShader(shaderType);
	glShaderSource(shader, 1, &source, nullptr);
	glCompileShader(shader);
	GLint compileResult = GL_FALSE;
	glGetShaderiv(shader, GL_COMPILE_STATUS, &compileResult);
	if (compileResult != GL_FALSE)
		return true;
	glDeleteShader(shader);
	shader = 0;
	return false;
}

bool InputFrameTexture::LinkProgram(GLuint vertexShader, GLuint fragmentShader, GLuint& program)
{
	program = glCreateProgram();
	glAttachShader(program, vertexShader);
	glAttachShader(program, fragmentShader);
	glLinkProgram(program);
	GLint linkResult = GL_FALSE;
	glGetProgramiv(program, GL_LINK_STATUS, &linkResult);
	if (linkResult != GL_FALSE)
		return true;
	glDeleteProgram(program);
	program = 0;
	return false;
}