#include "fixtures/D3D12TestFixture.h"

#include "XCEngine/RHI/D3D12/D3D12DescriptorSet.h"
#include "XCEngine/RHI/D3D12/D3D12PipelineLayout.h"
#include "XCEngine/RHI/RHIDescriptorPool.h"
#include "XCEngine/RHI/RHIEnums.h"
#include "XCEngine/RHI/RHIPipelineLayout.h"

using namespace XCEngine::RHI;

TEST_F(D3D12TestFixture, DescriptorSet_MixedBindings_AssignDescriptorIndicesByType) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 4;
    poolDesc.shaderVisible = true;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    ASSERT_NE(pool, nullptr);

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

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindings = bindings;
    layoutDesc.bindingCount = 3;

    RHIDescriptorSet* firstSet = pool->AllocateSet(layoutDesc);
    RHIDescriptorSet* secondSet = pool->AllocateSet(layoutDesc);
    ASSERT_NE(firstSet, nullptr);
    ASSERT_NE(secondSet, nullptr);

    auto* firstD3D12Set = static_cast<D3D12DescriptorSet*>(firstSet);
    auto* secondD3D12Set = static_cast<D3D12DescriptorSet*>(secondSet);
    EXPECT_EQ(firstD3D12Set->GetCount(), 2u);
    EXPECT_EQ(firstD3D12Set->GetDescriptorIndexForBinding(1), 0u);
    EXPECT_EQ(firstD3D12Set->GetDescriptorIndexForBinding(2), 1u);
    EXPECT_EQ(firstD3D12Set->GetDescriptorIndexForBinding(0), UINT32_MAX);
    EXPECT_EQ(firstD3D12Set->GetOffset(), 0u);
    EXPECT_EQ(secondD3D12Set->GetOffset(), 2u);

    firstSet->Shutdown();
    delete firstSet;
    secondSet->Shutdown();
    delete secondSet;
    pool->Shutdown();
    delete pool;
}

TEST_F(D3D12TestFixture, DescriptorSet_MultipleConstantBuffersUploadIndependently) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 1;
    poolDesc.shaderVisible = false;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    ASSERT_NE(pool, nullptr);

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

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindings = bindings;
    layoutDesc.bindingCount = 2;

    RHIDescriptorSet* set = pool->AllocateSet(layoutDesc);
    ASSERT_NE(set, nullptr);

    auto* d3d12Set = static_cast<D3D12DescriptorSet*>(set);
    const float firstData[16] = { 1.0f, 0.0f, 0.0f, 0.0f };
    const float secondData[16] = { 2.0f, 0.0f, 0.0f, 0.0f };
    d3d12Set->WriteConstant(0, firstData, sizeof(firstData));
    d3d12Set->WriteConstant(1, secondData, sizeof(secondData));

    ASSERT_TRUE(d3d12Set->UploadConstantBuffer(0));
    ASSERT_TRUE(d3d12Set->UploadConstantBuffer(1));

    const D3D12_GPU_VIRTUAL_ADDRESS firstAddress = d3d12Set->GetConstantBufferGPUAddress(0);
    const D3D12_GPU_VIRTUAL_ADDRESS secondAddress = d3d12Set->GetConstantBufferGPUAddress(1);
    EXPECT_NE(firstAddress, 0u);
    EXPECT_NE(secondAddress, 0u);
    EXPECT_NE(firstAddress, secondAddress);

    set->Shutdown();
    delete set;
    pool->Shutdown();
    delete pool;
}

TEST_F(D3D12TestFixture, PipelineLayout_TracksDistinctBindingClasses) {
    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_F(D3D12TestFixture, PipelineLayout_InfersBindingClassesFromSetLayouts) {
    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_F(D3D12TestFixture, PipelineLayout_SeparatesOverlappingBindingsAcrossSetSlots) {
    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;
}