#include "Rendering/Viewport/Passes/SceneViewportGridPass.h"

#include "Rendering/Viewport/SceneViewportResourcePaths.h"

#include <XCEngine/Rendering/RenderSurfacePipelineUtils.h>
#include <XCEngine/Rendering/ShaderVariantUtils.h>
#include <XCEngine/Rendering/Materials/RenderMaterialStateUtils.h>

#include <XCEngine/Core/Asset/ResourceManager.h>
#include <XCEngine/Core/Math/Matrix4.h>
#include <XCEngine/Core/Math/Vector4.h>
#include <XCEngine/Debug/Logger.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/Shader/Shader.h>

#include <algorithm>
#include <cmath>
#include <utility>
#include <vector>

namespace XCEngine::UI::Editor::App {

namespace {

constexpr float kCameraHeightScaleFactor = 0.50f;
constexpr float kTransitionStart = 0.65f;
constexpr float kTransitionEnd = 0.95f;
constexpr float kMinimumVerticalViewComponent = 0.15f;

struct GridConstants {
    ::XCEngine::Math::Matrix4x4 viewProjection =
        ::XCEngine::Math::Matrix4x4::Identity();
    ::XCEngine::Math::Vector4 cameraPositionAndScale =
        ::XCEngine::Math::Vector4::Zero();
    ::XCEngine::Math::Vector4 cameraRightAndFade =
        ::XCEngine::Math::Vector4::Zero();
    ::XCEngine::Math::Vector4 cameraUpAndTanHalfFov =
        ::XCEngine::Math::Vector4::Zero();
    ::XCEngine::Math::Vector4 cameraForwardAndAspect =
        ::XCEngine::Math::Vector4::Zero();
    ::XCEngine::Math::Vector4 viewportNearFar =
        ::XCEngine::Math::Vector4::Zero();
    ::XCEngine::Math::Vector4 gridTransition =
        ::XCEngine::Math::Vector4::Zero();
};

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

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

    return nullptr;
}

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

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

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

    return pipelineDesc;
}

float SnapGridSpacing(float targetSpacing) {
    const float clampedTarget = (std::max)(targetSpacing, 0.02f);
    const float exponent = std::floor(std::log10(clampedTarget));
    return std::pow(10.0f, exponent);
}

float ComputeTransitionBlend(float targetSpacing, float baseScale) {
    const float normalizedSpacing =
        (std::max)(targetSpacing / (std::max)(baseScale, 1e-4f), 1.0f);
    const float transitionPosition = std::log10(normalizedSpacing);
    const float t =
        (transitionPosition - kTransitionStart) /
        (kTransitionEnd - kTransitionStart);
    const float saturated = (std::clamp)(t, 0.0f, 1.0f);
    return saturated * saturated * (3.0f - 2.0f * saturated);
}

float ComputeViewDistanceToGridPlane(const SceneViewportGridPassData& data) {
    const float cameraHeight = std::abs(data.cameraPosition.y);
    const ::XCEngine::Math::Vector3 forward = data.cameraForward.Normalized();

    const bool lookingTowardGrid =
        (data.cameraPosition.y >= 0.0f && forward.y < 0.0f) ||
        (data.cameraPosition.y < 0.0f && forward.y > 0.0f);
    if (!lookingTowardGrid) {
        return cameraHeight;
    }

    const float verticalViewComponent =
        (std::max)(std::abs(forward.y), kMinimumVerticalViewComponent);
    return cameraHeight / verticalViewComponent;
}

::XCEngine::Math::Matrix4x4 BuildInfiniteGridViewMatrix(
    const SceneViewportGridPassData& data) {
    const ::XCEngine::Math::Vector3 right = data.cameraRight.Normalized();
    const ::XCEngine::Math::Vector3 up = data.cameraUp.Normalized();
    const ::XCEngine::Math::Vector3 forward = data.cameraForward.Normalized();

    ::XCEngine::Math::Matrix4x4 view =
        ::XCEngine::Math::Matrix4x4::Identity();
    view.m[0][0] = right.x;
    view.m[0][1] = right.y;
    view.m[0][2] = right.z;
    view.m[0][3] =
        -::XCEngine::Math::Vector3::Dot(right, data.cameraPosition);

    view.m[1][0] = up.x;
    view.m[1][1] = up.y;
    view.m[1][2] = up.z;
    view.m[1][3] = -::XCEngine::Math::Vector3::Dot(up, data.cameraPosition);

    view.m[2][0] = forward.x;
    view.m[2][1] = forward.y;
    view.m[2][2] = forward.z;
    view.m[2][3] =
        -::XCEngine::Math::Vector3::Dot(forward, data.cameraPosition);
    return view;
}

