#include "Passes/SceneViewportSelectionOutlinePass.h"

#include "Viewport/ViewportHostRenderTargets.h"

#include "Rendering/Internal/RenderSurfacePipelineUtils.h"
#include "Rendering/Internal/ShaderVariantUtils.h"
#include "Rendering/Materials/RenderMaterialStateUtils.h"

#include <XCEngine/Core/Asset/ResourceManager.h>
#include <XCEngine/Core/Math/Vector4.h>
#include <XCEngine/Debug/Logger.h>
#include <XCEngine/Rendering/FrameData/RenderSceneData.h>
#include <XCEngine/Rendering/FrameData/VisibleRenderItem.h>
#include <XCEngine/Rendering/Passes/BuiltinDepthStylePassBase.h>
#include <XCEngine/Rendering/Picking/ObjectIdCodec.h>
#include <XCEngine/RHI/RHICommandList.h>
#include <XCEngine/RHI/RHIDescriptorPool.h>
#include <XCEngine/RHI/RHIDescriptorSet.h>
#include <XCEngine/RHI/RHIDevice.h>
#include <XCEngine/RHI/RHIPipelineLayout.h>
#include <XCEngine/RHI/RHIPipelineState.h>
#include <XCEngine/Resources/BuiltinResources.h>
#include <XCEngine/Resources/Shader/Shader.h>

#include <algorithm>
#include <optional>
#include <utility>

