#include "stdafx.h"
#include "ShaderCompiler.h"

#include <cstring>
#include <fstream>
#include <regex>
#include <sstream>
#include <vector>

namespace
{
std::string ReplaceAll(std::string text, const std::string& from, const std::string& to)
{
	std::size_t offset = 0;
	while ((offset = text.find(from, offset)) != std::string::npos)
	{
		text.replace(offset, from.length(), to);
		offset += to.length();
	}
	return text;
}

std::string SlangTypeForParameter(ShaderParameterType type)
{
	switch (type)
	{
	case ShaderParameterType::Float: return "uniform float";
	case ShaderParameterType::Vec2: return "uniform float2";
	case ShaderParameterType::Color: return "uniform float4";
	case ShaderParameterType::Boolean: return "uniform bool";
	case ShaderParameterType::Enum: return "uniform int";
	}
	return "uniform float";
}
}

ShaderCompiler::ShaderCompiler(
	const std::filesystem::path& repoRoot,
	const std::filesystem::path& wrapperPath,
	const std::filesystem::path& generatedGlslPath,
	const std::filesystem::path& patchedGlslPath,
	unsigned maxTemporalHistoryFrames)
	: mRepoRoot(repoRoot),
	mWrapperPath(wrapperPath),
	mGeneratedGlslPath(generatedGlslPath),
	mPatchedGlslPath(patchedGlslPath),
	mMaxTemporalHistoryFrames(maxTemporalHistoryFrames)
{
}

bool ShaderCompiler::BuildLayerFragmentShaderSource(const ShaderPackage& shaderPackage, std::string& fragmentShaderSource, std::string& error) const
{
	const std::string wrapperSource = BuildWrapperSlangSource(shaderPackage);
	if (!WriteTextFile(mWrapperPath, wrapperSource, error))
		return false;

	if (!RunSlangCompiler(mWrapperPath, mGeneratedGlslPath, error))
		return false;

	fragmentShaderSource = ReadTextFile(mGeneratedGlslPath, error);
	if (fragmentShaderSource.empty())
		return false;

	if (!PatchGeneratedGlsl(fragmentShaderSource, error))
		return false;

	if (!WriteTextFile(mPatchedGlslPath, fragmentShaderSource, error))
		return false;

	return true;
}

