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

#include "NativeHandles.h"

#include <fstream>
#include <cctype>
#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 SlangCBufferTypeForParameter(ShaderParameterType type)
{
	switch (type)
	{
	case ShaderParameterType::Float: return "float";
	case ShaderParameterType::Vec2: return "float2";
	case ShaderParameterType::Color: return "float4";
	case ShaderParameterType::Boolean: return "bool";
	case ShaderParameterType::Enum: return "int";
	case ShaderParameterType::Text: return "";
	case ShaderParameterType::Trigger: return "int";
	}
	return "float";
}

std::string CapitalizeIdentifier(const std::string& identifier)
{
	if (identifier.empty())
		return identifier;
	std::string text = identifier;
	text[0] = static_cast<char>(std::toupper(static_cast<unsigned char>(text[0])));
	return text;
}

std::string BuildParameterUniforms(const std::vector<ShaderParameterDefinition>& parameters)
{
	std::ostringstream source;
	for (const ShaderParameterDefinition& definition : parameters)
	{
		if (definition.type == ShaderParameterType::Text)
			continue;
		if (definition.type == ShaderParameterType::Trigger)
		{
			source << "\tint " << definition.id << ";\n";
			source << "\tfloat " << definition.id << "Time;\n";
			continue;
		}
		source << "\t" << SlangCBufferTypeForParameter(definition.type) << " " << definition.id << ";\n";
	}
	return source.str();
}

std::string BuildHistorySamplerDeclarations(const std::string& samplerPrefix, unsigned historyLength)
{
	std::ostringstream source;
	for (unsigned index = 0; index < historyLength; ++index)
		source << "Sampler2D<float4> " << samplerPrefix << index << ";\n";
	return source.str();
}

std::string BuildTextureSamplerDeclarations(const std::vector<ShaderTextureAsset>& textureAssets)
{
	std::ostringstream source;
	for (const ShaderTextureAsset& textureAsset : textureAssets)
		source << "Sampler2D<float4> " << textureAsset.id << ";\n";
	if (!textureAssets.empty())
		source << "\n";
	return source.str();
}

std::string BuildTextSamplerDeclarations(const std::vector<ShaderParameterDefinition>& parameters)
{
	std::ostringstream source;
	for (const ShaderParameterDefinition& definition : parameters)
	{
		if (definition.type != ShaderParameterType::Text)
			continue;
		source << "Sampler2D<float4> " << definition.id << "Texture;\n";
	}
	if (source.tellp() > 0)
		source << "\n";
	return source.str();
}

std::string BuildTextHelpers(const std::vector<ShaderParameterDefinition>& parameters)
{
	std::ostringstream source;
	for (const ShaderParameterDefinition& definition : parameters)
	{
		if (definition.type != ShaderParameterType::Text)
			continue;
		const std::string suffix = CapitalizeIdentifier(definition.id);
		source
			<< "float sample" << suffix << "(float2 uv)\n"
			<< "{\n"
			<< "\tif (uv.x < 0.0 || uv.x > 1.0 || uv.y < 0.0 || uv.y > 1.0)\n"
			<< "\t\treturn 0.0;\n"
			<< "\treturn " << definition.id << "Texture.Sample(uv).r;\n"
			<< "}\n\n"
			<< "float4 draw" << suffix << "(float2 uv, float4 fillColor)\n"
			<< "{\n"
			<< "\tfloat alpha = sample" << suffix << "(uv) * fillColor.a;\n"
			<< "\treturn float4(fillColor.rgb * alpha, alpha);\n"
			<< "}\n\n";
	}
	return source.str();
}

std::string BuildHistorySwitchCases(const std::string& samplerPrefix, unsigned historyLength)
{
	std::ostringstream source;
	for (unsigned index = 0; index < historyLength; ++index)
		source << "\tcase " << index << ": return " << samplerPrefix << index << ".Sample(tc);\n";
	return source.str();
}
}

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
{
	std::string wrapperSource;
	if (!BuildWrapperSlangSource(shaderPackage, wrapperSource, error))
		return false;
	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;
}

bool ShaderCompiler::BuildWrapperSlangSource(const ShaderPackage& shaderPackage, std::string& wrapperSource, std::string& error) const
{
	const std::filesystem::path templatePath = mRepoRoot / "runtime" / "templates" / "shader_wrapper.slang.in";
	wrapperSource = ReadTextFile(templatePath, error);
	if (wrapperSource.empty())
		return false;

	wrapperSource = ReplaceAll(wrapperSource, "{{PARAMETER_UNIFORMS}}", BuildParameterUniforms(shaderPackage.parameters));
	const unsigned historySamplerCount = shaderPackage.temporal.enabled ? mMaxTemporalHistoryFrames : 0;
	wrapperSource = ReplaceAll(wrapperSource, "{{SOURCE_HISTORY_SAMPLERS}}", BuildHistorySamplerDeclarations("gSourceHistory", historySamplerCount));
	wrapperSource = ReplaceAll(wrapperSource, "{{TEMPORAL_HISTORY_SAMPLERS}}", BuildHistorySamplerDeclarations("gTemporalHistory", historySamplerCount));
	wrapperSource = ReplaceAll(wrapperSource, "{{TEXTURE_SAMPLERS}}", BuildTextureSamplerDeclarations(shaderPackage.textureAssets));
	wrapperSource = ReplaceAll(wrapperSource, "{{TEXT_SAMPLERS}}", BuildTextSamplerDeclarations(shaderPackage.parameters));
	wrapperSource = ReplaceAll(wrapperSource, "{{TEXT_HELPERS}}", BuildTextHelpers(shaderPackage.parameters));
	wrapperSource = ReplaceAll(wrapperSource, "{{SOURCE_HISTORY_SWITCH_CASES}}", BuildHistorySwitchCases("gSourceHistory", historySamplerCount));
	wrapperSource = ReplaceAll(wrapperSource, "{{TEMPORAL_HISTORY_SWITCH_CASES}}", BuildHistorySwitchCases("gTemporalHistory", historySamplerCount));
	wrapperSource = ReplaceAll(wrapperSource, "{{USER_SHADER_INCLUDE}}", shaderPackage.shaderPath.generic_string());
	wrapperSource = ReplaceAll(wrapperSource, "{{ENTRY_POINT_CALL}}", shaderPackage.entryPoint + "(context)");
	return true;
}

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;
	}

	UniqueHandle processHandle(processInfo.hProcess);
	UniqueHandle threadHandle(processInfo.hThread);

	WaitForSingleObject(processHandle.get(), INFINITE);

	DWORD exitCode = 0;
	GetExitCodeProcess(processHandle.get(), &exitCode);

	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();
}