namespace XCEngine {
namespace Editor {

namespace {

class SceneViewportSelectionMaskPass final
    : public Rendering::Passes::BuiltinDepthStylePassBase {
public:
    SceneViewportSelectionMaskPass()
        : BuiltinDepthStylePassBase(
            Rendering::BuiltinMaterialPass::SelectionMask,
            Resources::GetBuiltinSelectionMaskShaderPath()) {
    }

    const char* GetName() const override {
        return "SceneViewportSelectionMaskPass";
    }

    bool Render(
        const Rendering::RenderContext& context,
        const Rendering::RenderSurface& surface,
        const Rendering::RenderSceneData& sceneData,
        const std::vector<uint64_t>& selectedObjectIds) {
        m_selectedObjectIds.clear();
        m_selectedObjectIds.reserve(selectedObjectIds.size());
        for (uint64_t selectedObjectId : selectedObjectIds) {
            Rendering::RenderObjectId renderObjectId =
                Rendering::kInvalidRenderObjectId;
            if (Rendering::TryConvertRuntimeObjectIdToRenderObjectId(
                    selectedObjectId,
                    renderObjectId)) {
                m_selectedObjectIds.push_back(renderObjectId);
            }
        }

        if (m_selectedObjectIds.empty()) {
            return false;
        }

        Rendering::RenderSceneData selectionMaskSceneData = sceneData;
        selectionMaskSceneData.cameraData.clearFlags =
            Rendering::RenderClearFlags::Color;
        selectionMaskSceneData.cameraData.clearColor = Math::Color::Black();

        const Rendering::RenderPassContext passContext = {
            context,
            surface,
            selectionMaskSceneData,
            nullptr,
            nullptr,
            RHI::ResourceStates::Common
        };
        return Execute(passContext);
    }

protected:
    bool ShouldRenderVisibleItem(
        const Rendering::VisibleRenderItem& visibleItem) const override {
        if (!Rendering::IsValidRenderObjectId(visibleItem.renderObjectId)) {
            return false;
        }

        return std::find(
                   m_selectedObjectIds.begin(),
                   m_selectedObjectIds.end(),
                   visibleItem.renderObjectId) != m_selectedObjectIds.end();
    }

private:
    std::vector<Rendering::RenderObjectId> m_selectedObjectIds = {};
};

const Resources::ShaderPass* FindSelectionOutlineCompatiblePass(
    const Resources::Shader& shader,
    Resources::ShaderBackend backend) {
    const Resources::ShaderPass* outlinePass =
        shader.FindPass("SelectionOutline");
    if (outlinePass != nullptr &&
        Rendering::Internal::ShaderPassHasGraphicsVariants(
            shader,
            outlinePass->name,
            backend)) {
        return outlinePass;
    }

    const Resources::ShaderPass* editorOutlinePass =
        shader.FindPass("EditorSelectionOutline");
    if (editorOutlinePass != nullptr &&
        Rendering::Internal::ShaderPassHasGraphicsVariants(
            shader,
            editorOutlinePass->name,
            backend)) {
        return editorOutlinePass;
    }

    if (shader.GetPassCount() > 0 &&
        Rendering::Internal::ShaderPassHasGraphicsVariants(
            shader,
            shader.GetPasses()[0].name,
            backend)) {
        return &shader.GetPasses()[0];
    }

    return nullptr;
}

RHI::GraphicsPipelineDesc CreatePipelineDesc(
    RHI::RHIType backendType,
    RHI::RHIPipelineLayout* pipelineLayout,
    const Resources::Shader& shader,
    const Containers::String& passName,
    const Rendering::RenderSurface& surface) {
    RHI::GraphicsPipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = pipelineLayout;
    pipelineDesc.topologyType =
        static_cast<uint32_t>(RHI::PrimitiveTopologyType::Triangle);
    Rendering::Internal::ApplySingleColorAttachmentPropertiesToGraphicsPipelineDesc(
        surface,
        pipelineDesc);
    pipelineDesc.depthStencilFormat =
        static_cast<uint32_t>(RHI::Format::Unknown);

    const Resources::ShaderPass* shaderPass = shader.FindPass(passName);
    if (shaderPass != nullptr && shaderPass->hasFixedFunctionState) {
        Rendering::ApplyRenderState(
            shaderPass->fixedFunctionState,
            pipelineDesc);
    } else {
        Resources::MaterialRenderState fallbackState = {};
        fallbackState.cullMode = Resources::MaterialCullMode::None;
        fallbackState.depthWriteEnable = false;
        fallbackState.depthTestEnable = false;
        fallbackState.depthFunc = Resources::MaterialComparisonFunc::Always;
        fallbackState.blendEnable = true;
        fallbackState.srcBlend = Resources::MaterialBlendFactor::SrcAlpha;
        fallbackState.dstBlend = Resources::MaterialBlendFactor::InvSrcAlpha;
        fallbackState.srcBlendAlpha = Resources::MaterialBlendFactor::One;
        fallbackState.dstBlendAlpha =
            Resources::MaterialBlendFactor::InvSrcAlpha;
        fallbackState.blendOp = Resources::MaterialBlendOp::Add;
        fallbackState.blendOpAlpha = Resources::MaterialBlendOp::Add;
        fallbackState.colorWriteMask =
            static_cast<uint8_t>(RHI::ColorWriteMask::All);
        Rendering::ApplyRenderState(fallbackState, pipelineDesc);
    }

    const Resources::ShaderBackend backend =
        Rendering::Internal::ToShaderBackend(backendType);
    if (const Resources::ShaderStageVariant* vertexVariant =
            shader.FindVariant(
                passName,
                Resources::ShaderType::Vertex,
                backend)) {
        if (shaderPass != nullptr) {
            Rendering::Internal::ApplyShaderStageVariant(
                shader.GetPath(),
                *shaderPass,
                backend,
                *vertexVariant,
                pipelineDesc.vertexShader);
        }
    }
    if (const Resources::ShaderStageVariant* fragmentVariant =
            shader.FindVariant(
                passName,
                Resources::ShaderType::Fragment,
                backend)) {
        if (shaderPass != nullptr) {
            Rendering::Internal::ApplyShaderStageVariant(
                shader.GetPath(),
                *shaderPass,
                backend,
                *fragmentVariant,
                pipelineDesc.fragmentShader);
        }
    }

    return pipelineDesc;
}

class SceneViewportSelectionOutlinePass final : public Rendering::RenderPass {
public:
    SceneViewportSelectionOutlinePass(
        SceneViewportSelectionOutlinePassRenderer& renderer,
        ViewportRenderTargets* targets,
        std::vector<uint64_t> selectedObjectIds,
        const SceneViewportSelectionOutlineStyle& style)
        : m_renderer(renderer)
        , m_targets(targets)
        , m_selectedObjectIds(std::move(selectedObjectIds))
        , m_style(style) {
    }

    const char* GetName() const override {
        return "SceneViewportSelectionOutline";
    }

    bool Execute(const Rendering::RenderPassContext& context) override {
        return m_renderer.Render(
            context.renderContext,
            context.surface,
            context.sceneData,
            *m_targets,
            m_selectedObjectIds,
            m_style);
    }

private:
    SceneViewportSelectionOutlinePassRenderer& m_renderer;
    ViewportRenderTargets* m_targets = nullptr;
    std::vector<uint64_t> m_selectedObjectIds = {};
    SceneViewportSelectionOutlineStyle m_style = {};
};

} // namespace

class SceneViewportSelectionOutlinePassRenderer::Impl {
public:
    struct OutlineConstants {
        Math::Vector4 viewportSizeAndTexelSize = Math::Vector4::Zero();
        Math::Vector4 outlineColor = Math::Vector4::Zero();
        Math::Vector4 outlineInfo = Math::Vector4::Zero();
        Math::Vector4 depthParams = Math::Vector4::Zero();
    };

