#include "fixtures/RHITestFixture.h"
#include "XCEngine/RHI/RHIPipelineLayout.h"
#include "XCEngine/RHI/RHIPipelineState.h"
#include <cstring>

using namespace XCEngine::RHI;

TEST_P(RHITestFixture, PipelineState_Create_DefaultDesc) {
    GraphicsPipelineDesc desc = {};
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    if (pso != nullptr) {
        if (GetBackendType() == RHIType::D3D12) {
            EXPECT_FALSE(pso->IsValid());
        } else {
            EXPECT_TRUE(pso->IsValid());
        }
        pso->Shutdown();
        delete pso;
    }
}

TEST_P(RHITestFixture, PipelineState_SetGet_RasterizerState) {
    GraphicsPipelineDesc desc = {};
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    
    RasterizerDesc rasterizer = {};
    rasterizer.fillMode = 2;
    rasterizer.cullMode = 2;
    rasterizer.depthClipEnable = true;
    pso->SetRasterizerState(rasterizer);
    
    const RasterizerDesc& retrieved = pso->GetRasterizerState();
    EXPECT_EQ(retrieved.cullMode, 2u);
    EXPECT_EQ(retrieved.fillMode, 2u);
    
    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_SetGet_BlendState) {
    GraphicsPipelineDesc desc = {};
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    
    BlendDesc blend = {};
    blend.blendEnable = true;
    blend.srcBlend = 1;
    blend.dstBlend = 0;
    pso->SetBlendState(blend);
    
    const BlendDesc& retrieved = pso->GetBlendState();
    EXPECT_TRUE(retrieved.blendEnable);
    EXPECT_EQ(retrieved.srcBlend, 1u);
    
    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_SetGet_DepthStencilState) {
    GraphicsPipelineDesc desc = {};
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    
    DepthStencilStateDesc ds = {};
    ds.depthTestEnable = true;
    ds.depthWriteEnable = true;
    ds.depthFunc = 3;
    pso->SetDepthStencilState(ds);
    
    const DepthStencilStateDesc& retrieved = pso->GetDepthStencilState();
    EXPECT_TRUE(retrieved.depthTestEnable);
    EXPECT_TRUE(retrieved.depthWriteEnable);
    EXPECT_EQ(retrieved.depthFunc, 3u);
    
    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_SetGet_InputLayout) {
    GraphicsPipelineDesc desc = {};
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    
    InputElementDesc element = {};
    element.semanticName = "POSITION";
    element.semanticIndex = 0;
    element.format = static_cast<uint32_t>(Format::R32G32B32A32_Float);
    element.inputSlot = 0;
    element.alignedByteOffset = 0;
    
    InputLayoutDesc layoutDesc = {};
    layoutDesc.elements.push_back(element);
    pso->SetInputLayout(layoutDesc);
    
    const InputLayoutDesc& retrieved = pso->GetInputLayout();
    ASSERT_EQ(retrieved.elements.size(), 1u);
    EXPECT_EQ(retrieved.elements[0].semanticName, "POSITION");
    
    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_SetTopology) {
    GraphicsPipelineDesc desc = {};
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    
    pso->SetTopology(3);
    
    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_SetRenderTargetFormats) {
    GraphicsPipelineDesc desc = {};
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    
    uint32_t formats[] = { static_cast<uint32_t>(Format::R8G8B8A8_UNorm) };
    pso->SetRenderTargetFormats(1, formats, static_cast<uint32_t>(Format::D24_UNorm_S8_UInt));
    
    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_EnsureValid_IsValid) {
    GraphicsPipelineDesc desc = {};
    desc.renderTargetFormats[0] = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    desc.depthStencilFormat = static_cast<uint32_t>(Format::D24_UNorm_S8_UInt);
    
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    
    if (GetBackendType() == RHIType::D3D12) {
        EXPECT_FALSE(pso->IsValid());
        pso->EnsureValid();
        EXPECT_FALSE(pso->IsValid());
    } else {
        EXPECT_TRUE(pso->IsValid());
        pso->EnsureValid();
        EXPECT_TRUE(pso->IsValid());
    }
    
    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_Shutdown_Invalidates) {
    GraphicsPipelineDesc desc = {};
    desc.renderTargetFormats[0] = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    desc.depthStencilFormat = static_cast<uint32_t>(Format::D24_UNorm_S8_UInt);
    
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    
    pso->EnsureValid();
    pso->Shutdown();
    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_GetType) {
    GraphicsPipelineDesc desc = {};
    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    
    PipelineType type = pso->GetType();
    EXPECT_EQ(type, PipelineType::Graphics);
    
    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_Create_GraphicsShadersFromDesc) {
    GraphicsPipelineDesc desc = {};
    desc.topologyType = static_cast<uint32_t>(PrimitiveTopologyType::Triangle);
    desc.renderTargetFormats[0] = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    desc.depthStencilFormat = static_cast<uint32_t>(Format::Unknown);

    InputElementDesc position = {};
    position.semanticName = "POSITION";
    position.semanticIndex = 0;
    position.format = static_cast<uint32_t>(Format::R32G32B32A32_Float);
    position.inputSlot = 0;
    position.alignedByteOffset = 0;
    desc.inputLayout.elements.push_back(position);

    if (GetBackendType() == RHIType::D3D12) {
        const char* hlslSource = R"(
struct VSInput {
    float4 position : POSITION;
};

struct PSInput {
    float4 position : SV_POSITION;
};

PSInput MainVS(VSInput input) {
    PSInput output;
    output.position = input.position;
    return output;
}

float4 MainPS(PSInput input) : SV_TARGET {
    return float4(1.0f, 0.0f, 0.0f, 1.0f);
}
)";

        desc.vertexShader.source.assign(hlslSource, hlslSource + std::strlen(hlslSource));
        desc.vertexShader.entryPoint = L"MainVS";
        desc.vertexShader.profile = L"vs_5_0";
        desc.vertexShader.sourceLanguage = ShaderLanguage::HLSL;

        desc.fragmentShader.source.assign(hlslSource, hlslSource + std::strlen(hlslSource));
        desc.fragmentShader.entryPoint = L"MainPS";
        desc.fragmentShader.profile = L"ps_5_0";
        desc.fragmentShader.sourceLanguage = ShaderLanguage::HLSL;
    } else {
        const char* vertexSource = R"(#version 430
layout(location = 0) in vec4 aPosition;

void main() {
    gl_Position = aPosition;
}
)";
        const char* fragmentSource = R"(#version 430
