#include "fixtures/RHITestFixture.h"
#include "XCEngine/RHI/D3D12/D3D12PipelineLayout.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_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;
}