#include "GlShaderSources.h"

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

const char* kDecodeFragmentShaderSource =
	"#version 430 core\n"
	"layout(binding = 2) uniform sampler2D uPackedVideoInput;\n"
	"uniform vec2 uPackedVideoResolution;\n"
	"uniform vec2 uDecodedVideoResolution;\n"
	"uniform int uInputPixelFormat;\n"
	"in vec2 vTexCoord;\n"
	"layout(location = 0) out vec4 fragColor;\n"
	"vec4 rec709YCbCr2rgba(float Y, float Cb, float Cr, float a)\n"
	"{\n"
	"	Y = (Y * 256.0 - 16.0) / 219.0;\n"
	"	Cb = (Cb * 256.0 - 16.0) / 224.0 - 0.5;\n"
	"	Cr = (Cr * 256.0 - 16.0) / 224.0 - 0.5;\n"
	"	return vec4(Y + 1.5748 * Cr, Y - 0.1873 * Cb - 0.4681 * Cr, Y + 1.8556 * Cb, a);\n"
	"}\n"
	"vec4 rec709YCbCr10_2rgba(float Y, float Cb, float Cr, float a)\n"
	"{\n"
	"	Y = (Y - 64.0) / 876.0;\n"
	"	Cb = (Cb - 64.0) / 896.0 - 0.5;\n"
	"	Cr = (Cr - 64.0) / 896.0 - 0.5;\n"
	"	return vec4(Y + 1.5748 * Cr, Y - 0.1873 * Cb - 0.4681 * Cr, Y + 1.8556 * Cb, a);\n"
	"}\n"
	"uint loadV210Word(ivec2 coord)\n"
	"{\n"
	"	vec4 b = round(texelFetch(uPackedVideoInput, coord, 0) * 255.0);\n"
	"	return uint(b.r) | (uint(b.g) << 8) | (uint(b.b) << 16) | (uint(b.a) << 24);\n"
	"}\n"
	"float v210Component(uint word, int index)\n"
	"{\n"
	"	return float((word >> uint(index * 10)) & 1023u);\n"
	"}\n"
	"vec4 decodeUyvy8(ivec2 outputCoord, ivec2 packedSize)\n"
	"{\n"
	"	ivec2 packedCoord = ivec2(clamp(outputCoord.x / 2, 0, packedSize.x - 1), clamp(outputCoord.y, 0, packedSize.y - 1));\n"
	"	vec4 macroPixel = texelFetch(uPackedVideoInput, packedCoord, 0);\n"
	"	float ySample = (outputCoord.x & 1) != 0 ? macroPixel.a : macroPixel.g;\n"
	"	return rec709YCbCr2rgba(ySample, macroPixel.b, macroPixel.r, 1.0);\n"
	"}\n"
	"vec4 decodeV210(ivec2 outputCoord, ivec2 packedSize)\n"
	"{\n"
	"	int group = outputCoord.x / 6;\n"
	"	int pixel = outputCoord.x - group * 6;\n"
	"	int wordBase = group * 4;\n"
	"	ivec2 rowBase = ivec2(wordBase, clamp(outputCoord.y, 0, packedSize.y - 1));\n"
	"	uint w0 = loadV210Word(ivec2(min(rowBase.x + 0, packedSize.x - 1), rowBase.y));\n"
	"	uint w1 = loadV210Word(ivec2(min(rowBase.x + 1, packedSize.x - 1), rowBase.y));\n"
	"	uint w2 = loadV210Word(ivec2(min(rowBase.x + 2, packedSize.x - 1), rowBase.y));\n"
	"	uint w3 = loadV210Word(ivec2(min(rowBase.x + 3, packedSize.x - 1), rowBase.y));\n"
	"	float y0 = v210Component(w0, 1);\n"
	"	float y1 = v210Component(w1, 0);\n"
	"	float y2 = v210Component(w1, 2);\n"
	"	float y3 = v210Component(w2, 1);\n"
	"	float y4 = v210Component(w3, 0);\n"
	"	float y5 = v210Component(w3, 2);\n"
	"	float cb0 = v210Component(w0, 0);\n"
	"	float cr0 = v210Component(w0, 2);\n"
	"	float cb2 = v210Component(w1, 1);\n"
	"	float cr2 = v210Component(w2, 0);\n"
	"	float cb4 = v210Component(w2, 2);\n"
	"	float cr4 = v210Component(w3, 1);\n"
	"	float ySample = pixel == 0 ? y0 : pixel == 1 ? y1 : pixel == 2 ? y2 : pixel == 3 ? y3 : pixel == 4 ? y4 : y5;\n"
	"	float cbSample = pixel < 2 ? cb0 : pixel < 4 ? cb2 : cb4;\n"
	"	float crSample = pixel < 2 ? cr0 : pixel < 4 ? cr2 : cr4;\n"
	"	return rec709YCbCr10_2rgba(ySample, cbSample, crSample, 1.0);\n"
	"}\n"
	"void main()\n"
	"{\n"
	"	vec2 correctedUv = vec2(vTexCoord.x, 1.0 - vTexCoord.y);\n"
	"	ivec2 decodedSize = ivec2(max(uDecodedVideoResolution, vec2(1.0, 1.0)));\n"
	"	ivec2 outputCoord = clamp(ivec2(correctedUv * vec2(decodedSize)), ivec2(0, 0), decodedSize - ivec2(1, 1));\n"
	"	ivec2 packedSize = ivec2(max(uPackedVideoResolution, vec2(1.0, 1.0)));\n"
	"	fragColor = uInputPixelFormat == 1 ? decodeV210(outputCoord, packedSize) : decodeUyvy8(outputCoord, packedSize);\n"
	"}\n";