    Impl()
        : m_selectionMaskPass(std::make_unique<SceneViewportSelectionMaskPass>())
        , m_shaderPath(Resources::GetBuiltinSelectionOutlineShaderPath()) {
    }

    void Shutdown() {
        if (m_selectionMaskPass != nullptr) {
            m_selectionMaskPass->Shutdown();
        }
        DestroyResources();
    }

    bool Render(
        const Rendering::RenderContext& renderContext,
        const Rendering::RenderSurface& surface,
        const Rendering::RenderSceneData& sceneData,
        ViewportRenderTargets& targets,
        const std::vector<uint64_t>& selectedObjectIds,
        const SceneViewportSelectionOutlineStyle& style) {
        Rendering::RenderSurface selectionMaskSurface =
            BuildViewportSelectionMaskSurface(targets);
        selectionMaskSurface.SetRenderArea(surface.GetRenderArea());

        if (!m_selectionMaskPass->Render(
                renderContext,
                selectionMaskSurface,
                sceneData,
                selectedObjectIds)) {
            return false;
        }

        targets.selectionMaskState = RHI::ResourceStates::PixelShaderResource;
        return RenderOutline(
            renderContext,
            surface,
            targets.selectionMaskShaderView,
            targets.selectionMaskState,
            targets.depthShaderView,
            surface.GetDepthStateAfter(),
            style);
    }

private:
    bool RenderOutline(
        const Rendering::RenderContext& renderContext,
        const Rendering::RenderSurface& surface,
        RHI::RHIResourceView* selectionMaskTextureView,
        RHI::ResourceStates selectionMaskState,
        RHI::RHIResourceView* depthTextureView,
        RHI::ResourceStates depthTextureState,
        const SceneViewportSelectionOutlineStyle& style) {
        if (!renderContext.IsValid() ||
            selectionMaskTextureView == nullptr ||
            depthTextureView == nullptr) {
            return false;
        }

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

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

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

        OutlineConstants constants = {};
        constants.viewportSizeAndTexelSize = Math::Vector4(
            static_cast<float>(surface.GetWidth()),
            static_cast<float>(surface.GetHeight()),
            surface.GetWidth() > 0
                ? 1.0f / static_cast<float>(surface.GetWidth())
                : 0.0f,
            surface.GetHeight() > 0
                ? 1.0f / static_cast<float>(surface.GetHeight())
                : 0.0f);
        constants.outlineColor = Math::Vector4(
            style.outlineColor.r,
            style.outlineColor.g,
            style.outlineColor.b,
            style.outlineColor.a);
        constants.outlineInfo = Math::Vector4(
            style.debugSelectionMask ? 1.0f : 0.0f,
            style.outlineWidthPixels,
            0.0f,
            0.0f);
        constants.depthParams = Math::Vector4(1.0e-5f, 0.0f, 0.0f, 0.0f);

        m_constantSet->WriteConstant(0, &constants, sizeof(constants));
        m_textureSet->Update(0, selectionMaskTextureView);
        m_textureSet->Update(1, depthTextureView);

        RHI::RHICommandList* commandList = renderContext.commandList;
        RHI::RHIResourceView* renderTarget = colorAttachments[0];
        if (surface.IsAutoTransitionEnabled()) {
            commandList->TransitionBarrier(
                renderTarget,
                surface.GetColorStateAfter(),
                RHI::ResourceStates::RenderTarget);
            commandList->TransitionBarrier(
                selectionMaskTextureView,
                selectionMaskState,
                RHI::ResourceStates::PixelShaderResource);
            commandList->TransitionBarrier(
                depthTextureView,
                depthTextureState,
                RHI::ResourceStates::PixelShaderResource);
        }
        commandList->SetRenderTargets(1, &renderTarget, nullptr);

        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
        };

        commandList->SetViewport(viewport);
        commandList->SetScissorRect(scissorRect);
        commandList->SetPrimitiveTopology(RHI::PrimitiveTopology::TriangleList);
        commandList->SetPipelineState(m_pipelineState);

