#include "fixtures/RHITestFixture.h"
#include "XCEngine/RHI/D3D12/D3D12DescriptorSet.h"
#include "XCEngine/RHI/OpenGL/OpenGLDescriptorSet.h"
#include "XCEngine/RHI/OpenGL/OpenGLDevice.h"
#include "XCEngine/RHI/RHIDescriptorPool.h"
#include "XCEngine/RHI/RHIDescriptorSet.h"
#include "XCEngine/RHI/RHISampler.h"
#include "XCEngine/RHI/RHIResourceView.h"
#include "XCEngine/RHI/RHITexture.h"

#include <glad/glad.h>

using namespace XCEngine::RHI;

TEST_P(RHITestFixture, DescriptorSet_Allocate_Basic) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = true;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindings = nullptr;
    layoutDesc.bindingCount = 0;

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

    pool->Shutdown();
    delete pool;
}

TEST_P(RHITestFixture, DescriptorSet_Update_ResourceView) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = true;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    TextureDesc texDesc = {};
    texDesc.width = 256;
    texDesc.height = 256;
    texDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    texDesc.textureType = static_cast<uint32_t>(TextureType::Texture2D);
    texDesc.sampleCount = 1;

    RHITexture* texture = GetDevice()->CreateTexture(texDesc);
    if (texture == nullptr) {
        pool->Shutdown();
        delete pool;
        return;
    }

    RHIResourceView* srv = GetDevice()->CreateShaderResourceView(texture, {});
    if (srv == nullptr) {
        texture->Shutdown();
        delete texture;
        pool->Shutdown();
        delete pool;
        return;
    }

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindingCount = 1;
    DescriptorSetLayoutBinding binding = {};
    binding.binding = 0;
    binding.type = static_cast<uint32_t>(DescriptorType::SRV);
    binding.count = 1;
    layoutDesc.bindings = &binding;

    RHIDescriptorSet* set = pool->AllocateSet(layoutDesc);
    if (set != nullptr) {
        set->Update(0, srv);
        set->Shutdown();
        delete set;
    }

    srv->Shutdown();
    delete srv;
    texture->Shutdown();
    delete texture;
    pool->Shutdown();
    delete pool;
}

TEST_P(RHITestFixture, DescriptorSet_UpdateSampler) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::Sampler;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = true;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    SamplerDesc samplerDesc = {};
    samplerDesc.filter = static_cast<uint32_t>(FilterMode::Linear);
    samplerDesc.addressU = static_cast<uint32_t>(TextureAddressMode::Wrap);
    samplerDesc.addressV = static_cast<uint32_t>(TextureAddressMode::Wrap);
    samplerDesc.addressW = static_cast<uint32_t>(TextureAddressMode::Wrap);

    RHISampler* sampler = GetDevice()->CreateSampler(samplerDesc);
    if (sampler == nullptr) {
        pool->Shutdown();
        delete pool;
        return;
    }

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindingCount = 1;
    DescriptorSetLayoutBinding binding = {};
    binding.binding = 0;
    binding.type = static_cast<uint32_t>(DescriptorType::Sampler);
    binding.count = 1;
    layoutDesc.bindings = &binding;

    RHIDescriptorSet* set = pool->AllocateSet(layoutDesc);
    if (set != nullptr) {
        set->UpdateSampler(0, sampler);
        set->Shutdown();
        delete set;
    }

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

TEST_P(RHITestFixture, DescriptorSet_WriteConstant) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = false;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindingCount = 1;
    DescriptorSetLayoutBinding binding = {};
    binding.binding = 0;
    binding.type = static_cast<uint32_t>(DescriptorType::CBV);
    binding.count = 1;
    layoutDesc.bindings = &binding;

    RHIDescriptorSet* set = pool->AllocateSet(layoutDesc);
    if (set != nullptr) {
        float testData[16] = { 1.0f, 2.0f, 3.0f, 4.0f };
        set->WriteConstant(0, testData, sizeof(testData));

        void* constantData = set->GetConstantBufferData();
        EXPECT_NE(constantData, nullptr);

        size_t constantSize = set->GetConstantBufferSize();
        EXPECT_GE(constantSize, sizeof(testData));

        set->Shutdown();
        delete set;
    }

    pool->Shutdown();
    delete pool;
}