std::string ShaderCompiler::BuildWrapperSlangSource(const ShaderPackage& shaderPackage) const
{
	std::ostringstream source;
	source << "struct FragmentInput\n";
	source << "{\n";
	source << "\tfloat4 position : SV_Position;\n";
	source << "\tfloat2 texCoord : TEXCOORD0;\n";
	source << "};\n\n";
	source << "struct ShaderContext\n";
	source << "{\n";
	source << "\tfloat2 uv;\n";
	source << "\tfloat4 sourceColor;\n";
	source << "\tfloat2 inputResolution;\n";
	source << "\tfloat2 outputResolution;\n";
	source << "\tfloat time;\n";
	source << "\tfloat frameCount;\n";
	source << "\tfloat mixAmount;\n";
	source << "\tfloat bypass;\n";
	source << "\tint sourceHistoryLength;\n";
	source << "\tint temporalHistoryLength;\n";
	source << "};\n\n";
	source << "cbuffer GlobalParams\n";
	source << "{\n";
	source << "\tfloat gTime;\n";
	source << "\tfloat2 gInputResolution;\n";
	source << "\tfloat2 gOutputResolution;\n";
	source << "\tfloat gFrameCount;\n";
	source << "\tfloat gMixAmount;\n";
	source << "\tfloat gBypass;\n";
	source << "\tint gSourceHistoryLength;\n";
	source << "\tint gTemporalHistoryLength;\n";
	for (const ShaderParameterDefinition& definition : shaderPackage.parameters)
		source << "\t" << SlangTypeForParameter(definition.type).substr(strlen("uniform ")) << " " << definition.id << ";\n";
	source << "};\n\n";
	source << "Sampler2D<float4> gVideoInput;\n";
	for (unsigned index = 0; index < mMaxTemporalHistoryFrames; ++index)
		source << "Sampler2D<float4> gSourceHistory" << index << ";\n";
	for (unsigned index = 0; index < mMaxTemporalHistoryFrames; ++index)
		source << "Sampler2D<float4> gTemporalHistory" << index << ";\n";
	for (const ShaderTextureAsset& textureAsset : shaderPackage.textureAssets)
		source << "Sampler2D<float4> " << textureAsset.id << ";\n";
	source << "\n";
	source << "float4 sampleVideo(float2 tc)\n";
	source << "{\n";
	source << "\treturn gVideoInput.Sample(tc);\n";
	source << "}\n\n";
	source << "float4 sampleSourceHistory(int framesAgo, float2 tc)\n";
	source << "{\n";
	source << "\tif (gSourceHistoryLength <= 0)\n";
	source << "\t\treturn sampleVideo(tc);\n";
	source << "\tint clampedIndex = framesAgo;\n";
	source << "\tif (clampedIndex < 0)\n";
	source << "\t\tclampedIndex = 0;\n";
	source << "\tif (clampedIndex >= gSourceHistoryLength)\n";
	source << "\t\tclampedIndex = gSourceHistoryLength - 1;\n";
	source << "\tswitch (clampedIndex)\n";
	source << "\t{\n";
	for (unsigned index = 0; index < mMaxTemporalHistoryFrames; ++index)
		source << "\tcase " << index << ": return gSourceHistory" << index << ".Sample(tc);\n";
	source << "\tdefault: return sampleVideo(tc);\n";
	source << "\t}\n";
	source << "}\n\n";
	source << "float4 sampleTemporalHistory(int framesAgo, float2 tc)\n";
	source << "{\n";
	source << "\tif (gTemporalHistoryLength <= 0)\n";
	source << "\t\treturn sampleVideo(tc);\n";
	source << "\tint clampedIndex = framesAgo;\n";
	source << "\tif (clampedIndex < 0)\n";
	source << "\t\tclampedIndex = 0;\n";
	source << "\tif (clampedIndex >= gTemporalHistoryLength)\n";
	source << "\t\tclampedIndex = gTemporalHistoryLength - 1;\n";
	source << "\tswitch (clampedIndex)\n";
	source << "\t{\n";
	for (unsigned index = 0; index < mMaxTemporalHistoryFrames; ++index)
		source << "\tcase " << index << ": return gTemporalHistory" << index << ".Sample(tc);\n";
	source << "\tdefault: return sampleVideo(tc);\n";
	source << "\t}\n";
	source << "}\n\n";
	source << "#include \"" << shaderPackage.shaderPath.generic_string() << "\"\n\n";
	source << "[shader(\"fragment\")]\n";
	source << "float4 fragmentMain(FragmentInput input) : SV_Target\n";
	source << "{\n";
	source << "\tShaderContext context;\n";
	source << "\tcontext.uv = input.texCoord;\n";
	source << "\tcontext.sourceColor = sampleVideo(context.uv);\n";
	source << "\tcontext.inputResolution = gInputResolution;\n";
	source << "\tcontext.outputResolution = gOutputResolution;\n";
	source << "\tcontext.time = gTime;\n";
	source << "\tcontext.frameCount = gFrameCount;\n";
	source << "\tcontext.mixAmount = gMixAmount;\n";
	source << "\tcontext.bypass = gBypass;\n";
	source << "\tcontext.sourceHistoryLength = gSourceHistoryLength;\n";
	source << "\tcontext.temporalHistoryLength = gTemporalHistoryLength;\n";
	source << "\tfloat4 effectedColor = " << shaderPackage.entryPoint << "(context);\n";
	source << "\tfloat mixValue = clamp(gBypass > 0.5 ? 0.0 : gMixAmount, 0.0, 1.0);\n";
	source << "\treturn lerp(context.sourceColor, effectedColor, mixValue);\n";
	source << "}\n";
	return source.str();
}