        RHI::RHIDescriptorSet* descriptorSets[] = {
            m_constantSet,
            m_textureSet
        };
        commandList->SetGraphicsDescriptorSets(
            0,
            2,
            descriptorSets,
            m_pipelineLayout);
        commandList->Draw(3, 1, 0, 0);
        commandList->EndRenderPass();

        if (surface.IsAutoTransitionEnabled()) {
            commandList->TransitionBarrier(
                renderTarget,
                RHI::ResourceStates::RenderTarget,
                surface.GetColorStateAfter());
            commandList->TransitionBarrier(
                selectionMaskTextureView,
                RHI::ResourceStates::PixelShaderResource,
                selectionMaskState);
            commandList->TransitionBarrier(
                depthTextureView,
                RHI::ResourceStates::PixelShaderResource,
                depthTextureState);
        }
        return true;
    }

    bool EnsureInitialized(
        const Rendering::RenderContext& renderContext,
        const Rendering::RenderSurface& surface) {
        const RHI::Format renderTargetFormat =
            Rendering::Internal::ResolveSurfaceColorFormat(surface, 0u);
        const uint32_t renderTargetSampleCount =
            Rendering::Internal::ResolveSurfaceSampleCount(surface);
        if (m_pipelineLayout != nullptr &&
            m_pipelineState != nullptr &&
            m_constantPool != nullptr &&
            m_constantSet != nullptr &&
            m_texturePool != nullptr &&
            m_textureSet != nullptr &&
            m_device == renderContext.device &&
            m_backendType == renderContext.backendType &&
            m_renderTargetFormat == renderTargetFormat &&
            m_renderTargetSampleCount == renderTargetSampleCount) {
            return true;
        }

        if (HasCreatedResources()) {
            DestroyResources();
        }
        return CreateResources(renderContext, surface);
    }

    bool CreateResources(
        const Rendering::RenderContext& renderContext,
        const Rendering::RenderSurface& surface) {
        if (!renderContext.IsValid()) {
            return false;
        }

        if (!Rendering::Internal::HasSingleColorAttachment(surface) ||
            Rendering::Internal::ResolveSurfaceColorFormat(surface, 0u) ==
                RHI::Format::Unknown) {
            return false;
        }

        if (m_shaderPath.Empty()) {
            Debug::Logger::Get().Error(
                Debug::LogCategory::Rendering,
                "SceneViewportSelectionOutlinePassRenderer requires a valid shader path");
            return false;
        }

        Resources::ResourceHandle<Resources::Shader> shader =
            Resources::ResourceManager::Get().Load<Resources::Shader>(
                m_shaderPath);
        if (!shader.IsValid()) {
            Debug::Logger::Get().Error(
                Debug::LogCategory::Rendering,
                "SceneViewportSelectionOutlinePassRenderer failed to load selection-outline shader");
            ResetState();
            return false;
        }

        m_device = renderContext.device;
        m_backendType = renderContext.backendType;
        m_shader.emplace(std::move(shader));

        const Resources::ShaderBackend backend =
            Rendering::Internal::ToShaderBackend(m_backendType);
        const Resources::ShaderPass* outlinePass =
            FindSelectionOutlineCompatiblePass(*m_shader->Get(), backend);
        if (outlinePass == nullptr) {
            Debug::Logger::Get().Error(
                Debug::LogCategory::Rendering,
                "SceneViewportSelectionOutlinePassRenderer could not resolve SelectionOutline pass");
            DestroyResources();
            return false;
        }

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

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

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

        RHI::DescriptorSetLayoutDesc textureLayout = {};
        textureLayout.bindings = textureBindings;
        textureLayout.bindingCount = 2;

        RHI::DescriptorSetLayoutDesc setLayouts[2] = {};
        setLayouts[0] = constantLayout;
        setLayouts[1] = textureLayout;

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

        RHI::DescriptorPoolDesc constantPoolDesc = {};
        constantPoolDesc.type = RHI::DescriptorHeapType::CBV_SRV_UAV;
        constantPoolDesc.descriptorCount = 1;
        constantPoolDesc.shaderVisible = false;
        m_constantPool = m_device->CreateDescriptorPool(constantPoolDesc);
        if (m_constantPool == nullptr) {
            DestroyResources();
            return false;
        }

        m_constantSet = m_constantPool->AllocateSet(constantLayout);
        if (m_constantSet == nullptr) {
            DestroyResources();
            return false;
        }

        RHI::DescriptorPoolDesc texturePoolDesc = {};
        texturePoolDesc.type = RHI::DescriptorHeapType::CBV_SRV_UAV;
        texturePoolDesc.descriptorCount = 2;
        texturePoolDesc.shaderVisible = true;
        m_texturePool = m_device->CreateDescriptorPool(texturePoolDesc);
        if (m_texturePool == nullptr) {
            DestroyResources();
            return false;
        }

        m_textureSet = m_texturePool->AllocateSet(textureLayout);
        if (m_textureSet == nullptr) {
            DestroyResources();
            return false;
        }

        m_pipelineState = m_device->CreatePipelineState(
            CreatePipelineDesc(
                m_backendType,
                m_pipelineLayout,
                *m_shader->Get(),
                outlinePass->name,
                surface));
        if (m_pipelineState == nullptr || !m_pipelineState->IsValid()) {
            DestroyResources();
            return false;
        }

        m_renderTargetFormat =
            Rendering::Internal::ResolveSurfaceColorFormat(surface, 0u);
        m_renderTargetSampleCount =
            Rendering::Internal::ResolveSurfaceSampleCount(surface);
        return true;
    }

