#include <Nazara/Core/File.hpp>
#include <Nazara/Shader/GlslWriter.hpp>
#include <Nazara/Shader/ShaderBuilder.hpp>
#include <Nazara/Shader/SpirvPrinter.hpp>
#include <Nazara/Shader/SpirvWriter.hpp>
#include <Catch/catch.hpp>
#include <cctype>

std::string_view Trim(std::string_view str)
{
	while (!str.empty() && std::isspace(str.front()))
		str.remove_prefix(1);

	while (!str.empty() && std::isspace(str.back()))
		str.remove_suffix(1);

	return str;
}

void ExpectingGLSL(Nz::ShaderAst::StatementPtr& shader, std::string_view expectedOutput)
{
	Nz::GlslWriter writer;

	std::string output = writer.Generate(shader);
	std::size_t funcOffset = output.find("void main()");
	std::string_view subset = Trim(output).substr(funcOffset);
	expectedOutput = Trim(expectedOutput);

	REQUIRE(subset == expectedOutput);
}

void ExpectingSpirV(Nz::ShaderAst::StatementPtr& shader, std::string_view expectedOutput)
{
	Nz::SpirvWriter writer;
	auto spirv = writer.Generate(shader);

	Nz::SpirvPrinter printer;

	Nz::SpirvPrinter::Settings settings;
	settings.printHeader = false;
	settings.printParameters = false;

	std::string output = printer.Print(spirv.data(), spirv.size(), settings);
	std::size_t funcOffset = output.find("OpFunction");
	std::string_view subset = Trim(output).substr(funcOffset);
	expectedOutput = Trim(expectedOutput);

	REQUIRE(subset == expectedOutput);
}

SCENARIO("Shader generation", "[Shader]")
{
	SECTION("Nested member loading")
	{
		std::vector<Nz::ShaderAst::StatementPtr> statements;

		Nz::ShaderAst::StructDescription innerStructDesc;
		{
			innerStructDesc.name = "innerStruct";
			auto& member = innerStructDesc.members.emplace_back();
			member.name = "field";
			member.type = Nz::ShaderAst::VectorType{ 3, Nz::ShaderAst::PrimitiveType::Float32 };
		}
		statements.push_back(Nz::ShaderBuilder::DeclareStruct(std::move(innerStructDesc)));

		Nz::ShaderAst::StructDescription outerStruct;
		{
			outerStruct.name = "outerStruct";
			auto& member = outerStruct.members.emplace_back();
			member.name = "s";
			member.type = Nz::ShaderAst::IdentifierType{ "innerStruct" };
		}
		statements.push_back(Nz::ShaderBuilder::DeclareStruct(std::move(outerStruct)));

		auto external = std::make_unique<Nz::ShaderAst::DeclareExternalStatement>();
		external->externalVars.push_back({
			std::nullopt,
			"ubo",
			Nz::ShaderAst::IdentifierType{ "outerStruct" }
		});
		statements.push_back(std::move(external));

		SECTION("Nested AccessMember")
		{
			auto ubo = Nz::ShaderBuilder::Identifier("ubo");
			auto firstAccess = Nz::ShaderBuilder::AccessMember(std::move(ubo), { "s" });
			auto secondAccess = Nz::ShaderBuilder::AccessMember(std::move(firstAccess), { "field" });

			auto swizzle = Nz::ShaderBuilder::Swizzle(std::move(secondAccess), { Nz::ShaderAst::SwizzleComponent::Third });
			auto varDecl = Nz::ShaderBuilder::DeclareVariable("result", Nz::ShaderAst::PrimitiveType::Float32, std::move(swizzle));

			statements.push_back(Nz::ShaderBuilder::DeclareFunction("main", std::move(varDecl)));

			Nz::ShaderAst::StatementPtr shader = Nz::ShaderBuilder::MultiStatement(std::move(statements));

			SECTION("Generating GLSL")
			{
				ExpectingGLSL(shader, R"(
void main()
{
	float result = ubo.s.field.z;
}
)");
			}
			SECTION("Generating Spir-V")
			{
				ExpectingSpirV(shader, R"(
OpFunction
OpLabel
OpVariable
OpAccessChain
OpAccessChain
OpLoad
OpCompositeExtract
OpStore
OpReturn
OpFunctionEnd)");
			}
		}

		SECTION("AccessMember with multiples fields")
		{
			auto ubo = Nz::ShaderBuilder::Identifier("ubo");
			auto access = Nz::ShaderBuilder::AccessMember(std::move(ubo), { "s", "field" });

			auto swizzle = Nz::ShaderBuilder::Swizzle(std::move(access), { Nz::ShaderAst::SwizzleComponent::Third });
			auto varDecl = Nz::ShaderBuilder::DeclareVariable("result", Nz::ShaderAst::PrimitiveType::Float32, std::move(swizzle));

			statements.push_back(Nz::ShaderBuilder::DeclareFunction("main", std::move(varDecl)));

			Nz::ShaderAst::StatementPtr shader = Nz::ShaderBuilder::MultiStatement(std::move(statements));

			SECTION("Generating GLSL")
			{
				ExpectingGLSL(shader, R"(
void main()
{
	float result = ubo.s.field.z;
}
)");
			}
			SECTION("Generating Spir-V")
			{
				ExpectingSpirV(shader, R"(
OpFunction
OpLabel
OpVariable
OpAccessChain
OpLoad
OpCompositeExtract
OpStore
OpReturn
OpFunctionEnd)");
			}
		}
	}
}