#if defined(XCENGINE_SUPPORT_VULKAN)

#include "fixtures/VulkanTestFixture.h"
#include "XCEngine/RHI/RHIBuffer.h"
#include "XCEngine/RHI/RHIDescriptorPool.h"
#include "XCEngine/RHI/RHIDescriptorSet.h"
#include "XCEngine/RHI/RHIPipelineLayout.h"
#include "XCEngine/RHI/RHIPipelineState.h"
#include "XCEngine/RHI/RHIResourceView.h"
#include "XCEngine/RHI/RHISampler.h"
#include "XCEngine/RHI/Vulkan/VulkanTexture.h"

#include <cstdint>
#include <cstring>
#include <vector>

using namespace XCEngine::RHI;

namespace {

struct Vertex {
    float pos[4];
    float uv[2];
};

constexpr Vertex kQuadVertices[] = {
    {{-0.5f, -0.5f, 0.0f, 1.0f}, {0.0f, 1.0f}},
    {{-0.5f,  0.5f, 0.0f, 1.0f}, {0.0f, 0.0f}},
    {{ 0.5f, -0.5f, 0.0f, 1.0f}, {1.0f, 1.0f}},
    {{ 0.5f,  0.5f, 0.0f, 1.0f}, {1.0f, 0.0f}},
};

constexpr uint32_t kQuadIndices[] = {0, 1, 2, 2, 1, 3};

const char kVertexShader[] = R"(#version 450
layout(location = 0) in vec4 aPosition;
layout(location = 1) in vec2 aTexCoord;

layout(location = 0) out vec2 vTexCoord;

void main() {
    gl_Position = aPosition;
    vTexCoord = aTexCoord;
}
)";

const char kFragmentShader[] = R"(#version 450
layout(set = 1, binding = 0) uniform texture2D uTexture;
layout(set = 2, binding = 0) uniform sampler uSampler;

layout(location = 0) in vec2 vTexCoord;
layout(location = 0) out vec4 fragColor;

void main() {
    fragColor = texture(sampler2D(uTexture, uSampler), vTexCoord);
}
)";

GraphicsPipelineDesc CreatePipelineDesc(RHIPipelineLayout* pipelineLayout) {
    GraphicsPipelineDesc desc = {};
    desc.pipelineLayout = pipelineLayout;
    desc.topologyType = static_cast<uint32_t>(PrimitiveTopologyType::Triangle);
    desc.renderTargetFormats[0] = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    desc.depthStencilFormat = static_cast<uint32_t>(Format::Unknown);
    desc.sampleCount = 1;

    desc.rasterizerState.fillMode = static_cast<uint32_t>(FillMode::Solid);
    desc.rasterizerState.cullMode = static_cast<uint32_t>(CullMode::None);
    desc.rasterizerState.frontFace = static_cast<uint32_t>(FrontFace::CounterClockwise);
    desc.rasterizerState.depthClipEnable = true;

    desc.depthStencilState.depthTestEnable = false;
    desc.depthStencilState.depthWriteEnable = false;
    desc.depthStencilState.stencilEnable = false;

    InputElementDesc position = {};
    position.semanticName = "POSITION";
    position.semanticIndex = 0;
    position.format = static_cast<uint32_t>(Format::R32G32B32A32_Float);
    position.inputSlot = 0;
    position.alignedByteOffset = 0;
    desc.inputLayout.elements.push_back(position);

    InputElementDesc texcoord = {};
    texcoord.semanticName = "TEXCOORD";
    texcoord.semanticIndex = 0;
    texcoord.format = static_cast<uint32_t>(Format::R32G32_Float);
    texcoord.inputSlot = 0;
    texcoord.alignedByteOffset = sizeof(float) * 4;
    desc.inputLayout.elements.push_back(texcoord);

    desc.vertexShader.source.assign(kVertexShader, kVertexShader + std::strlen(kVertexShader));
    desc.vertexShader.sourceLanguage = ShaderLanguage::GLSL;
    desc.vertexShader.profile = L"vs";

    desc.fragmentShader.source.assign(kFragmentShader, kFragmentShader + std::strlen(kFragmentShader));
    desc.fragmentShader.sourceLanguage = ShaderLanguage::GLSL;
    desc.fragmentShader.profile = L"fs";
    return desc;
}

