#include "Rendering/Passes/BuiltinObjectIdPass.h"

#include "Components/GameObject.h"
#include "RHI/RHICommandList.h"
#include "RHI/RHIPipelineLayout.h"
#include "RHI/RHIPipelineState.h"
#include "Resources/Mesh/Mesh.h"

#include <cstddef>
#include <cstring>

namespace XCEngine {
namespace Rendering {
namespace Passes {

namespace {

const char kBuiltinObjectIdHlsl[] = R"(
cbuffer PerObjectConstants : register(b0) {
    float4x4 gProjectionMatrix;
    float4x4 gViewMatrix;
    float4x4 gModelMatrix;
    float4 gObjectIdColor;
};

struct VSInput {
    float3 position : POSITION;
};

struct PSInput {
    float4 position : SV_POSITION;
};

PSInput MainVS(VSInput input) {
    PSInput output;
    float4 positionWS = mul(gModelMatrix, float4(input.position, 1.0));
    float4 positionVS = mul(gViewMatrix, positionWS);
    output.position = mul(gProjectionMatrix, positionVS);
    return output;
}

float4 MainPS(PSInput input) : SV_TARGET {
    return gObjectIdColor;
}
)";

const char kBuiltinObjectIdVertexShader[] = R"(#version 430
layout(location = 0) in vec3 aPosition;

layout(std140, binding = 0) uniform PerObjectConstants {
    mat4 gProjectionMatrix;
    mat4 gViewMatrix;
    mat4 gModelMatrix;
    vec4 gObjectIdColor;
};

void main() {
    vec4 positionWS = gModelMatrix * vec4(aPosition, 1.0);
    vec4 positionVS = gViewMatrix * positionWS;
    gl_Position = gProjectionMatrix * positionVS;
}
)";

const char kBuiltinObjectIdFragmentShader[] = R"(#version 430
layout(std140, binding = 0) uniform PerObjectConstants {
    mat4 gProjectionMatrix;
    mat4 gViewMatrix;
    mat4 gModelMatrix;
    vec4 gObjectIdColor;
};

layout(location = 0) out vec4 fragColor;

void main() {
    fragColor = gObjectIdColor;
}
)";

RHI::InputLayoutDesc BuildInputLayout() {
    RHI::InputLayoutDesc inputLayout = {};

    RHI::InputElementDesc position = {};
    position.semanticName = "POSITION";
    position.semanticIndex = 0;
    position.format = static_cast<uint32_t>(RHI::Format::R32G32B32_Float);
    position.inputSlot = 0;
    position.alignedByteOffset = static_cast<uint32_t>(offsetof(Resources::StaticMeshVertex, position));
    inputLayout.elements.push_back(position);
    return inputLayout;
}

RHI::GraphicsPipelineDesc CreatePipelineDesc(
    RHI::RHIType backendType,
    RHI::RHIPipelineLayout* pipelineLayout) {
    RHI::GraphicsPipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = pipelineLayout;
    pipelineDesc.topologyType = static_cast<uint32_t>(RHI::PrimitiveTopologyType::Triangle);
    pipelineDesc.renderTargetCount = 1;
    pipelineDesc.renderTargetFormats[0] = static_cast<uint32_t>(RHI::Format::R8G8B8A8_UNorm);
    pipelineDesc.depthStencilFormat = static_cast<uint32_t>(RHI::Format::D24_UNorm_S8_UInt);
    pipelineDesc.sampleCount = 1;
    pipelineDesc.inputLayout = BuildInputLayout();

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

    pipelineDesc.blendState.blendEnable = false;
    pipelineDesc.blendState.colorWriteMask = static_cast<uint8_t>(RHI::ColorWriteMask::All);

    pipelineDesc.depthStencilState.depthTestEnable = true;
    pipelineDesc.depthStencilState.depthWriteEnable = false;
    pipelineDesc.depthStencilState.depthFunc = static_cast<uint32_t>(RHI::ComparisonFunc::LessEqual);

    if (backendType == RHI::RHIType::D3D12) {
        pipelineDesc.vertexShader.source.assign(
            kBuiltinObjectIdHlsl,
            kBuiltinObjectIdHlsl + std::strlen(kBuiltinObjectIdHlsl));
        pipelineDesc.vertexShader.sourceLanguage = RHI::ShaderLanguage::HLSL;
        pipelineDesc.vertexShader.entryPoint = L"MainVS";
        pipelineDesc.vertexShader.profile = L"vs_5_0";

        pipelineDesc.fragmentShader.source.assign(
            kBuiltinObjectIdHlsl,
            kBuiltinObjectIdHlsl + std::strlen(kBuiltinObjectIdHlsl));
        pipelineDesc.fragmentShader.sourceLanguage = RHI::ShaderLanguage::HLSL;
        pipelineDesc.fragmentShader.entryPoint = L"MainPS";
        pipelineDesc.fragmentShader.profile = L"ps_5_0";
    } else {
        pipelineDesc.vertexShader.source.assign(
            kBuiltinObjectIdVertexShader,
            kBuiltinObjectIdVertexShader + std::strlen(kBuiltinObjectIdVertexShader));
        pipelineDesc.vertexShader.sourceLanguage = RHI::ShaderLanguage::GLSL;
        pipelineDesc.vertexShader.profile = L"vs_4_30";

        pipelineDesc.fragmentShader.source.assign(
            kBuiltinObjectIdFragmentShader,
            kBuiltinObjectIdFragmentShader + std::strlen(kBuiltinObjectIdFragmentShader));
        pipelineDesc.fragmentShader.sourceLanguage = RHI::ShaderLanguage::GLSL;
        pipelineDesc.fragmentShader.profile = L"fs_4_30";
    }

    return pipelineDesc;
}

} // namespace