::XCEngine::Math::Matrix4x4 BuildInfiniteGridProjectionMatrix(
    const SceneViewportGridPassData& data,
    float viewportWidth,
    float viewportHeight) {
    const float aspect = viewportHeight > 0.0f
        ? viewportWidth / viewportHeight
        : 1.0f;
    return ::XCEngine::Math::Matrix4x4::Perspective(
        data.verticalFovDegrees * ::XCEngine::Math::DEG_TO_RAD,
        aspect,
        data.nearClipPlane,
        data.farClipPlane);
}

class SceneViewportGridPass final : public ::XCEngine::Rendering::RenderPass {
public:
    SceneViewportGridPass(
        SceneViewportGridPassRenderer& renderer,
        const SceneViewportGridPassData& data)
        : m_renderer(renderer)
        , m_data(data) {
    }

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

    bool Execute(
        const ::XCEngine::Rendering::RenderPassContext& context) override {
        return m_renderer.Render(
            context.renderContext,
            context.surface,
            m_data);
    }

private:
    SceneViewportGridPassRenderer& m_renderer;
    SceneViewportGridPassData m_data = {};
};

} // namespace

InfiniteGridParameters BuildInfiniteGridParameters(
    const SceneViewportGridPassData& data) {
    InfiniteGridParameters parameters = {};
    if (!data.valid) {
        return parameters;
    }

    const float cameraHeight = std::abs(data.cameraPosition.y);
    const float viewDistance = ComputeViewDistanceToGridPlane(data);
    const float targetSpacing =
        (std::max)(cameraHeight * kCameraHeightScaleFactor, 0.1f);

    parameters.valid = true;
    parameters.baseScale = SnapGridSpacing(targetSpacing);
    parameters.transitionBlend =
        ComputeTransitionBlend(targetSpacing, parameters.baseScale);
    parameters.fadeDistance =
        (std::max)(parameters.baseScale * 320.0f, viewDistance * 80.0f);
    return parameters;
}

class SceneViewportGridPassRenderer::Impl {
public:
    Impl()
        : m_shaderPath(GetSceneViewportInfiniteGridShaderPath()) {
    }

    void Shutdown() {
        DestroyResources();
    }

    bool Render(
        const ::XCEngine::Rendering::RenderContext& renderContext,
        const ::XCEngine::Rendering::RenderSurface& surface,
        const SceneViewportGridPassData& data) {
        if (!data.valid || !renderContext.IsValid()) {
            return false;
        }

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

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

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

        const InfiniteGridParameters parameters =
            BuildInfiniteGridParameters(data);
        if (!parameters.valid) {
            return false;
        }

        const ::XCEngine::Math::Matrix4x4 viewProjection =
            BuildInfiniteGridProjectionMatrix(
                data,
                static_cast<float>(renderArea.width),
                static_cast<float>(renderArea.height)) *
            BuildInfiniteGridViewMatrix(data);

        const float aspect = renderArea.height > 0
            ? static_cast<float>(renderArea.width) /
                static_cast<float>(renderArea.height)
            : 1.0f;

        GridConstants constants = {};
        constants.viewProjection = viewProjection.Transpose();
        constants.cameraPositionAndScale =
            ::XCEngine::Math::Vector4(
                data.cameraPosition,
                parameters.baseScale);
        constants.cameraRightAndFade =
            ::XCEngine::Math::Vector4(
                data.cameraRight,
                parameters.fadeDistance);
        constants.cameraUpAndTanHalfFov = ::XCEngine::Math::Vector4(
            data.cameraUp,
            std::tan(
                data.verticalFovDegrees *
                ::XCEngine::Math::DEG_TO_RAD *
                0.5f));
        constants.cameraForwardAndAspect =
            ::XCEngine::Math::Vector4(data.cameraForward, aspect);
        constants.viewportNearFar = ::XCEngine::Math::Vector4(
            static_cast<float>(surface.GetWidth()),
            static_cast<float>(surface.GetHeight()),
            data.nearClipPlane,
            data.farClipPlane);
        constants.gridTransition =
            ::XCEngine::Math::Vector4(
                parameters.transitionBlend,
                0.0f,
                0.0f,
                0.0f);

        m_constantSet->WriteConstant(0, &constants, sizeof(constants));

        ::XCEngine::RHI::RHICommandList* commandList =
            renderContext.commandList;
        ::XCEngine::RHI::RHIResourceView* renderTarget = colorAttachments[0];
        ::XCEngine::RHI::RHIResourceView* depthAttachment =
            surface.GetDepthAttachment();
        if (surface.IsAutoTransitionEnabled()) {
            commandList->TransitionBarrier(
                renderTarget,
                surface.GetColorStateAfter(),
                ::XCEngine::RHI::ResourceStates::RenderTarget);
            commandList->TransitionBarrier(
                depthAttachment,
                surface.GetDepthStateAfter(),
                ::XCEngine::RHI::ResourceStates::DepthWrite);
        }

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

        const ::XCEngine::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 ::XCEngine::RHI::Rect scissorRect = {
            renderArea.x,
            renderArea.y,
            renderArea.x + renderArea.width,
            renderArea.y + renderArea.height
        };

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

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

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

        return true;
    }

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