TEST_P(RHITestFixture, DescriptorSet_WriteConstant_PartialUpdate) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = false;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindingCount = 1;
    DescriptorSetLayoutBinding binding = {};
    binding.binding = 0;
    binding.type = static_cast<uint32_t>(DescriptorType::CBV);
    binding.count = 1;
    layoutDesc.bindings = &binding;

    RHIDescriptorSet* set = pool->AllocateSet(layoutDesc);
    if (set != nullptr) {
        float offsetData[4] = { 100.0f, 200.0f };
        set->WriteConstant(0, offsetData, sizeof(offsetData), 0);

        EXPECT_TRUE(set->IsConstantDirty());

        set->MarkConstantClean();
        EXPECT_FALSE(set->IsConstantDirty());

        set->Shutdown();
        delete set;
    }

    pool->Shutdown();
    delete pool;
}

TEST_P(RHITestFixture, DescriptorSet_Bind_Unbind) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = true;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindingCount = 0;

    RHIDescriptorSet* set = pool->AllocateSet(layoutDesc);
    if (set != nullptr) {
        set->Bind();
        set->Unbind();
        set->Shutdown();
        delete set;
    }

    pool->Shutdown();
    delete pool;
}

TEST_P(RHITestFixture, DescriptorSet_GetBindings) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = true;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindingCount = 2;
    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::SRV);
    bindings[1].count = 2;
    layoutDesc.bindings = bindings;

    RHIDescriptorSet* set = pool->AllocateSet(layoutDesc);
    if (set != nullptr) {
        EXPECT_EQ(set->GetBindingCount(), 2u);
        const DescriptorSetLayoutBinding* retrievedBindings = set->GetBindings();
        ASSERT_NE(retrievedBindings, nullptr);
        EXPECT_EQ(retrievedBindings[0].binding, 0u);
        EXPECT_EQ(retrievedBindings[1].binding, 1u);
        set->Shutdown();
        delete set;
    }

    pool->Shutdown();
    delete pool;
}

TEST_P(RHITestFixture, DescriptorSet_MultipleAllocations) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 100;
    poolDesc.shaderVisible = true;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    RHIDescriptorSet* sets[10] = { nullptr };

    for (int i = 0; i < 10; ++i) {
        DescriptorSetLayoutDesc layoutDesc = {};
        layoutDesc.bindingCount = 0;
        sets[i] = pool->AllocateSet(layoutDesc);
    }

    for (int i = 0; i < 10; ++i) {
        if (sets[i] != nullptr) {
            sets[i]->Shutdown();
            delete sets[i];
        }
    }

    pool->Shutdown();
    delete pool;
}

TEST_P(RHITestFixture, DescriptorSet_FreeSet_ViaPool) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = true;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindingCount = 0;

    RHIDescriptorSet* set = pool->AllocateSet(layoutDesc);
    if (set != nullptr) {
        pool->FreeSet(set);
        set = nullptr;
    }

    pool->Shutdown();
    delete pool;
}

TEST_P(RHITestFixture, DescriptorSet_ConstantBufferSize) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 10;
    poolDesc.shaderVisible = false;

    RHIDescriptorPool* pool = GetDevice()->CreateDescriptorPool(poolDesc);
    if (pool == nullptr) {
        return;
    }

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindingCount = 1;
    DescriptorSetLayoutBinding binding = {};
    binding.binding = 0;
    binding.type = static_cast<uint32_t>(DescriptorType::CBV);
    binding.count = 1;
    layoutDesc.bindings = &binding;

    RHIDescriptorSet* set = pool->AllocateSet(layoutDesc);
    if (set != nullptr) {
        float data[64] = {};
        set->WriteConstant(0, data, sizeof(data));

        EXPECT_GE(set->GetConstantBufferSize(), sizeof(data));

        set->Shutdown();
        delete set;
    }

    pool->Shutdown();
    delete pool;
}