BuiltinObjectIdPass::~BuiltinObjectIdPass() {
    Shutdown();
}

bool BuiltinObjectIdPass::Render(
    const RenderContext& context,
    const RenderSurface& surface,
    const RenderSceneData& sceneData) {
    if (!context.IsValid()) {
        return false;
    }

    const std::vector<RHI::RHIResourceView*>& colorAttachments = surface.GetColorAttachments();
    if (colorAttachments.empty() || colorAttachments[0] == nullptr || surface.GetDepthAttachment() == nullptr) {
        return false;
    }

    const Math::RectInt renderArea = surface.GetRenderArea();
    if (renderArea.width <= 0 || renderArea.height <= 0) {
        return false;
    }

    if (!EnsureInitialized(context)) {
        return false;
    }

    RHI::RHICommandList* commandList = context.commandList;
    RHI::RHIResourceView* renderTarget = colorAttachments[0];

    if (surface.IsAutoTransitionEnabled()) {
        commandList->TransitionBarrier(
            renderTarget,
            surface.GetColorStateBefore(),
            RHI::ResourceStates::RenderTarget);
    }

    commandList->SetRenderTargets(1, &renderTarget, surface.GetDepthAttachment());

    const RHI::Viewport viewport = {
        static_cast<float>(renderArea.x),
        static_cast<float>(renderArea.y),
        static_cast<float>(renderArea.width),
        static_cast<float>(renderArea.height),
        0.0f,
        1.0f
    };
    const RHI::Rect scissorRect = {
        renderArea.x,
        renderArea.y,
        renderArea.x + renderArea.width,
        renderArea.y + renderArea.height
    };
    const float clearColor[4] = { 0.0f, 0.0f, 0.0f, 0.0f };
    const RHI::Rect clearRects[] = { scissorRect };

    commandList->SetViewport(viewport);
    commandList->SetScissorRect(scissorRect);
    commandList->ClearRenderTarget(renderTarget, clearColor, 1, clearRects);
    commandList->SetPrimitiveTopology(RHI::PrimitiveTopology::TriangleList);
    commandList->SetPipelineState(m_pipelineState);

    for (const VisibleRenderItem& visibleItem : sceneData.visibleItems) {
        DrawVisibleItem(context, sceneData, visibleItem);
    }

    if (surface.IsAutoTransitionEnabled()) {
        commandList->TransitionBarrier(
            renderTarget,
            RHI::ResourceStates::RenderTarget,
            surface.GetColorStateAfter());
    }

    return true;
}

void BuiltinObjectIdPass::Shutdown() {
    DestroyResources();
}

bool BuiltinObjectIdPass::EnsureInitialized(const RenderContext& context) {
    if (!context.IsValid()) {
        return false;
    }

    if (m_pipelineLayout != nullptr &&
        m_pipelineState != nullptr &&
        m_device == context.device &&
        m_backendType == context.backendType) {
        return true;
    }

    DestroyResources();
    return CreateResources(context);
}

bool BuiltinObjectIdPass::CreateResources(const RenderContext& context) {
    m_device = context.device;
    m_backendType = context.backendType;

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

    RHI::DescriptorSetLayoutDesc constantLayout = {};
    constantLayout.bindings = &constantBinding;
    constantLayout.bindingCount = 1;

    RHI::RHIPipelineLayoutDesc pipelineLayoutDesc = {};
    pipelineLayoutDesc.setLayouts = &constantLayout;
    pipelineLayoutDesc.setLayoutCount = 1;
    m_pipelineLayout = m_device->CreatePipelineLayout(pipelineLayoutDesc);
    if (m_pipelineLayout == nullptr) {
        DestroyResources();
        return false;
    }

    m_pipelineState = m_device->CreatePipelineState(CreatePipelineDesc(m_backendType, m_pipelineLayout));
    if (m_pipelineState == nullptr || !m_pipelineState->IsValid()) {
        if (m_pipelineState != nullptr) {
            m_pipelineState->Shutdown();
            delete m_pipelineState;
            m_pipelineState = nullptr;
        }
        DestroyResources();
        return false;
    }

    return true;
}

