#include <gtest/gtest.h>
#include <XCEngine/Resources/Shader/Shader.h>
#include <XCEngine/Core/Asset/ResourceTypes.h>
#include <XCEngine/Core/Containers/Array.h>
#include <XCEngine/RHI/RHITypes.h>
#include "Rendering/Internal/ShaderVariantUtils.h"

using namespace XCEngine::Resources;
using namespace XCEngine::Containers;

namespace {

TEST(Shader, DefaultConstructor) {
    Shader shader;
    EXPECT_EQ(shader.GetPassCount(), 0u);
    EXPECT_FALSE(shader.IsValid());
    EXPECT_EQ(shader.GetMemorySize(), 0u);
}

TEST(Shader, GetType) {
    Shader shader;
    EXPECT_EQ(shader.GetType(), ResourceType::Shader);
}

TEST(Shader, AddPassVariantStoresStageLanguageAndPayload) {
    Shader shader;
    ShaderStageVariant variant = {};
    variant.stage = ShaderType::Vertex;
    variant.language = ShaderLanguage::HLSL;
    variant.backend = ShaderBackend::D3D12;
    variant.entryPoint = "MainVS";
    variant.profile = "vs_5_0";
    variant.sourceCode = "float4 MainVS() : SV_POSITION { return 0; }";
    variant.compiledBinary = { 0x00, 0x01, 0x02, 0x03 };

    shader.AddPassVariant("ForwardLit", variant);

    const ShaderStageVariant* storedVariant =
        shader.FindVariant("ForwardLit", ShaderType::Vertex, ShaderBackend::D3D12);
    ASSERT_NE(storedVariant, nullptr);
    EXPECT_EQ(storedVariant->stage, ShaderType::Vertex);
    EXPECT_EQ(storedVariant->language, ShaderLanguage::HLSL);
    EXPECT_EQ(storedVariant->entryPoint, "MainVS");
    EXPECT_EQ(storedVariant->profile, "vs_5_0");
    EXPECT_EQ(storedVariant->sourceCode, "float4 MainVS() : SV_POSITION { return 0; }");
    ASSERT_EQ(storedVariant->compiledBinary.Size(), 4u);
    EXPECT_EQ(storedVariant->compiledBinary[3], 0x03);
}

TEST(Shader, AddGetUniforms) {
    Shader shader;

    ShaderUniform uniform1;
    uniform1.name = "uModelView";
    uniform1.location = 0;
    uniform1.size = 1;
    uniform1.type = 4;

    shader.AddUniform(uniform1);

    const auto& uniforms = shader.GetUniforms();
    EXPECT_EQ(uniforms.Size(), 1u);
    EXPECT_EQ(uniforms[0].name, "uModelView");
}

TEST(Shader, AddGetAttributes) {
    Shader shader;

    ShaderAttribute attr1;
    attr1.name = "aPosition";
    attr1.location = 0;
    attr1.size = 1;
    attr1.type = 3;

    shader.AddAttribute(attr1);

    const auto& attributes = shader.GetAttributes();
    EXPECT_EQ(attributes.Size(), 1u);
    EXPECT_EQ(attributes[0].name, "aPosition");
}

TEST(Shader, StoresBackendCompiledBinariesPerBackend) {
    ShaderStageVariant variant = {};
    variant.stage = ShaderType::Fragment;
    variant.language = ShaderLanguage::HLSL;
    variant.backend = ShaderBackend::Generic;

    Array<XCEngine::Core::uint8> d3d12Payload;
    d3d12Payload.PushBack(0x01);
    d3d12Payload.PushBack(0x02);
    variant.SetCompiledBinaryForBackend(ShaderBackend::D3D12, d3d12Payload);

    Array<XCEngine::Core::uint8> vulkanPayload;
    vulkanPayload.PushBack(0x03);
    vulkanPayload.PushBack(0x04);
    vulkanPayload.PushBack(0x05);
    variant.SetCompiledBinaryForBackend(ShaderBackend::Vulkan, vulkanPayload);

    const Array<XCEngine::Core::uint8>* d3d12Binary =
        variant.GetCompiledBinaryForBackend(ShaderBackend::D3D12);
    ASSERT_NE(d3d12Binary, nullptr);
    ASSERT_EQ(d3d12Binary->Size(), 2u);
    EXPECT_EQ((*d3d12Binary)[1], 0x02);

    const Array<XCEngine::Core::uint8>* vulkanBinary =
        variant.GetCompiledBinaryForBackend(ShaderBackend::Vulkan);
    ASSERT_NE(vulkanBinary, nullptr);
    ASSERT_EQ(vulkanBinary->Size(), 3u);
    EXPECT_EQ((*vulkanBinary)[2], 0x05);

    EXPECT_EQ(variant.GetCompiledBinaryForBackend(ShaderBackend::OpenGL), nullptr);
}

TEST(Shader, ApplyShaderStageVariantCarriesMatchedBackendCompiledBinary) {
    ShaderPass pass = {};
    pass.name = "ForwardLit";

    ShaderStageVariant variant = {};
    variant.stage = ShaderType::Fragment;
    variant.language = ShaderLanguage::HLSL;
    variant.backend = ShaderBackend::Generic;
    variant.entryPoint = "MainPS";
    variant.profile = "ps_5_1";
    variant.sourceCode = "float4 MainPS() : SV_TARGET { return 1; }";

    Array<XCEngine::Core::uint8> vulkanPayload;
    vulkanPayload.PushBack(0x07);
    vulkanPayload.PushBack(0x08);
    vulkanPayload.PushBack(0x09);
    variant.SetCompiledBinaryForBackend(ShaderBackend::Vulkan, vulkanPayload);

    XCEngine::RHI::ShaderCompileDesc compileDesc = {};
    ::XCEngine::Rendering::Internal::ApplyShaderStageVariant(
        pass,
        ShaderBackend::Vulkan,
        variant,
        compileDesc);

    EXPECT_EQ(compileDesc.compiledBinaryBackend, XCEngine::RHI::ShaderBinaryBackend::Vulkan);
    ASSERT_EQ(compileDesc.compiledBinary.size(), 3u);
    EXPECT_EQ(compileDesc.compiledBinary[0], 0x07);
    EXPECT_EQ(compileDesc.compiledBinary[2], 0x09);
}

TEST(Shader, FindsBackendSpecificVariantAndFallsBackToGeneric) {
    Shader shader;

    ShaderStageVariant genericFragment = {};
    genericFragment.stage = ShaderType::Fragment;
    genericFragment.language = ShaderLanguage::GLSL;
    genericFragment.backend = ShaderBackend::Generic;
    genericFragment.sourceCode = "generic fragment";
    shader.AddPassVariant("ForwardLit", genericFragment);

    ShaderStageVariant d3d12Fragment = {};
    d3d12Fragment.stage = ShaderType::Fragment;
    d3d12Fragment.language = ShaderLanguage::HLSL;
    d3d12Fragment.backend = ShaderBackend::D3D12;
    d3d12Fragment.sourceCode = "d3d12 fragment";
    shader.AddPassVariant("ForwardLit", d3d12Fragment);

    const ShaderStageVariant* d3d12Variant =
        shader.FindVariant("ForwardLit", ShaderType::Fragment, ShaderBackend::D3D12);
    ASSERT_NE(d3d12Variant, nullptr);
    EXPECT_EQ(d3d12Variant->sourceCode, "d3d12 fragment");

    const ShaderStageVariant* openglVariant =
        shader.FindVariant("ForwardLit", ShaderType::Fragment, ShaderBackend::OpenGL);
    ASSERT_NE(openglVariant, nullptr);
    EXPECT_EQ(openglVariant->sourceCode, "generic fragment");
}

TEST(Shader, FindsMostSpecificKeywordVariantAndPrefersExactBackend) {
    Shader shader;

    ShaderStageVariant baseFragment = {};
    baseFragment.stage = ShaderType::Fragment;
    baseFragment.language = ShaderLanguage::GLSL;
    baseFragment.backend = ShaderBackend::Generic;
    baseFragment.sourceCode = "generic base";
    shader.AddPassVariant("ForwardLit", baseFragment);

    ShaderStageVariant genericKeywordFragment = {};
    genericKeywordFragment.stage = ShaderType::Fragment;
    genericKeywordFragment.language = ShaderLanguage::GLSL;
    genericKeywordFragment.backend = ShaderBackend::Generic;
    genericKeywordFragment.requiredKeywords.enabledKeywords.PushBack("XC_DEBUG");
    genericKeywordFragment.requiredKeywords.enabledKeywords.PushBack("XC_ALPHA_TEST");
    genericKeywordFragment.requiredKeywords.enabledKeywords.PushBack("XC_DEBUG");
    genericKeywordFragment.sourceCode = "generic keyword";
    shader.AddPassVariant("ForwardLit", genericKeywordFragment);

    ShaderStageVariant d3d12KeywordFragment = {};
    d3d12KeywordFragment.stage = ShaderType::Fragment;
    d3d12KeywordFragment.language = ShaderLanguage::HLSL;
    d3d12KeywordFragment.backend = ShaderBackend::D3D12;
    d3d12KeywordFragment.requiredKeywords.enabledKeywords.PushBack("XC_ALPHA_TEST");
    d3d12KeywordFragment.sourceCode = "d3d12 keyword";
    shader.AddPassVariant("ForwardLit", d3d12KeywordFragment);

    ShaderKeywordSet enabledKeywords = {};
    enabledKeywords.enabledKeywords.PushBack("XC_ALPHA_TEST");
    enabledKeywords.enabledKeywords.PushBack("XC_DEBUG");

    const ShaderStageVariant* d3d12Variant =
        shader.FindVariant(
            "ForwardLit",
            ShaderType::Fragment,
            ShaderBackend::D3D12,
            enabledKeywords);
    ASSERT_NE(d3d12Variant, nullptr);
    EXPECT_EQ(d3d12Variant->sourceCode, "d3d12 keyword");

    const ShaderStageVariant* openglVariant =
        shader.FindVariant(
            "ForwardLit",
            ShaderType::Fragment,
            ShaderBackend::OpenGL,
            enabledKeywords);
    ASSERT_NE(openglVariant, nullptr);
    EXPECT_EQ(openglVariant->sourceCode, "generic keyword");
    ASSERT_EQ(openglVariant->requiredKeywords.enabledKeywords.Size(), 2u);
    EXPECT_EQ(openglVariant->requiredKeywords.enabledKeywords[0], "XC_ALPHA_TEST");
    EXPECT_EQ(openglVariant->requiredKeywords.enabledKeywords[1], "XC_DEBUG");
}

TEST(Shader, RejectsVariantWhenRequiredKeywordsAreMissing) {
    Shader shader;

    ShaderStageVariant keywordFragment = {};
    keywordFragment.stage = ShaderType::Fragment;
    keywordFragment.language = ShaderLanguage::GLSL;
    keywordFragment.backend = ShaderBackend::OpenGL;
    keywordFragment.requiredKeywords.enabledKeywords.PushBack("XC_CLIP");
    keywordFragment.sourceCode = "clip fragment";
    shader.AddPassVariant("ForwardLit", keywordFragment);

    EXPECT_EQ(
        shader.FindVariant("ForwardLit", ShaderType::Fragment, ShaderBackend::OpenGL),
        nullptr);

    ShaderKeywordSet enabledKeywords = {};
    enabledKeywords.enabledKeywords.PushBack("XC_CLIP");

    const ShaderStageVariant* matchedVariant =
        shader.FindVariant(
            "ForwardLit",
            ShaderType::Fragment,
            ShaderBackend::OpenGL,
            enabledKeywords);
    ASSERT_NE(matchedVariant, nullptr);
    EXPECT_EQ(matchedVariant->sourceCode, "clip fragment");
}

TEST(Shader, StoresPerPassTags) {
    Shader shader;
    shader.SetPassTag("ForwardLit", "LightMode", "ForwardBase");
    shader.SetPassTag("ForwardLit", "Queue", "Geometry");

    const ShaderPass* pass = shader.FindPass("ForwardLit");
    ASSERT_NE(pass, nullptr);
    ASSERT_EQ(pass->tags.Size(), 2u);
    EXPECT_EQ(pass->tags[0].name, "LightMode");
    EXPECT_EQ(pass->tags[0].value, "ForwardBase");
    EXPECT_EQ(pass->tags[1].name, "Queue");
    EXPECT_EQ(pass->tags[1].value, "Geometry");
}

TEST(Shader, StoresShaderPropertiesAndPassResources) {
    Shader shader;

    ShaderPropertyDesc baseColor = {};
    baseColor.name = "_BaseColor";
    baseColor.displayName = "Base Color";
    baseColor.type = ShaderPropertyType::Color;
    baseColor.defaultValue = "(1,1,1,1)";
    baseColor.semantic = "BaseColor";
    shader.AddProperty(baseColor);

    ShaderResourceBindingDesc perObject = {};
    perObject.name = "PerObjectConstants";
    perObject.type = ShaderResourceType::ConstantBuffer;
    perObject.set = 1;
    perObject.binding = 0;
    perObject.semantic = "PerObject";
    shader.AddPassResourceBinding("ForwardLit", perObject);

    ASSERT_EQ(shader.GetProperties().Size(), 1u);
    const ShaderPropertyDesc* storedProperty = shader.FindProperty("_BaseColor");
    ASSERT_NE(storedProperty, nullptr);
    EXPECT_EQ(storedProperty->displayName, "Base Color");
    EXPECT_EQ(storedProperty->type, ShaderPropertyType::Color);
    EXPECT_EQ(storedProperty->semantic, "BaseColor");

    const ShaderResourceBindingDesc* storedBinding =
        shader.FindPassResourceBinding("ForwardLit", "PerObjectConstants");
    ASSERT_NE(storedBinding, nullptr);
    EXPECT_EQ(storedBinding->type, ShaderResourceType::ConstantBuffer);
    EXPECT_EQ(storedBinding->set, 1u);
    EXPECT_EQ(storedBinding->binding, 0u);
    EXPECT_EQ(storedBinding->semantic, "PerObject");
}

TEST(Shader, StoresPassKeywordDeclarationsAndQueriesDeclaredKeywords) {
    Shader shader;

    ShaderKeywordDeclaration globalKeywords = {};
    globalKeywords.type = ShaderKeywordDeclarationType::MultiCompile;
    globalKeywords.options.PushBack("_");
    globalKeywords.options.PushBack("XC_MAIN_LIGHT_SHADOWS");
    globalKeywords.options.PushBack("XC_ADDITIONAL_LIGHTS");
    shader.AddPassKeywordDeclaration("ForwardLit", globalKeywords);

    ShaderKeywordDeclaration localKeywords = {};
    localKeywords.type = ShaderKeywordDeclarationType::ShaderFeatureLocal;
    localKeywords.options.PushBack("__");
    localKeywords.options.PushBack("XC_ALPHA_TEST");
    shader.AddPassKeywordDeclaration("ShadowCaster", localKeywords);

    const ShaderPass* forwardPass = shader.FindPass("ForwardLit");
    ASSERT_NE(forwardPass, nullptr);
    ASSERT_EQ(forwardPass->keywordDeclarations.Size(), 1u);
    EXPECT_EQ(
        forwardPass->keywordDeclarations[0].type,
        ShaderKeywordDeclarationType::MultiCompile);
    ASSERT_EQ(forwardPass->keywordDeclarations[0].options.Size(), 3u);
    EXPECT_EQ(forwardPass->keywordDeclarations[0].options[1], "XC_MAIN_LIGHT_SHADOWS");

    EXPECT_TRUE(shader.PassDeclaresKeyword("ForwardLit", "XC_MAIN_LIGHT_SHADOWS"));
    EXPECT_TRUE(shader.PassDeclaresKeyword("ShadowCaster", "XC_ALPHA_TEST"));
    EXPECT_TRUE(shader.DeclaresKeyword("XC_ADDITIONAL_LIGHTS"));
    EXPECT_FALSE(shader.DeclaresKeyword("_"));
    EXPECT_FALSE(shader.DeclaresKeyword("XC_UNKNOWN"));
}

TEST(Shader, ReleaseClearsPassRuntimeData) {
    Shader shader;
    ShaderPropertyDesc property = {};
    property.name = "_BaseColor";
    property.type = ShaderPropertyType::Color;
    shader.AddProperty(property);
    ShaderStageVariant variant = {};
    variant.stage = ShaderType::Fragment;
    variant.sourceCode = "fragment";
    shader.AddPassVariant("ForwardLit", variant);

    shader.Release();

    EXPECT_EQ(shader.GetProperties().Size(), 0u);
    EXPECT_EQ(shader.GetPassCount(), 0u);
    EXPECT_FALSE(shader.IsValid());
}

} // namespace