bool ShaderCompiler::FindSlangCompiler(std::filesystem::path& compilerPath, std::string& error) const
{
	std::filesystem::path thirdPartyRoot = mRepoRoot / "3rdParty";
	if (!std::filesystem::exists(thirdPartyRoot))
	{
		error = "3rdParty directory was not found under the repository root.";
		return false;
	}

	for (const auto& entry : std::filesystem::directory_iterator(thirdPartyRoot))
	{
		if (!entry.is_directory())
			continue;
		std::filesystem::path candidate = entry.path() / "bin" / "slangc.exe";
		if (std::filesystem::exists(candidate))
		{
			compilerPath = candidate;
			return true;
		}
	}

	error = "Could not find slangc.exe under 3rdParty.";
	return false;
}

bool ShaderCompiler::RunSlangCompiler(const std::filesystem::path& wrapperPath, const std::filesystem::path& outputPath, std::string& error) const
{
	std::filesystem::path compilerPath;
	if (!FindSlangCompiler(compilerPath, error))
		return false;

	std::string commandLine = "\"" + compilerPath.string() + "\" \"" + wrapperPath.string()
		+ "\" -target glsl -profile glsl_430 -entry fragmentMain -stage fragment -o \"" + outputPath.string() + "\"";

	STARTUPINFOA startupInfo = {};
	PROCESS_INFORMATION processInfo = {};
	startupInfo.cb = sizeof(startupInfo);
	std::vector<char> mutableCommandLine(commandLine.begin(), commandLine.end());
	mutableCommandLine.push_back('\0');

	if (!CreateProcessA(NULL, mutableCommandLine.data(), NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, mRepoRoot.string().c_str(), &startupInfo, &processInfo))
	{
		error = "Failed to launch slangc.exe.";
		return false;
	}

	WaitForSingleObject(processInfo.hProcess, INFINITE);

	DWORD exitCode = 0;
	GetExitCodeProcess(processInfo.hProcess, &exitCode);
	CloseHandle(processInfo.hThread);
	CloseHandle(processInfo.hProcess);

	if (exitCode != 0)
	{
		error = "slangc.exe returned a non-zero exit code while compiling the active shader package.";
		return false;
	}

	return true;
}

bool ShaderCompiler::PatchGeneratedGlsl(std::string& shaderText, std::string& error) const
{
	if (shaderText.find("#version 450") == std::string::npos)
	{
		error = "Generated GLSL did not include the expected version header.";
		return false;
	}

	shaderText = ReplaceAll(shaderText, "#version 450", "#version 430 core");
	shaderText = std::regex_replace(shaderText, std::regex(R"(#extension GL_EXT_samplerless_texture_functions : require\r?\n)"), "");
	shaderText = std::regex_replace(shaderText, std::regex(R"(layout\(row_major\) uniform;\r?\n)"), "");
	shaderText = std::regex_replace(shaderText, std::regex(R"(layout\(row_major\) buffer;\r?\n)"), "");
	shaderText = std::regex_replace(shaderText, std::regex(R"(layout\(location = 0\)\s*in vec2 ([A-Za-z0-9_]+);)"), "in vec2 vTexCoord;");
	shaderText = ReplaceAll(shaderText, "input_texCoord_0", "vTexCoord");

	std::smatch match;
	std::regex outRegex(R"(layout\(location = 0\)\s*out vec4 ([A-Za-z0-9_]+);)");
	if (std::regex_search(shaderText, match, outRegex))
	{
		const std::string outputName = match[1].str();
		shaderText = std::regex_replace(shaderText, outRegex, "layout(location = 0) out vec4 fragColor;");
		shaderText = ReplaceAll(shaderText, outputName + " =", "fragColor =");
	}

	return true;
}

std::string ShaderCompiler::ReadTextFile(const std::filesystem::path& path, std::string& error) const
{
	std::ifstream input(path, std::ios::binary);
	if (!input)
	{
		error = "Could not open file: " + path.string();
		return std::string();
	}

	std::ostringstream buffer;
	buffer << input.rdbuf();
	return buffer.str();
}

bool ShaderCompiler::WriteTextFile(const std::filesystem::path& path, const std::string& contents, std::string& error) const
{
	std::error_code fsError;
	std::filesystem::create_directories(path.parent_path(), fsError);

	std::ofstream output(path, std::ios::binary);
	if (!output)
	{
		error = "Could not write file: " + path.string();
		return false;
	}

	output << contents;
	return output.good();
}