const char* kOutputPackFragmentShaderSource =
	"#version 430 core\n"
	"layout(binding = 0) uniform sampler2D uOutputRgb;\n"
	"uniform vec2 uOutputVideoResolution;\n"
	"uniform float uActiveV210Words;\n"
	"uniform int uOutputPackFormat;\n"
	"in vec2 vTexCoord;\n"
	"layout(location = 0) out vec4 fragColor;\n"
	"vec4 rgbaAt(int x, int y)\n"
	"{\n"
	"	ivec2 size = ivec2(max(uOutputVideoResolution, vec2(1.0, 1.0)));\n"
	"	return clamp(texelFetch(uOutputRgb, ivec2(clamp(x, 0, size.x - 1), clamp(y, 0, size.y - 1)), 0), vec4(0.0), vec4(1.0));\n"
	"}\n"
	"vec3 rgbAt(int x, int y)\n"
	"{\n"
	"	return rgbaAt(x, y).rgb;\n"
	"}\n"
	"vec3 rgbToLegalYcbcr10(vec3 rgb)\n"
	"{\n"
	"	float y = dot(rgb, vec3(0.2126, 0.7152, 0.0722));\n"
	"	float cb = (rgb.b - y) / 1.8556 + 0.5;\n"
	"	float cr = (rgb.r - y) / 1.5748 + 0.5;\n"
	"	return vec3(clamp(round(64.0 + y * 876.0), 64.0, 940.0), clamp(round(64.0 + cb * 896.0), 64.0, 960.0), clamp(round(64.0 + cr * 896.0), 64.0, 960.0));\n"
	"}\n"
	"uint makeWord(float a, float b, float c)\n"
	"{\n"
	"	return (uint(a) & 1023u) | ((uint(b) & 1023u) << 10) | ((uint(c) & 1023u) << 20);\n"
	"}\n"
	"vec4 wordToBytes(uint word)\n"
	"{\n"
	"	return vec4(float(word & 255u), float((word >> 8) & 255u), float((word >> 16) & 255u), float((word >> 24) & 255u)) / 255.0;\n"
	"}\n"
	"vec4 bigEndianWordToBytes(uint word)\n"
	"{\n"
	"	return vec4(float((word >> 24) & 255u), float((word >> 16) & 255u), float((word >> 8) & 255u), float(word & 255u)) / 255.0;\n"
	"}\n"
	"vec4 packAy10Word(ivec2 outCoord)\n"
	"{\n"
	"	ivec2 size = ivec2(max(uOutputVideoResolution, vec2(1.0, 1.0)));\n"
	"	if (outCoord.x >= size.x)\n"
	"		return vec4(0.0);\n"
	"	int pixelBase = (outCoord.x / 2) * 2;\n"
	"	int y = outCoord.y;\n"
	"	vec4 rgba0 = rgbaAt(pixelBase + 0, y);\n"
	"	vec4 rgba1 = rgbaAt(pixelBase + 1, y);\n"
	"	vec3 c0 = rgbToLegalYcbcr10(rgba0.rgb);\n"
	"	vec3 c1 = rgbToLegalYcbcr10(rgba1.rgb);\n"
	"	float chroma = (outCoord.x & 1) == 0 ? round((c0.y + c1.y) * 0.5) : round((c0.z + c1.z) * 0.5);\n"
	"	float alpha = round(clamp(((outCoord.x & 1) == 0 ? rgba0.a : rgba1.a), 0.0, 1.0) * 1023.0);\n"
	"	float luma = (outCoord.x & 1) == 0 ? c0.x : c1.x;\n"
	"	uint word = ((uint(luma) & 1023u) << 22) | ((uint(chroma) & 1023u) << 12) | ((uint(alpha) & 1023u) << 2);\n"
	"	return bigEndianWordToBytes(word);\n"
	"}\n"
	"void main()\n"
	"{\n"
	"	ivec2 outCoord = ivec2(gl_FragCoord.xy);\n"
	"	if (uOutputPackFormat == 2)\n"
	"	{\n"
	"		fragColor = packAy10Word(outCoord);\n"
	"		return;\n"
	"	}\n"
	"	if (float(outCoord.x) >= uActiveV210Words)\n"
	"	{\n"
	"		fragColor = vec4(0.0);\n"
	"		return;\n"
	"	}\n"
	"	int group = outCoord.x / 4;\n"
	"	int wordIndex = outCoord.x - group * 4;\n"
	"	int pixelBase = group * 6;\n"
	"	int y = outCoord.y;\n"
	"	vec3 c0 = rgbToLegalYcbcr10(rgbAt(pixelBase + 0, y));\n"
	"	vec3 c1 = rgbToLegalYcbcr10(rgbAt(pixelBase + 1, y));\n"
	"	vec3 c2 = rgbToLegalYcbcr10(rgbAt(pixelBase + 2, y));\n"
	"	vec3 c3 = rgbToLegalYcbcr10(rgbAt(pixelBase + 3, y));\n"
	"	vec3 c4 = rgbToLegalYcbcr10(rgbAt(pixelBase + 4, y));\n"
	"	vec3 c5 = rgbToLegalYcbcr10(rgbAt(pixelBase + 5, y));\n"
	"	float cb0 = round((c0.y + c1.y) * 0.5);\n"
	"	float cr0 = round((c0.z + c1.z) * 0.5);\n"
	"	float cb2 = round((c2.y + c3.y) * 0.5);\n"
	"	float cr2 = round((c2.z + c3.z) * 0.5);\n"
	"	float cb4 = round((c4.y + c5.y) * 0.5);\n"
	"	float cr4 = round((c4.z + c5.z) * 0.5);\n"
	"	uint word = wordIndex == 0 ? makeWord(cb0, c0.x, cr0) : wordIndex == 1 ? makeWord(c1.x, cb2, c2.x) : wordIndex == 2 ? makeWord(cr2, c3.x, cb4) : makeWord(c4.x, cr4, c5.x);\n"
	"	fragColor = wordToBytes(word);\n"
	"}\n";
