#include "fixtures/RHITestFixture.h"
#include "XCEngine/RHI/D3D12/D3D12PipelineLayout.h"
#include "XCEngine/RHI/OpenGL/OpenGLPipelineLayout.h"
#include "XCEngine/RHI/RHIPipelineLayout.h"
#include "XCEngine/RHI/RHIDescriptorSet.h"

using namespace XCEngine::RHI;

TEST_P(RHITestFixture, PipelineLayout_Create_Basic) {
    RHIPipelineLayoutDesc desc = {};
    desc.constantBufferCount = 1;
    desc.textureCount = 0;
    desc.samplerCount = 0;
    desc.uavCount = 0;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(desc);
    ASSERT_NE(layout, nullptr);
    EXPECT_NE(layout->GetNativeHandle(), nullptr);
    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_Create_WithTextures) {
    RHIPipelineLayoutDesc desc = {};
    desc.constantBufferCount = 0;
    desc.textureCount = 4;
    desc.samplerCount = 0;
    desc.uavCount = 0;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(desc);
    ASSERT_NE(layout, nullptr);
    EXPECT_NE(layout->GetNativeHandle(), nullptr);
    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_Create_WithSamplers) {
    RHIPipelineLayoutDesc desc = {};
    desc.constantBufferCount = 0;
    desc.textureCount = 0;
    desc.samplerCount = 2;
    desc.uavCount = 0;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(desc);
    ASSERT_NE(layout, nullptr);
    EXPECT_NE(layout->GetNativeHandle(), nullptr);
    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_Create_WithUAVs) {
    RHIPipelineLayoutDesc desc = {};
    desc.constantBufferCount = 0;
    desc.textureCount = 0;
    desc.samplerCount = 0;
    desc.uavCount = 3;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(desc);
    ASSERT_NE(layout, nullptr);
    EXPECT_NE(layout->GetNativeHandle(), nullptr);
    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_Create_Complex) {
    RHIPipelineLayoutDesc desc = {};
    desc.constantBufferCount = 2;
    desc.textureCount = 4;
    desc.samplerCount = 2;
    desc.uavCount = 1;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(desc);
    ASSERT_NE(layout, nullptr);
    EXPECT_NE(layout->GetNativeHandle(), nullptr);
    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_Create_ZeroCounts) {
    RHIPipelineLayoutDesc desc = {};
    desc.constantBufferCount = 0;
    desc.textureCount = 0;
    desc.samplerCount = 0;
    desc.uavCount = 0;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(desc);
    ASSERT_NE(layout, nullptr);
    EXPECT_NE(layout->GetNativeHandle(), nullptr);
    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_Shutdown) {
    RHIPipelineLayoutDesc desc = {};
    desc.constantBufferCount = 1;
    desc.textureCount = 2;
    desc.samplerCount = 1;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(desc);
    ASSERT_NE(layout, nullptr);
    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_DoubleShutdown) {
    RHIPipelineLayoutDesc desc = {};
    desc.constantBufferCount = 1;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(desc);
    ASSERT_NE(layout, nullptr);
    layout->Shutdown();
    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_DescriptorSetAllocation) {
    RHIPipelineLayoutDesc layoutDesc = {};
    layoutDesc.constantBufferCount = 1;
    layoutDesc.textureCount = 2;

    RHIPipelineLayout* layout = GetDevice()->CreatePipelineLayout(layoutDesc);
    if (layout == nullptr) {
        return;
    }

    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = true;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool != nullptr) {
        DescriptorSetLayoutDesc setDesc = {};
        setDesc.bindingCount = 0;

        RHIDescriptorSet* set = pool->AllocateSet(setDesc);
        if (set != nullptr) {
            EXPECT_GE(set->GetBindingCount(), 0u);
            set->Shutdown();
            delete set;
        }
        pool->Shutdown();
        delete pool;
    }

    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_DeepCopiesSetLayoutsAndInfersCounts) {
    DescriptorSetLayoutBinding set0Bindings[1] = {};
    set0Bindings[0].binding = 0;
    set0Bindings[0].type = static_cast<uint32_t>(DescriptorType::CBV);
    set0Bindings[0].count = 1;

    DescriptorSetLayoutBinding set1Bindings[2] = {};
    set1Bindings[0].binding = 1;
    set1Bindings[0].type = static_cast<uint32_t>(DescriptorType::SRV);
    set1Bindings[0].count = 2;
    set1Bindings[1].binding = 3;
    set1Bindings[1].type = static_cast<uint32_t>(DescriptorType::Sampler);
    set1Bindings[1].count = 1;

    DescriptorSetLayoutDesc setLayouts[2] = {};
    setLayouts[0].bindings = set0Bindings;
    setLayouts[0].bindingCount = 1;
    setLayouts[1].bindings = set1Bindings;
    setLayouts[1].bindingCount = 2;

    RHIPipelineLayoutDesc desc = {};
    desc.setLayouts = setLayouts;
    desc.setLayoutCount = 2;

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

    set0Bindings[0].binding = 99;
    set1Bindings[0].count = 7;
    setLayouts[1].bindingCount = 0;

    const RHIPipelineLayoutDesc* storedDesc = nullptr;
    if (GetBackendType() == RHIType::D3D12) {
        storedDesc = &static_cast<D3D12PipelineLayout*>(layout)->GetDesc();
    } else {
        storedDesc = &static_cast<OpenGLPipelineLayout*>(layout)->GetDesc();
    }

    ASSERT_NE(storedDesc, nullptr);
    ASSERT_EQ(storedDesc->setLayoutCount, 2u);
    ASSERT_NE(storedDesc->setLayouts, nullptr);
    ASSERT_EQ(storedDesc->setLayouts[0].bindingCount, 1u);
    ASSERT_EQ(storedDesc->setLayouts[1].bindingCount, 2u);
    ASSERT_NE(storedDesc->setLayouts[0].bindings, nullptr);
    ASSERT_NE(storedDesc->setLayouts[1].bindings, nullptr);
    EXPECT_EQ(storedDesc->setLayouts[0].bindings[0].binding, 0u);
    EXPECT_EQ(storedDesc->setLayouts[1].bindings[0].count, 2u);
    EXPECT_EQ(storedDesc->constantBufferCount, 1u);
    EXPECT_EQ(storedDesc->textureCount, 2u);
    EXPECT_EQ(storedDesc->samplerCount, 1u);
    EXPECT_EQ(storedDesc->uavCount, 0u);

    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_D3D12TracksDistinctBindingClasses) {
    if (GetBackendType() != RHIType::D3D12) {
        GTEST_SKIP() << "D3D12-specific root parameter verification";
    }

    RHIPipelineLayoutDesc desc = {};
    desc.constantBufferCount = 2;
    desc.textureCount = 1;
    desc.uavCount = 1;
    desc.samplerCount = 1;

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

    auto* d3d12Layout = static_cast<D3D12PipelineLayout*>(layout);
    EXPECT_TRUE(d3d12Layout->HasConstantBufferBinding(0));
    EXPECT_TRUE(d3d12Layout->HasConstantBufferBinding(1));
    EXPECT_FALSE(d3d12Layout->HasConstantBufferBinding(2));

    EXPECT_TRUE(d3d12Layout->HasShaderResourceTable());
    EXPECT_TRUE(d3d12Layout->HasUnorderedAccessTable());
    EXPECT_TRUE(d3d12Layout->HasSamplerTable());

    EXPECT_NE(
        d3d12Layout->GetConstantBufferRootParameterIndex(0),
        d3d12Layout->GetConstantBufferRootParameterIndex(1));
    EXPECT_NE(
        d3d12Layout->GetShaderResourceTableRootParameterIndex(),
        d3d12Layout->GetUnorderedAccessTableRootParameterIndex());
    EXPECT_NE(
        d3d12Layout->GetUnorderedAccessTableRootParameterIndex(),
        d3d12Layout->GetSamplerTableRootParameterIndex());

    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_D3D12InfersBindingClassesFromSetLayouts) {
    if (GetBackendType() != RHIType::D3D12) {
        GTEST_SKIP() << "D3D12-specific root parameter verification";
    }

    DescriptorSetLayoutBinding set0Bindings[1] = {};
    set0Bindings[0].binding = 0;
    set0Bindings[0].type = static_cast<uint32_t>(DescriptorType::CBV);
    set0Bindings[0].count = 1;

    DescriptorSetLayoutBinding set1Bindings[1] = {};
    set1Bindings[0].binding = 0;
    set1Bindings[0].type = static_cast<uint32_t>(DescriptorType::SRV);
    set1Bindings[0].count = 1;

    DescriptorSetLayoutBinding set2Bindings[1] = {};
    set2Bindings[0].binding = 0;
    set2Bindings[0].type = static_cast<uint32_t>(DescriptorType::UAV);
    set2Bindings[0].count = 1;

    DescriptorSetLayoutBinding set3Bindings[1] = {};
    set3Bindings[0].binding = 0;
    set3Bindings[0].type = static_cast<uint32_t>(DescriptorType::Sampler);
    set3Bindings[0].count = 1;

    DescriptorSetLayoutDesc setLayouts[4] = {};
    setLayouts[0].bindings = set0Bindings;
    setLayouts[0].bindingCount = 1;
    setLayouts[1].bindings = set1Bindings;
    setLayouts[1].bindingCount = 1;
    setLayouts[2].bindings = set2Bindings;
    setLayouts[2].bindingCount = 1;
    setLayouts[3].bindings = set3Bindings;
    setLayouts[3].bindingCount = 1;

    RHIPipelineLayoutDesc desc = {};
    desc.setLayouts = setLayouts;
    desc.setLayoutCount = 4;

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

    auto* d3d12Layout = static_cast<D3D12PipelineLayout*>(layout);
    EXPECT_TRUE(d3d12Layout->UsesSetLayouts());
    EXPECT_EQ(d3d12Layout->GetSetLayoutCount(), 4u);
    EXPECT_TRUE(d3d12Layout->HasConstantBufferBinding(0, 0));
    EXPECT_TRUE(d3d12Layout->HasShaderResourceTable(1));
    EXPECT_TRUE(d3d12Layout->HasUnorderedAccessTable(2));
    EXPECT_TRUE(d3d12Layout->HasSamplerTable(3));
    EXPECT_EQ(d3d12Layout->GetDesc().constantBufferCount, 1u);
    EXPECT_EQ(d3d12Layout->GetDesc().textureCount, 1u);
    EXPECT_EQ(d3d12Layout->GetDesc().uavCount, 1u);
    EXPECT_EQ(d3d12Layout->GetDesc().samplerCount, 1u);

    layout->Shutdown();
    delete layout;
}

