#include "Rendering/Passes/BuiltinObjectIdOutlinePass.h"

#include "Core/Asset/ResourceManager.h"
#include "Debug/Logger.h"
#include "Rendering/Internal/RenderSurfacePipelineUtils.h"
#include "Rendering/Internal/ShaderVariantUtils.h"
#include "Rendering/Materials/RenderMaterialStateUtils.h"
#include "Rendering/Picking/ObjectIdCodec.h"
#include "RHI/RHICommandList.h"
#include "RHI/RHIDevice.h"
#include <XCEngine/Resources/BuiltinResources.h>

#include <algorithm>
#include <utility>

namespace XCEngine {
namespace Rendering {
namespace Passes {

namespace {

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

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

    if (shader.GetPassCount() > 0 &&
        ::XCEngine::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 RenderSurface& surface) {
    RHI::GraphicsPipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = pipelineLayout;
    pipelineDesc.topologyType = static_cast<uint32_t>(RHI::PrimitiveTopologyType::Triangle);
    ::XCEngine::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) {
        ::XCEngine::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);
        ::XCEngine::Rendering::ApplyRenderState(fallbackState, pipelineDesc);
    }

    const Resources::ShaderBackend backend = ::XCEngine::Rendering::Internal::ToShaderBackend(backendType);
    if (const Resources::ShaderStageVariant* vertexVariant =
            shader.FindVariant(passName, Resources::ShaderType::Vertex, backend)) {
        if (shaderPass != nullptr) {
            ::XCEngine::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) {
            ::XCEngine::Rendering::Internal::ApplyShaderStageVariant(
                shader.GetPath(),
                *shaderPass,
                backend,
                *fragmentVariant,
                pipelineDesc.fragmentShader);
        }
    }

    return pipelineDesc;
}

} // namespace

BuiltinObjectIdOutlinePass::BuiltinObjectIdOutlinePass(Containers::String shaderPath)
    : m_shaderPath(shaderPath.Empty() ? Resources::GetBuiltinObjectIdOutlineShaderPath() : std::move(shaderPath)) {
    ResetState();
}

void BuiltinObjectIdOutlinePass::SetShaderPath(const Containers::String& shaderPath) {
    if (m_shaderPath == shaderPath) {
        return;
    }

    DestroyResources();
    m_shaderPath = shaderPath;
}

const Containers::String& BuiltinObjectIdOutlinePass::GetShaderPath() const {
    return m_shaderPath;
}

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

bool BuiltinObjectIdOutlinePass::Render(
    const RenderContext& renderContext,
    const RenderSurface& surface,
    const ObjectIdOutlinePassInputs& inputs,
    const std::vector<uint64_t>& selectedObjectIds,
    const ObjectIdOutlineStyle& style) {
    if (!renderContext.IsValid() ||
        inputs.objectIdTextureView == nullptr ||
        selectedObjectIds.empty()) {
        return false;
    }

    const std::vector<RHI::RHIResourceView*>& colorAttachments = surface.GetColorAttachments();
    if (!::XCEngine::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);

    uint32_t selectedCount = 0u;
    for (uint64_t selectedObjectId : selectedObjectIds) {
        if (selectedCount >= kMaxSelectedObjectCount) {
            break;
        }

        RenderObjectId renderObjectId = kInvalidRenderObjectId;
        if (!TryConvertRuntimeObjectIdToRenderObjectId(selectedObjectId, renderObjectId)) {
            continue;
        }

        constants.selectedObjectColors[selectedCount] =
            EncodeRenderObjectIdToColor(renderObjectId);
        ++selectedCount;
    }
    if (selectedCount == 0u) {
        return false;
    }

    constants.selectedInfo = Math::Vector4(
        static_cast<float>(selectedCount),
        style.debugSelectionMask ? 1.0f : 0.0f,
        style.outlineWidthPixels,
        0.0f);

    m_constantSet->WriteConstant(0, &constants, sizeof(constants));
    m_textureSet->Update(0, inputs.objectIdTextureView);

    RHI::RHICommandList* commandList = renderContext.commandList;
    RHI::RHIResourceView* renderTarget = colorAttachments[0];
    if (surface.IsAutoTransitionEnabled()) {
        commandList->TransitionBarrier(
            renderTarget,
            surface.GetColorStateAfter(),
            RHI::ResourceStates::RenderTarget);
        commandList->TransitionBarrier(
            inputs.objectIdTextureView,
            inputs.objectIdTextureState,
            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(
            inputs.objectIdTextureView,
            RHI::ResourceStates::PixelShaderResource,
            inputs.objectIdTextureState);
    }

    return true;
}

bool BuiltinObjectIdOutlinePass::EnsureInitialized(const RenderContext& renderContext, const RenderSurface& surface) {
    const RHI::Format renderTargetFormat =
        ::XCEngine::Rendering::Internal::ResolveSurfaceColorFormat(surface, 0u);
    const uint32_t renderTargetSampleCount =
        ::XCEngine::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 BuiltinObjectIdOutlinePass::CreateResources(const RenderContext& renderContext, const RenderSurface& surface) {
    if (!renderContext.IsValid()) {
        return false;
    }

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

    if (m_shaderPath.Empty()) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinObjectIdOutlinePass requires an injected shader path before resource creation");
        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,
            "BuiltinObjectIdOutlinePass failed to load configured object-id-outline shader resource");
        ResetState();
        return false;
    }

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

    const Resources::ShaderBackend backend = ::XCEngine::Rendering::Internal::ToShaderBackend(m_backendType);
    const Resources::ShaderPass* outlinePass =
        FindObjectIdOutlineCompatiblePass(*m_builtinObjectIdOutlineShader->Get(), backend);
    if (outlinePass == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinObjectIdOutlinePass could not resolve a valid ObjectIdOutline shader pass");
        DestroyResources();
        return false;
    }

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

    RHI::DescriptorSetLayoutDesc constantLayout = {};
    constantLayout.bindings = &setBindings[0];
    constantLayout.bindingCount = 1;

    RHI::DescriptorSetLayoutDesc textureLayout = {};
    textureLayout.bindings = &setBindings[1];
    textureLayout.bindingCount = 1;

    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 = 1;
    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_builtinObjectIdOutlineShader->Get(),
            outlinePass->name,
            surface));
    if (m_pipelineState == nullptr || !m_pipelineState->IsValid()) {
        DestroyResources();
        return false;
    }

    m_renderTargetFormat = ::XCEngine::Rendering::Internal::ResolveSurfaceColorFormat(surface, 0u);
    m_renderTargetSampleCount = ::XCEngine::Rendering::Internal::ResolveSurfaceSampleCount(surface);

    return true;
}

bool BuiltinObjectIdOutlinePass::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_builtinObjectIdOutlineShader.has_value();
}

void BuiltinObjectIdOutlinePass::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_builtinObjectIdOutlineShader.has_value()) {
        m_builtinObjectIdOutlineShader.reset();
    }

    ResetState();
}

void BuiltinObjectIdOutlinePass::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_builtinObjectIdOutlineShader.reset();
    m_renderTargetFormat = RHI::Format::Unknown;
    m_renderTargetSampleCount = 1u;
}

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