        DestroyResources();
        return CreateResources(renderContext, surface);
    }

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

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

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

        m_device = renderContext.device;
        m_backendType = renderContext.backendType;
        m_shader = ::XCEngine::Resources::ResourceManager::Get()
                       .Load<::XCEngine::Resources::Shader>(m_shaderPath);
        if (!m_shader.IsValid()) {
            ::XCEngine::Debug::Logger::Get().Error(
                ::XCEngine::Debug::LogCategory::Rendering,
                "SceneViewportGridPassRenderer failed to load infinite-grid shader");
            DestroyResources();
            return false;
        }

        const ::XCEngine::Resources::ShaderBackend backend =
            ::XCEngine::Rendering::ToShaderBackend(m_backendType);
        const ::XCEngine::Resources::ShaderPass* infiniteGridPass =
            FindInfiniteGridCompatiblePass(*m_shader.Get(), backend);
        if (infiniteGridPass == nullptr) {
            ::XCEngine::Debug::Logger::Get().Error(
                ::XCEngine::Debug::LogCategory::Rendering,
                "SceneViewportGridPassRenderer could not resolve InfiniteGrid pass");
            DestroyResources();
            return false;
        }

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

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

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

        ::XCEngine::RHI::DescriptorPoolDesc constantPoolDesc = {};
        constantPoolDesc.type =
            ::XCEngine::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;
        }

        const ::XCEngine::RHI::GraphicsPipelineDesc pipelineDesc =
            CreatePipelineDesc(
                m_backendType,
                m_pipelineLayout,
                *m_shader.Get(),
                infiniteGridPass->name,
                surface);
        m_pipelineState = m_device->CreatePipelineState(pipelineDesc);
        if (m_pipelineState == nullptr || !m_pipelineState->IsValid()) {
            DestroyResources();
            return false;
        }

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

    void DestroyResources() {
        if (m_pipelineState != nullptr) {
            m_pipelineState->Shutdown();
            delete m_pipelineState;
            m_pipelineState = 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;
        }

        m_device = nullptr;
        m_backendType = ::XCEngine::RHI::RHIType::D3D12;
        m_shader.Reset();
        m_renderTargetFormat = ::XCEngine::RHI::Format::Unknown;
        m_depthStencilFormat = ::XCEngine::RHI::Format::Unknown;
        m_renderTargetSampleCount = 1u;
    }

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

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

SceneViewportGridPassRenderer::~SceneViewportGridPassRenderer() = default;

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

bool SceneViewportGridPassRenderer::Render(
    const ::XCEngine::Rendering::RenderContext& renderContext,
    const ::XCEngine::Rendering::RenderSurface& surface,
    const SceneViewportGridPassData& data) {
    return m_impl->Render(renderContext, surface, data);
}

std::unique_ptr<::XCEngine::Rendering::RenderPass> CreateSceneViewportGridPass(
    SceneViewportGridPassRenderer& renderer,
    const SceneViewportGridPassData& data) {
    return std::make_unique<SceneViewportGridPass>(renderer, data);
}

} // namespace XCEngine::UI::Editor::App