TEST_P(RHITestFixture, PipelineLayout_D3D12SeparatesOverlappingBindingsAcrossSetSlots) {
    if (GetBackendType() != RHIType::D3D12) {
        GTEST_SKIP() << "D3D12-specific root parameter verification";
    }

    DescriptorSetLayoutBinding set0Bindings[2] = {};
    set0Bindings[0].binding = 0;
    set0Bindings[0].type = static_cast<uint32_t>(DescriptorType::CBV);
    set0Bindings[0].count = 1;
    set0Bindings[1].binding = 0;
    set0Bindings[1].type = static_cast<uint32_t>(DescriptorType::SRV);
    set0Bindings[1].count = 1;

    DescriptorSetLayoutBinding set1Bindings[2] = {};
    set1Bindings[0].binding = 0;
    set1Bindings[0].type = static_cast<uint32_t>(DescriptorType::CBV);
    set1Bindings[0].count = 1;
    set1Bindings[1].binding = 0;
    set1Bindings[1].type = static_cast<uint32_t>(DescriptorType::SRV);
    set1Bindings[1].count = 1;

    DescriptorSetLayoutBinding set2Bindings[1] = {};
    set2Bindings[0].binding = 0;
    set2Bindings[0].type = static_cast<uint32_t>(DescriptorType::Sampler);
    set2Bindings[0].count = 1;

    DescriptorSetLayoutDesc setLayouts[3] = {};
    setLayouts[0].bindings = set0Bindings;
    setLayouts[0].bindingCount = 2;
    setLayouts[1].bindings = set1Bindings;
    setLayouts[1].bindingCount = 2;
    setLayouts[2].bindings = set2Bindings;
    setLayouts[2].bindingCount = 1;

    RHIPipelineLayoutDesc desc = {};
    desc.setLayouts = setLayouts;
    desc.setLayoutCount = 3;

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

    auto* d3d12Layout = static_cast<D3D12PipelineLayout*>(layout);
    EXPECT_TRUE(d3d12Layout->UsesSetLayouts());
    EXPECT_TRUE(d3d12Layout->HasConstantBufferBinding(0, 0));
    EXPECT_TRUE(d3d12Layout->HasConstantBufferBinding(1, 0));
    EXPECT_FALSE(d3d12Layout->HasConstantBufferBinding(2, 0));
    EXPECT_NE(
        d3d12Layout->GetConstantBufferRootParameterIndex(0, 0),
        d3d12Layout->GetConstantBufferRootParameterIndex(1, 0));

    EXPECT_TRUE(d3d12Layout->HasShaderResourceTable(0));
    EXPECT_TRUE(d3d12Layout->HasShaderResourceTable(1));
    EXPECT_FALSE(d3d12Layout->HasShaderResourceTable(2));
    EXPECT_NE(
        d3d12Layout->GetShaderResourceTableRootParameterIndex(0),
        d3d12Layout->GetShaderResourceTableRootParameterIndex(1));

    EXPECT_FALSE(d3d12Layout->HasSamplerTable(0));
    EXPECT_FALSE(d3d12Layout->HasSamplerTable(1));
    EXPECT_TRUE(d3d12Layout->HasSamplerTable(2));

    layout->Shutdown();
    delete layout;
}