TEST_P(RHITestFixture, DescriptorSet_MultipleAllocations_AdvanceDescriptorOffsets) {
    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 8;
    poolDesc.shaderVisible = true;

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

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

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

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

    if (GetBackendType() == RHIType::D3D12) {
        auto* firstD3D12Set = static_cast<D3D12DescriptorSet*>(firstSet);
        auto* secondD3D12Set = static_cast<D3D12DescriptorSet*>(secondSet);
        EXPECT_EQ(firstD3D12Set->GetOffset(), 0u);
        EXPECT_EQ(secondD3D12Set->GetOffset(), 2u);
    }

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

TEST_P(RHITestFixture, DescriptorSet_Update_UsesBindingNumberOnOpenGL) {
    if (GetBackendType() != RHIType::OpenGL) {
        GTEST_SKIP() << "OpenGL-specific descriptor binding verification";
    }

    auto* openGLDevice = static_cast<OpenGLDevice*>(GetDevice());
    ASSERT_NE(openGLDevice, nullptr);
    ASSERT_TRUE(openGLDevice->MakeContextCurrent());

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

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

    TextureDesc textureDesc = {};
    textureDesc.width = 1;
    textureDesc.height = 1;
    textureDesc.depth = 1;
    textureDesc.mipLevels = 1;
    textureDesc.arraySize = 1;
    textureDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    textureDesc.textureType = static_cast<uint32_t>(TextureType::Texture2D);
    textureDesc.sampleCount = 1;

    const uint8_t pixel[4] = { 255, 255, 255, 255 };
    RHITexture* texture = GetDevice()->CreateTexture(textureDesc, pixel, sizeof(pixel), 4);
    ASSERT_NE(texture, nullptr);

    ResourceViewDesc srvDesc = {};
    srvDesc.format = textureDesc.format;
    srvDesc.dimension = ResourceViewDimension::Texture2D;
    RHIResourceView* srv = GetDevice()->CreateShaderResourceView(texture, srvDesc);
    ASSERT_NE(srv, nullptr);

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

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

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

    auto* openGLSet = static_cast<OpenGLDescriptorSet*>(set);
    const uint32_t untouchedUnit = openGLSet->GetBindingPoint(3);
    const uint32_t updatedUnit = openGLSet->GetBindingPoint(7);
    ASSERT_NE(untouchedUnit, updatedUnit);

    set->Update(7, srv);
    set->Bind();

    GLint updatedTexture = 0;
    glActiveTexture(GL_TEXTURE0 + updatedUnit);
    glGetIntegerv(GL_TEXTURE_BINDING_2D, &updatedTexture);
    EXPECT_EQ(static_cast<unsigned int>(updatedTexture), static_cast<unsigned int>(reinterpret_cast<uintptr_t>(srv->GetNativeHandle())));

    GLint untouchedTexture = 0;
    glActiveTexture(GL_TEXTURE0 + untouchedUnit);
    glGetIntegerv(GL_TEXTURE_BINDING_2D, &untouchedTexture);
    EXPECT_EQ(untouchedTexture, 0);

    set->Unbind();

    GLint unboundTexture = 0;
    glActiveTexture(GL_TEXTURE0 + updatedUnit);
    glGetIntegerv(GL_TEXTURE_BINDING_2D, &unboundTexture);
    EXPECT_EQ(unboundTexture, 0);

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

TEST_P(RHITestFixture, DescriptorSet_BindConstantBuffer_UsesBindingNumberOnOpenGL) {
    if (GetBackendType() != RHIType::OpenGL) {
        GTEST_SKIP() << "OpenGL-specific constant buffer binding verification";
    }

    auto* openGLDevice = static_cast<OpenGLDevice*>(GetDevice());
    ASSERT_NE(openGLDevice, nullptr);
    ASSERT_TRUE(openGLDevice->MakeContextCurrent());

    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 1;
    poolDesc.shaderVisible = false;

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

    DescriptorSetLayoutBinding binding = {};
    binding.binding = 3;
    binding.type = static_cast<uint32_t>(DescriptorType::CBV);
    binding.count = 1;

    DescriptorSetLayoutDesc layoutDesc = {};
    layoutDesc.bindings = &binding;
    layoutDesc.bindingCount = 1;

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

    const float matrixData[16] = {
        1.0f, 0.0f, 0.0f, 0.0f,
        0.0f, 1.0f, 0.0f, 0.0f,
        0.0f, 0.0f, 1.0f, 0.0f,
        0.0f, 0.0f, 0.0f, 1.0f,
    };
    set->WriteConstant(3, matrixData, sizeof(matrixData));
    set->Bind();

    GLint boundBuffer = 0;
    glGetIntegeri_v(GL_UNIFORM_BUFFER_BINDING, 3, &boundBuffer);
    EXPECT_NE(boundBuffer, 0);

    set->Unbind();

    GLint unboundBuffer = -1;
    glGetIntegeri_v(GL_UNIFORM_BUFFER_BINDING, 3, &unboundBuffer);
    EXPECT_EQ(unboundBuffer, 0);

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