    bool HasCreatedResources() const {
        return m_device != nullptr ||
            m_pipelineLayout != nullptr ||
            m_pipelineState != nullptr ||
            m_constantPool != nullptr ||
            m_constantSet != nullptr ||
            m_texturePool != nullptr ||
            m_textureSet != nullptr ||
            m_shader.has_value();
    }

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

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

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

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

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

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

        if (m_shader.has_value()) {
            m_shader.reset();
        }

        ResetState();
    }

    void ResetState() {
        m_device = nullptr;
        m_backendType = RHI::RHIType::D3D12;
        m_pipelineLayout = nullptr;
        m_pipelineState = nullptr;
        m_constantPool = nullptr;
        m_constantSet = nullptr;
        m_texturePool = nullptr;
        m_textureSet = nullptr;
        m_shader.reset();
        m_renderTargetFormat = RHI::Format::Unknown;
        m_renderTargetSampleCount = 1u;
    }

    std::unique_ptr<SceneViewportSelectionMaskPass> m_selectionMaskPass;
    RHI::RHIDevice* m_device = nullptr;
    RHI::RHIType m_backendType = RHI::RHIType::D3D12;
    RHI::RHIPipelineLayout* m_pipelineLayout = nullptr;
    RHI::RHIPipelineState* m_pipelineState = nullptr;
    RHI::RHIDescriptorPool* m_constantPool = nullptr;
    RHI::RHIDescriptorSet* m_constantSet = nullptr;
    RHI::RHIDescriptorPool* m_texturePool = nullptr;
    RHI::RHIDescriptorSet* m_textureSet = nullptr;
    Containers::String m_shaderPath = {};
    std::optional<Resources::ResourceHandle<Resources::Shader>> m_shader = {};
    RHI::Format m_renderTargetFormat = RHI::Format::Unknown;
    uint32_t m_renderTargetSampleCount = 1u;
};

SceneViewportSelectionOutlinePassRenderer::SceneViewportSelectionOutlinePassRenderer()
    : m_impl(std::make_unique<Impl>()) {
}

SceneViewportSelectionOutlinePassRenderer::~SceneViewportSelectionOutlinePassRenderer() = default;

void SceneViewportSelectionOutlinePassRenderer::Shutdown() {
    m_impl->Shutdown();
}

bool SceneViewportSelectionOutlinePassRenderer::Render(
    const Rendering::RenderContext& renderContext,
    const Rendering::RenderSurface& surface,
    const Rendering::RenderSceneData& sceneData,
    ViewportRenderTargets& targets,
    const std::vector<uint64_t>& selectedObjectIds,
    const SceneViewportSelectionOutlineStyle& style) {
    return m_impl->Render(
        renderContext,
        surface,
        sceneData,
        targets,
        selectedObjectIds,
        style);
}

std::unique_ptr<Rendering::RenderPass> CreateSceneViewportSelectionOutlinePass(
    SceneViewportSelectionOutlinePassRenderer& renderer,
    ViewportRenderTargets* targets,
    const std::vector<uint64_t>& selectedObjectIds,
    const SceneViewportSelectionOutlineStyle& style) {
    return std::make_unique<SceneViewportSelectionOutlinePass>(
        renderer,
        targets,
        selectedObjectIds,
        style);
}

} // namespace Editor
} // namespace XCEngine