#include "Rendering/Passes/BuiltinObjectIdPass.h"

#include "Core/Asset/ResourceManager.h"
#include "RHI/RHICommandList.h"
#include <XCEngine/Rendering/RenderPassGraphContract.h>
#include "Rendering/FrameData/RendererListUtils.h"
#include "Rendering/Internal/RenderSurfacePipelineUtils.h"
#include "Rendering/Extraction/RenderSceneExtractor.h"
#include "Rendering/RenderSurface.h"
#include "Resources/Mesh/Mesh.h"

#include <cstddef>
#include <vector>

namespace XCEngine {
namespace Rendering {
namespace Passes {

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

const char* BuiltinObjectIdPass::GetName() const {
    return "BuiltinObjectIdPass";
}

RHI::InputLayoutDesc BuiltinObjectIdPass::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);

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

    RHI::InputElementDesc texcoord = {};
    texcoord.semanticName = "TEXCOORD";
    texcoord.semanticIndex = 0;
    texcoord.format = static_cast<uint32_t>(RHI::Format::R32G32_Float);
    texcoord.inputSlot = 0;
    texcoord.alignedByteOffset = static_cast<uint32_t>(offsetof(Resources::StaticMeshVertex, uv0));
    inputLayout.elements.push_back(texcoord);

    return inputLayout;
}

bool BuiltinObjectIdPass::Initialize(const RenderContext& context) {
    return context.IsValid();
}

bool BuiltinObjectIdPass::SupportsRenderGraph() const {
    return true;
}

bool BuiltinObjectIdPass::RecordRenderGraph(
    const RenderPassRenderGraphContext& context) {
    return RecordRasterRenderPass(
        *this,
        context,
        {
            false,
            true,
            true
        });
}

bool BuiltinObjectIdPass::Execute(const RenderPassContext& context) {
    if (!context.renderContext.IsValid()) {
        return false;
    }

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

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

    if (!EnsureInitialized(context.renderContext, context.surface)) {
        return false;
    }

    RHI::RHICommandList* commandList = context.renderContext.commandList;
    RHI::RHIResourceView* renderTarget = colorAttachments[0];
    RHI::RHIResourceView* depthAttachment = context.surface.GetDepthAttachment();

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

    commandList->SetRenderTargets(1, &renderTarget, depthAttachment);

    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->ClearDepthStencil(depthAttachment, 1.0f, 0, 1, clearRects);
    commandList->SetPrimitiveTopology(RHI::PrimitiveTopology::TriangleList);
    commandList->SetPipelineState(m_pipelineState);

    VisitRendererListVisibleItems(
        context.sceneData,
        RendererListType::ObjectId,
        [&](const VisibleRenderItem& visibleItem) {
            DrawVisibleItem(context.renderContext, context.sceneData, visibleItem);
        });

    commandList->EndRenderPass();

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

    return true;
}

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

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