#include "Rendering/Passes/BuiltinInfiniteGridPass.h"

#include <XCEngine/Core/Math/Matrix4.h>
#include <XCEngine/Core/Math/Vector4.h>
#include <XCEngine/Core/Asset/ResourceManager.h>
#include <XCEngine/Debug/Logger.h>
#include <XCEngine/RHI/RHICommandList.h>
#include <XCEngine/RHI/RHIDevice.h>
#include <XCEngine/Resources/BuiltinResources.h>

#include "Rendering/Detail/ShaderVariantUtils.h"

#include <algorithm>
#include <cmath>

namespace XCEngine {
namespace Rendering {
namespace Passes {

namespace {

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

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

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

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

    return nullptr;
}

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 InfiniteGridPassData& data) {
    const float cameraHeight = std::abs(data.cameraPosition.y);
    const 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;
}

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

    Math::Matrix4x4 view = 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] = -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] = -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] = -Math::Vector3::Dot(forward, data.cameraPosition);
    return view;
}

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

} // namespace

InfiniteGridParameters BuildInfiniteGridParameters(const InfiniteGridPassData& 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;
}

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

bool BuiltinInfiniteGridPass::Render(
    const RenderContext& renderContext,
    const RenderSurface& surface,
    const InfiniteGridPassData& data) {
    if (!data.valid || !renderContext.IsValid() || renderContext.backendType != RHI::RHIType::D3D12) {
        return false;
    }

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

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

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

    const Math::Matrix4x4 viewProjection =
        BuildInfiniteGridProjectionMatrix(
            data,
            static_cast<float>(surface.GetWidth()),
            static_cast<float>(surface.GetHeight())) *
        BuildInfiniteGridViewMatrix(data);

    const float aspect = surface.GetHeight() > 0
        ? static_cast<float>(surface.GetWidth()) / static_cast<float>(surface.GetHeight())
        : 1.0f;

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

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

    RHI::RHICommandList* commandList = renderContext.commandList;
    RHI::RHIResourceView* renderTarget = colorAttachments[0];
    commandList->SetRenderTargets(1, &renderTarget, surface.GetDepthAttachment());

    const RHI::Viewport viewport = {
        0.0f,
        0.0f,
        static_cast<float>(surface.GetWidth()),
        static_cast<float>(surface.GetHeight()),
        0.0f,
        1.0f
    };
    const RHI::Rect scissorRect = {
        0,
        0,
        static_cast<int32_t>(surface.GetWidth()),
        static_cast<int32_t>(surface.GetHeight())
    };

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

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

bool BuiltinInfiniteGridPass::EnsureInitialized(const RenderContext& renderContext) {
    if (m_pipelineState != nullptr &&
        m_pipelineLayout != nullptr &&
        m_constantPool != nullptr &&
        m_constantSet != nullptr &&
        m_device == renderContext.device &&
        m_backendType == renderContext.backendType) {
        return true;
    }

    DestroyResources();
    return CreateResources(renderContext);
}

bool BuiltinInfiniteGridPass::CreateResources(const RenderContext& renderContext) {
    if (!renderContext.IsValid() || renderContext.backendType != RHI::RHIType::D3D12) {
        return false;
    }

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

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

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

    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::GraphicsPipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = m_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.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 = true;
    pipelineDesc.blendState.srcBlend = static_cast<uint32_t>(RHI::BlendFactor::SrcAlpha);
    pipelineDesc.blendState.dstBlend = static_cast<uint32_t>(RHI::BlendFactor::InvSrcAlpha);
    pipelineDesc.blendState.srcBlendAlpha = static_cast<uint32_t>(RHI::BlendFactor::One);
    pipelineDesc.blendState.dstBlendAlpha = static_cast<uint32_t>(RHI::BlendFactor::InvSrcAlpha);
    pipelineDesc.blendState.blendOp = static_cast<uint32_t>(RHI::BlendOp::Add);
    pipelineDesc.blendState.blendOpAlpha = static_cast<uint32_t>(RHI::BlendOp::Add);
    pipelineDesc.blendState.colorWriteMask = 0xF;

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

    if (const Resources::ShaderStageVariant* vertexVariant = m_builtinInfiniteGridShader->FindVariant(
            infiniteGridPass->name,
            Resources::ShaderType::Vertex,
            backend)) {
        ::XCEngine::Rendering::Detail::ApplyShaderStageVariant(*vertexVariant, pipelineDesc.vertexShader);
    }
    if (const Resources::ShaderStageVariant* fragmentVariant = m_builtinInfiniteGridShader->FindVariant(
            infiniteGridPass->name,
            Resources::ShaderType::Fragment,
            backend)) {
        ::XCEngine::Rendering::Detail::ApplyShaderStageVariant(*fragmentVariant, pipelineDesc.fragmentShader);
    }

    m_pipelineState = m_device->CreatePipelineState(pipelineDesc);
    if (m_pipelineState == nullptr || !m_pipelineState->IsValid()) {
        DestroyResources();
        return false;
    }

    return true;
}

void BuiltinInfiniteGridPass::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 = RHI::RHIType::D3D12;
    m_builtinInfiniteGridShader.Reset();
}

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