TEST_F(VulkanGraphicsFixture, DrawIndexedTexturedQuadWithFirstSetOffsetWritesExpectedPixels) {
    constexpr uint32_t kWidth = 64;
    constexpr uint32_t kHeight = 64;

    TextureDesc colorDesc = CreateColorTextureDesc(kWidth, kHeight);
    RHITexture* colorTexture = m_device->CreateTexture(colorDesc);
    ASSERT_NE(colorTexture, nullptr);

    RHIResourceView* renderTargetView = m_device->CreateRenderTargetView(colorTexture, {});
    ASSERT_NE(renderTargetView, nullptr);

    BufferDesc vertexBufferDesc = {};
    vertexBufferDesc.size = sizeof(kQuadVertices);
    vertexBufferDesc.stride = sizeof(Vertex);
    vertexBufferDesc.bufferType = static_cast<uint32_t>(BufferType::Vertex);
    RHIBuffer* vertexBuffer = m_device->CreateBuffer(vertexBufferDesc);
    ASSERT_NE(vertexBuffer, nullptr);
    vertexBuffer->SetData(kQuadVertices, sizeof(kQuadVertices));
    vertexBuffer->SetStride(sizeof(Vertex));
    vertexBuffer->SetBufferType(BufferType::Vertex);

    ResourceViewDesc vertexViewDesc = {};
    vertexViewDesc.dimension = ResourceViewDimension::Buffer;
    vertexViewDesc.structureByteStride = sizeof(Vertex);
    RHIResourceView* vertexView = m_device->CreateVertexBufferView(vertexBuffer, vertexViewDesc);
    ASSERT_NE(vertexView, nullptr);

    BufferDesc indexBufferDesc = {};
    indexBufferDesc.size = sizeof(kQuadIndices);
    indexBufferDesc.stride = sizeof(uint32_t);
    indexBufferDesc.bufferType = static_cast<uint32_t>(BufferType::Index);
    RHIBuffer* indexBuffer = m_device->CreateBuffer(indexBufferDesc);
    ASSERT_NE(indexBuffer, nullptr);
    indexBuffer->SetData(kQuadIndices, sizeof(kQuadIndices));
    indexBuffer->SetStride(sizeof(uint32_t));
    indexBuffer->SetBufferType(BufferType::Index);

    ResourceViewDesc indexViewDesc = {};
    indexViewDesc.dimension = ResourceViewDimension::Buffer;
    indexViewDesc.format = static_cast<uint32_t>(Format::R32_UInt);
    RHIResourceView* indexView = m_device->CreateIndexBufferView(indexBuffer, indexViewDesc);
    ASSERT_NE(indexView, nullptr);

    const uint8_t texturePixel[] = {255, 0, 0, 255};
    TextureDesc textureDesc = CreateColorTextureDesc(1, 1);
    RHITexture* texture = m_device->CreateTexture(textureDesc, texturePixel, sizeof(texturePixel), 4);
    ASSERT_NE(texture, nullptr);

    RHIResourceView* shaderResourceView = m_device->CreateShaderResourceView(texture, {});
    ASSERT_NE(shaderResourceView, nullptr);

    SamplerDesc samplerDesc = {};
    samplerDesc.filter = static_cast<uint32_t>(FilterMode::Linear);
    samplerDesc.addressU = static_cast<uint32_t>(TextureAddressMode::Clamp);
    samplerDesc.addressV = static_cast<uint32_t>(TextureAddressMode::Clamp);
    samplerDesc.addressW = static_cast<uint32_t>(TextureAddressMode::Clamp);
    samplerDesc.comparisonFunc = static_cast<uint32_t>(ComparisonFunc::Always);
    samplerDesc.maxAnisotropy = 1;
    samplerDesc.minLod = 0.0f;
    samplerDesc.maxLod = 1000.0f;
    RHISampler* sampler = m_device->CreateSampler(samplerDesc);
    ASSERT_NE(sampler, nullptr);

    DescriptorPoolDesc texturePoolDesc = {};
    texturePoolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    texturePoolDesc.descriptorCount = 1;
    texturePoolDesc.shaderVisible = true;
    RHIDescriptorPool* texturePool = m_device->CreateDescriptorPool(texturePoolDesc);
    ASSERT_NE(texturePool, nullptr);

    DescriptorSetLayoutBinding textureBinding = {};
    textureBinding.binding = 0;
    textureBinding.type = static_cast<uint32_t>(DescriptorType::SRV);
    textureBinding.count = 1;
    textureBinding.visibility = static_cast<uint32_t>(ShaderVisibility::Pixel);

    DescriptorSetLayoutDesc textureLayoutDesc = {};
    textureLayoutDesc.bindings = &textureBinding;
    textureLayoutDesc.bindingCount = 1;

    RHIDescriptorSet* textureSet = texturePool->AllocateSet(textureLayoutDesc);
    ASSERT_NE(textureSet, nullptr);
    textureSet->Update(0, shaderResourceView);

    DescriptorPoolDesc samplerPoolDesc = {};
    samplerPoolDesc.type = DescriptorHeapType::Sampler;
    samplerPoolDesc.descriptorCount = 1;
    samplerPoolDesc.shaderVisible = true;
    RHIDescriptorPool* samplerPool = m_device->CreateDescriptorPool(samplerPoolDesc);
    ASSERT_NE(samplerPool, nullptr);

    DescriptorSetLayoutBinding samplerBinding = {};
    samplerBinding.binding = 0;
    samplerBinding.type = static_cast<uint32_t>(DescriptorType::Sampler);
    samplerBinding.count = 1;
    samplerBinding.visibility = static_cast<uint32_t>(ShaderVisibility::Pixel);

    DescriptorSetLayoutDesc samplerLayoutDesc = {};
    samplerLayoutDesc.bindings = &samplerBinding;
    samplerLayoutDesc.bindingCount = 1;

    RHIDescriptorSet* samplerSet = samplerPool->AllocateSet(samplerLayoutDesc);
    ASSERT_NE(samplerSet, nullptr);
    samplerSet->UpdateSampler(0, sampler);

    DescriptorSetLayoutBinding reservedBinding = {};
    reservedBinding.binding = 0;
    reservedBinding.type = static_cast<uint32_t>(DescriptorType::CBV);
    reservedBinding.count = 1;
    reservedBinding.visibility = static_cast<uint32_t>(ShaderVisibility::Vertex);

    DescriptorSetLayoutDesc reservedLayoutDesc = {};
    reservedLayoutDesc.bindings = &reservedBinding;
    reservedLayoutDesc.bindingCount = 1;

    DescriptorSetLayoutDesc setLayouts[3] = {};
    setLayouts[0] = reservedLayoutDesc;
    setLayouts[1] = textureLayoutDesc;
    setLayouts[2] = samplerLayoutDesc;

    RHIPipelineLayoutDesc pipelineLayoutDesc = {};
    pipelineLayoutDesc.setLayouts = setLayouts;
    pipelineLayoutDesc.setLayoutCount = 3;
    RHIPipelineLayout* pipelineLayout = m_device->CreatePipelineLayout(pipelineLayoutDesc);
    ASSERT_NE(pipelineLayout, nullptr);

    RHIPipelineState* pipelineState = m_device->CreatePipelineState(CreatePipelineDesc(pipelineLayout));
    ASSERT_NE(pipelineState, nullptr);
    ASSERT_TRUE(pipelineState->IsValid());

    RHICommandList* commandList = CreateCommandList();
    ASSERT_NE(commandList, nullptr);

    commandList->Reset();
    commandList->SetRenderTargets(1, &renderTargetView, nullptr);
    commandList->SetViewport({0.0f, 0.0f, static_cast<float>(kWidth), static_cast<float>(kHeight), 0.0f, 1.0f});
    commandList->SetScissorRect({0, 0, static_cast<int32_t>(kWidth), static_cast<int32_t>(kHeight)});
    commandList->Clear(0.0f, 0.0f, 1.0f, 1.0f, 1);
    commandList->SetPipelineState(pipelineState);

    RHIDescriptorSet* descriptorSets[] = {textureSet, samplerSet};
    commandList->SetGraphicsDescriptorSets(1, 2, descriptorSets, pipelineLayout);
    commandList->SetPrimitiveTopology(PrimitiveTopology::TriangleList);

    RHIResourceView* vertexViews[] = {vertexView};
    uint64_t offsets[] = {0};
    uint32_t strides[] = {sizeof(Vertex)};
    commandList->SetVertexBuffers(0, 1, vertexViews, offsets, strides);
    commandList->SetIndexBuffer(indexView, 0);
    commandList->DrawIndexed(static_cast<uint32_t>(sizeof(kQuadIndices) / sizeof(kQuadIndices[0])));
    SubmitAndWait(commandList);

    const std::vector<uint8_t> pixels = ReadTextureRgba8(static_cast<VulkanTexture*>(colorTexture));
    ASSERT_EQ(pixels.size(), static_cast<size_t>(kWidth * kHeight * 4));

    auto PixelAt = [&](uint32_t x, uint32_t y) -> const uint8_t* {
        return pixels.data() + (static_cast<size_t>(y) * kWidth + x) * 4;
    };

    const uint8_t* centerPixel = PixelAt(kWidth / 2, kHeight / 2);
    EXPECT_EQ(centerPixel[0], 255u);
    EXPECT_EQ(centerPixel[1], 0u);
    EXPECT_EQ(centerPixel[2], 0u);
    EXPECT_EQ(centerPixel[3], 255u);

    const uint8_t* cornerPixel = PixelAt(2, 2);
    EXPECT_EQ(cornerPixel[0], 0u);
    EXPECT_EQ(cornerPixel[1], 0u);
    EXPECT_EQ(cornerPixel[2], 255u);
    EXPECT_EQ(cornerPixel[3], 255u);

    commandList->Shutdown();
    delete commandList;
    pipelineState->Shutdown();
    delete pipelineState;
    pipelineLayout->Shutdown();
    delete pipelineLayout;
    textureSet->Shutdown();
    delete textureSet;
    samplerSet->Shutdown();
    delete samplerSet;
    texturePool->Shutdown();
    delete texturePool;
    samplerPool->Shutdown();
    delete samplerPool;
    sampler->Shutdown();
    delete sampler;
    shaderResourceView->Shutdown();
    delete shaderResourceView;
    texture->Shutdown();
    delete texture;
    indexView->Shutdown();
    delete indexView;
    vertexView->Shutdown();
    delete vertexView;
    indexBuffer->Shutdown();
    delete indexBuffer;
    vertexBuffer->Shutdown();
    delete vertexBuffer;
    renderTargetView->Shutdown();
    delete renderTargetView;
    colorTexture->Shutdown();
    delete colorTexture;
}

} // namespace

#endif