#include "ShaderCompiler.h"
#include "ShaderPackageRegistry.h"

#include <filesystem>
#include <iostream>
#include <map>
#include <string>
#include <vector>

namespace
{
int gFailures = 0;

void Fail(const std::string& message)
{
	std::cerr << "FAIL: " << message << "\n";
	++gFailures;
}

std::filesystem::path FindRepoRoot()
{
	std::filesystem::path current = std::filesystem::current_path();
	for (;;)
	{
		if (std::filesystem::exists(current / "shaders") &&
			std::filesystem::exists(current / "runtime" / "templates" / "shader_wrapper.slang.in"))
		{
			return current;
		}

		if (!current.has_parent_path() || current.parent_path() == current)
			return std::filesystem::path();
		current = current.parent_path();
	}
}

bool IsExpectedInvalidPackage(const ShaderPackageStatus& status)
{
	return status.id.find("broken-shader-example") != std::string::npos &&
		status.error.find("Unsupported parameter type") != std::string::npos;
}
}

int main()
{
	const std::filesystem::path repoRoot = FindRepoRoot();
	if (repoRoot.empty())
	{
		Fail("Could not locate repository root from current working directory.");
		return 1;
	}

	std::map<std::string, ShaderPackage> packagesById;
	std::vector<std::string> packageOrder;
	std::vector<ShaderPackageStatus> packageStatuses;
	std::string error;

	ShaderPackageRegistry registry(12);
	if (!registry.Scan(repoRoot / "shaders", packagesById, packageOrder, packageStatuses, error))
	{
		Fail("Shader package scan failed: " + error);
		return 1;
	}

	for (const ShaderPackageStatus& status : packageStatuses)
	{
		if (status.available || IsExpectedInvalidPackage(status))
			continue;
		Fail("Unexpected invalid shader package '" + status.id + "': " + status.error);
	}

	const std::filesystem::path validationRoot = repoRoot / "runtime" / "shader_validation";
	const std::filesystem::path wrapperPath = validationRoot / "validation_shader_wrapper.slang";
	const std::filesystem::path generatedGlslPath = validationRoot / "validation_shader.raw.frag";
	const std::filesystem::path patchedGlslPath = validationRoot / "validation_shader.frag";

	ShaderCompiler compiler(repoRoot, wrapperPath, generatedGlslPath, patchedGlslPath, 12);
	for (const std::string& packageId : packageOrder)
	{
		auto packageIt = packagesById.find(packageId);
		if (packageIt == packagesById.end())
			continue;

		const ShaderPackage& shaderPackage = packageIt->second;
		for (const ShaderPassDefinition& pass : shaderPackage.passes)
		{
			std::string fragmentShaderSource;
			std::string compileError;
			if (!compiler.BuildPassFragmentShaderSource(shaderPackage, pass, fragmentShaderSource, compileError))
			{
				Fail("Shader package '" + packageId + "' pass '" + pass.id + "' failed Slang validation: " + compileError);
				continue;
			}
			if (fragmentShaderSource.find("#version 430 core") == std::string::npos)
				Fail("Shader package '" + packageId + "' pass '" + pass.id + "' generated GLSL without the expected patched GLSL version header.");
		}
	}

	std::error_code removeError;
	std::filesystem::remove_all(validationRoot, removeError);

	if (gFailures != 0)
	{
		std::cerr << gFailures << " shader Slang validation failure(s).\n";
		return 1;
	}

	std::cout << "Validated " << packagesById.size() << " shader package(s) through Slang.\n";
	return 0;
}