layout(location = 0) out vec4 fragColor;

void main() {
    fragColor = vec4(1.0, 0.0, 0.0, 1.0);
}
)";

        desc.vertexShader.source.assign(vertexSource, vertexSource + std::strlen(vertexSource));
        desc.vertexShader.sourceLanguage = ShaderLanguage::GLSL;
        desc.vertexShader.profile = L"vs_4_30";

        desc.fragmentShader.source.assign(fragmentSource, fragmentSource + std::strlen(fragmentSource));
        desc.fragmentShader.sourceLanguage = ShaderLanguage::GLSL;
        desc.fragmentShader.profile = L"fs_4_30";
    }

    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    EXPECT_TRUE(pso->IsValid());
    EXPECT_NE(pso->GetNativeHandle(), nullptr);

    pso->Shutdown();
    delete pso;
}

TEST_P(RHITestFixture, PipelineState_Create_WithPipelineLayout) {
    RHIPipelineLayoutDesc layoutDesc = {};
    layoutDesc.textureCount = 1;
    layoutDesc.samplerCount = 1;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(layoutDesc);
    ASSERT_NE(layout, nullptr);
    ASSERT_NE(layout->GetNativeHandle(), nullptr);

    GraphicsPipelineDesc desc = {};
    desc.pipelineLayout = layout;
    desc.topologyType = static_cast<uint32_t>(PrimitiveTopologyType::Triangle);
    desc.renderTargetFormats[0] = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    desc.depthStencilFormat = static_cast<uint32_t>(Format::Unknown);

    InputElementDesc position = {};
    position.semanticName = "POSITION";
    position.semanticIndex = 0;
    position.format = static_cast<uint32_t>(Format::R32G32B32A32_Float);
    position.inputSlot = 0;
    position.alignedByteOffset = 0;
    desc.inputLayout.elements.push_back(position);

    if (GetBackendType() == RHIType::D3D12) {
        const char* hlslSource = R"(
struct VSInput {
    float4 position : POSITION;
};

struct PSInput {
    float4 position : SV_POSITION;
};

PSInput MainVS(VSInput input) {
    PSInput output;
    output.position = input.position;
    return output;
}

float4 MainPS(PSInput input) : SV_TARGET {
    return float4(1.0f, 0.0f, 0.0f, 1.0f);
}
)";

        desc.vertexShader.source.assign(hlslSource, hlslSource + std::strlen(hlslSource));
        desc.vertexShader.entryPoint = L"MainVS";
        desc.vertexShader.profile = L"vs_5_0";
        desc.vertexShader.sourceLanguage = ShaderLanguage::HLSL;

        desc.fragmentShader.source.assign(hlslSource, hlslSource + std::strlen(hlslSource));
        desc.fragmentShader.entryPoint = L"MainPS";
        desc.fragmentShader.profile = L"ps_5_0";
        desc.fragmentShader.sourceLanguage = ShaderLanguage::HLSL;
    } else {
        const char* vertexSource = R"(#version 430
layout(location = 0) in vec4 aPosition;

void main() {
    gl_Position = aPosition;
}
)";
        const char* fragmentSource = R"(#version 430
layout(location = 0) out vec4 fragColor;

void main() {
    fragColor = vec4(1.0, 0.0, 0.0, 1.0);
}
)";

        desc.vertexShader.source.assign(vertexSource, vertexSource + std::strlen(vertexSource));
        desc.vertexShader.sourceLanguage = ShaderLanguage::GLSL;
        desc.vertexShader.profile = L"vs_4_30";

        desc.fragmentShader.source.assign(fragmentSource, fragmentSource + std::strlen(fragmentSource));
        desc.fragmentShader.sourceLanguage = ShaderLanguage::GLSL;
        desc.fragmentShader.profile = L"fs_4_30";
    }

    RHIPipelineState* pso = GetDevice()->CreatePipelineState(desc);
    ASSERT_NE(pso, nullptr);
    EXPECT_TRUE(pso->IsValid());
    EXPECT_NE(pso->GetNativeHandle(), nullptr);

    pso->Shutdown();
    delete pso;
    layout->Shutdown();
    delete layout;
}