#include "fixtures/RHITestFixture.h"
#include "XCEngine/RHI/D3D12/D3D12PipelineLayout.h"
#include "XCEngine/RHI/OpenGL/OpenGLPipelineLayout.h"
#if defined(XCENGINE_SUPPORT_VULKAN)
#include "XCEngine/RHI/Vulkan/VulkanPipelineLayout.h"
#endif
#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();
#if defined(XCENGINE_SUPPORT_VULKAN)
    } else if (GetBackendType() == RHIType::Vulkan) {
        storedDesc = &static_cast<VulkanPipelineLayout*>(layout)->GetDesc();
#endif
    } 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;
}