void BuiltinObjectIdPass::DestroyResources() {
    m_resourceCache.Shutdown();

    for (auto& descriptorSetEntry : m_perObjectSets) {
        DestroyOwnedDescriptorSet(descriptorSetEntry.second);
    }
    m_perObjectSets.clear();

    if (m_pipelineState != nullptr) {
        m_pipelineState->Shutdown();
        delete m_pipelineState;
        m_pipelineState = nullptr;
    }

    if (m_pipelineLayout != nullptr) {
        m_pipelineLayout->Shutdown();
        delete m_pipelineLayout;
        m_pipelineLayout = nullptr;
    }

    m_device = nullptr;
    m_backendType = RHI::RHIType::D3D12;
}

RHI::RHIDescriptorSet* BuiltinObjectIdPass::GetOrCreatePerObjectSet(uint64_t objectId) {
    const auto existing = m_perObjectSets.find(objectId);
    if (existing != m_perObjectSets.end()) {
        return existing->second.set;
    }

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

    OwnedDescriptorSet descriptorSet = {};
    descriptorSet.pool = m_device->CreateDescriptorPool(poolDesc);
    if (descriptorSet.pool == nullptr) {
        return nullptr;
    }

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

    RHI::DescriptorSetLayoutDesc layout = {};
    layout.bindings = &binding;
    layout.bindingCount = 1;
    descriptorSet.set = descriptorSet.pool->AllocateSet(layout);
    if (descriptorSet.set == nullptr) {
        DestroyOwnedDescriptorSet(descriptorSet);
        return nullptr;
    }

    const auto result = m_perObjectSets.emplace(objectId, descriptorSet);
    return result.first->second.set;
}

void BuiltinObjectIdPass::DestroyOwnedDescriptorSet(OwnedDescriptorSet& descriptorSet) {
    if (descriptorSet.set != nullptr) {
        descriptorSet.set->Shutdown();
        delete descriptorSet.set;
        descriptorSet.set = nullptr;
    }

    if (descriptorSet.pool != nullptr) {
        descriptorSet.pool->Shutdown();
        delete descriptorSet.pool;
        descriptorSet.pool = nullptr;
    }
}

bool BuiltinObjectIdPass::DrawVisibleItem(
    const RenderContext& context,
    const RenderSceneData& sceneData,
    const VisibleRenderItem& visibleItem) {
    if (visibleItem.mesh == nullptr || visibleItem.gameObject == nullptr) {
        return false;
    }

    const RenderResourceCache::CachedMesh* cachedMesh =
        m_resourceCache.GetOrCreateMesh(m_device, visibleItem.mesh);
    if (cachedMesh == nullptr || cachedMesh->vertexBufferView == nullptr) {
        return false;
    }

    RHI::RHICommandList* commandList = context.commandList;

    RHI::RHIResourceView* vertexBuffers[] = { cachedMesh->vertexBufferView };
    const uint64_t offsets[] = { 0 };
    const uint32_t strides[] = { cachedMesh->vertexStride };
    commandList->SetVertexBuffers(0, 1, vertexBuffers, offsets, strides);
    if (cachedMesh->indexBufferView != nullptr) {
        commandList->SetIndexBuffer(cachedMesh->indexBufferView, 0);
    }

    const uint64_t objectId = visibleItem.gameObject->GetID();
    RHI::RHIDescriptorSet* constantSet = GetOrCreatePerObjectSet(objectId);
    if (constantSet == nullptr) {
        return false;
    }

    const PerObjectConstants constants = {
        sceneData.cameraData.projection,
        sceneData.cameraData.view,
        visibleItem.localToWorld.Transpose(),
        EncodeObjectIdToColor(objectId)
    };
    constantSet->WriteConstant(0, &constants, sizeof(constants));

    RHI::RHIDescriptorSet* descriptorSets[] = { constantSet };
    commandList->SetGraphicsDescriptorSets(0, 1, descriptorSets, m_pipelineLayout);

    if (visibleItem.hasSection) {
        const Containers::Array<Resources::MeshSection>& sections = visibleItem.mesh->GetSections();
        if (visibleItem.sectionIndex >= sections.Size()) {
            return false;
        }

        const Resources::MeshSection& section = sections[visibleItem.sectionIndex];
        if (cachedMesh->indexBufferView != nullptr && section.indexCount > 0) {
            commandList->DrawIndexed(section.indexCount, 1, section.startIndex, 0, 0);
        } else if (section.vertexCount > 0) {
            commandList->Draw(section.vertexCount, 1, section.baseVertex, 0);
        }
        return true;
    }

    if (cachedMesh->indexBufferView != nullptr && cachedMesh->indexCount > 0) {
        commandList->DrawIndexed(cachedMesh->indexCount, 1, 0, 0, 0);
    } else if (cachedMesh->vertexCount > 0) {
        commandList->Draw(cachedMesh->vertexCount, 1, 0, 0);
    }

    return true;
}

} // namespace Passes
} // namespace Rendering
} // namespace XCEngine