#include "SceneViewportInfiniteGridPass.h"

#include "SceneViewportGrid.h"
#include "SceneViewportMath.h"

#include <XCEngine/Core/Math/Vector4.h>
#include <XCEngine/RHI/RHICommandList.h>
#include <XCEngine/RHI/RHIDevice.h>

#include <cstring>

namespace XCEngine {
namespace Editor {

namespace {

const char kInfiniteGridHlsl[] = R"(
cbuffer GridConstants : register(b0) {
    float4x4 gViewProjectionMatrix;
    float4 gCameraPositionAndScale;
    float4 gCameraRightAndFade;
    float4 gCameraUpAndTanHalfFov;
    float4 gCameraForwardAndAspect;
    float4 gViewportNearFar;
    float4 gGridTransition;
};

struct VSOutput {
    float4 position : SV_POSITION;
};

VSOutput MainVS(uint vertexId : SV_VertexID) {
    static const float2 positions[3] = {
        float2(-1.0, -1.0),
        float2(-1.0,  3.0),
        float2( 3.0, -1.0)
    };

    VSOutput output;
    output.position = float4(positions[vertexId], 0.0, 1.0);
    return output;
}

float PristineGridLine(float2 uv) {
    float2 deriv = max(fwidth(uv), float2(1e-6, 1e-6));
    float2 uvMod = frac(uv);
    float2 uvDist = min(uvMod, 1.0 - uvMod);
    float2 distInPixels = uvDist / deriv;
    float2 lineAlpha = 1.0 - smoothstep(0.0, 1.0, distInPixels);
    float density = max(deriv.x, deriv.y);
    float densityFade = 1.0 - smoothstep(0.5, 1.0, density);
    return max(lineAlpha.x, lineAlpha.y) * densityFade;
}

float AxisLineAA(float coord, float deriv) {
    float distInPixels = abs(coord) / max(deriv, 1e-6);
    return 1.0 - smoothstep(0.0, 1.5, distInPixels);
}

struct GridLayer {
    float minor;
    float major;
};

GridLayer SampleGridLayer(float2 worldPos2D, float baseScale) {
    GridLayer layer;
    const float2 gridCoord1 = worldPos2D / baseScale;
    const float2 gridCoord10 = worldPos2D / (baseScale * 10.0);
    const float grid1 = PristineGridLine(gridCoord1);
    const float grid10 = PristineGridLine(gridCoord10);
    const float2 deriv1 = fwidth(gridCoord1);
    const float lodFactor = smoothstep(0.3, 0.6, max(deriv1.x, deriv1.y));

    layer.major = max(grid10, grid1 * 0.35);
    layer.minor = grid1 * (1.0 - lodFactor);
    return layer;
}

struct PSOutput {
    float4 color : SV_TARGET0;
    float depth : SV_Depth;
};

PSOutput MainPS(VSOutput input) {
    const float2 viewportSize = max(gViewportNearFar.xy, float2(1.0, 1.0));
    const float scale = max(gCameraPositionAndScale.w, 1e-4);
    const float fadeDistance = max(gCameraRightAndFade.w, scale * 10.0);
    const float tanHalfFov = max(gCameraUpAndTanHalfFov.w, 1e-4);
    const float aspect = max(gCameraForwardAndAspect.w, 1e-4);
    const float transitionBlend = saturate(gGridTransition.x);
    const float nearClip = gViewportNearFar.z;
    const float sceneFarClip = gViewportNearFar.w;

    const float2 ndc = float2(
        (input.position.x / viewportSize.x) * 2.0 - 1.0,
        1.0 - (input.position.y / viewportSize.y) * 2.0);

    const float3 cameraPosition = gCameraPositionAndScale.xyz;
    const float3 rayDirection = normalize(
        gCameraForwardAndAspect.xyz +
        ndc.x * aspect * tanHalfFov * gCameraRightAndFade.xyz +
        ndc.y * tanHalfFov * gCameraUpAndTanHalfFov.xyz);

    if (abs(rayDirection.y) < 1e-5) {
        discard;
    }

    const float t = -cameraPosition.y / rayDirection.y;
    if (t <= nearClip) {
        discard;
    }

    const float3 worldPosition = cameraPosition + rayDirection * t;
    float depth = 0.999999;
    if (t < sceneFarClip) {
        const float4 clipPosition = mul(gViewProjectionMatrix, float4(worldPosition, 1.0));
        if (clipPosition.w <= 1e-6) {
            discard;
        }

        depth = clipPosition.z / clipPosition.w;
        if (depth <= 0.0 || depth >= 1.0) {
            discard;
        }
    }

    const float radialFade =
        1.0 - smoothstep(fadeDistance * 0.3, fadeDistance, length(worldPosition - cameraPosition));
    const float normalFade = smoothstep(0.0, 0.15, abs(rayDirection.y));
    const float fadeFactor = radialFade * normalFade;
    if (fadeFactor < 1e-3) {
        discard;
    }

    const float2 worldPos2D = worldPosition.xz;
    const GridLayer baseLayer = SampleGridLayer(worldPos2D, scale);
    const GridLayer nextLayer = SampleGridLayer(worldPos2D, scale * 10.0);
    const float minorGridIntensity = lerp(baseLayer.minor, nextLayer.minor, transitionBlend);
    const float majorGridIntensity = lerp(baseLayer.major, nextLayer.major, transitionBlend);
    float3 finalColor = float3(0.56, 0.56, 0.56);
    float finalAlpha = max(
        0.13 * minorGridIntensity * fadeFactor,
        0.28 * majorGridIntensity * fadeFactor);

    const float2 worldDeriv = max(fwidth(worldPos2D), float2(1e-6, 1e-6));
    const float xAxisAlpha = AxisLineAA(worldPos2D.y, worldDeriv.y) * fadeFactor;
    const float zAxisAlpha = AxisLineAA(worldPos2D.x, worldDeriv.x) * fadeFactor;

    const float axisAlpha = max(xAxisAlpha, zAxisAlpha);
    finalAlpha = max(finalAlpha, 0.34 * saturate(axisAlpha));

    if (finalAlpha < 1e-3) {
        discard;
    }

    PSOutput output;
    output.color = float4(finalColor, finalAlpha);
    output.depth = depth;
    return output;
}
)";

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

} // namespace

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

bool SceneViewportInfiniteGridPass::Render(
    const Rendering::RenderContext& renderContext,
    const Rendering::RenderSurface& surface,
    const SceneViewportOverlayData& overlay) {
    if (!overlay.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 SceneGridParameters parameters = BuildSceneGridParameters(overlay);
    if (!parameters.valid) {
        return false;
    }

    const Math::Matrix4x4 viewProjection =
        BuildSceneViewportProjectionMatrix(
            overlay,
            static_cast<float>(surface.GetWidth()),
            static_cast<float>(surface.GetHeight())) *
        BuildSceneViewportViewMatrix(overlay);

    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(overlay.cameraPosition, parameters.baseScale);
    constants.cameraRightAndFade = Math::Vector4(overlay.cameraRight, parameters.fadeDistance);
    constants.cameraUpAndTanHalfFov = Math::Vector4(
        overlay.cameraUp,
        std::tan(overlay.verticalFovDegrees * Math::DEG_TO_RAD * 0.5f));
    constants.cameraForwardAndAspect = Math::Vector4(overlay.cameraForward, aspect);
    constants.viewportNearFar = Math::Vector4(
        static_cast<float>(surface.GetWidth()),
        static_cast<float>(surface.GetHeight()),
        overlay.nearClipPlane,
        overlay.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 SceneViewportInfiniteGridPass::EnsureInitialized(const Rendering::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 SceneViewportInfiniteGridPass::CreateResources(const Rendering::RenderContext& renderContext) {
    if (!renderContext.IsValid() || renderContext.backendType != RHI::RHIType::D3D12) {
        return false;
    }

    m_device = renderContext.device;
    m_backendType = renderContext.backendType;

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

    pipelineDesc.vertexShader.source.assign(
        kInfiniteGridHlsl,
        kInfiniteGridHlsl + std::strlen(kInfiniteGridHlsl));
    pipelineDesc.vertexShader.sourceLanguage = RHI::ShaderLanguage::HLSL;
    pipelineDesc.vertexShader.entryPoint = L"MainVS";
    pipelineDesc.vertexShader.profile = L"vs_5_0";

    pipelineDesc.fragmentShader.source.assign(
        kInfiniteGridHlsl,
        kInfiniteGridHlsl + std::strlen(kInfiniteGridHlsl));
    pipelineDesc.fragmentShader.sourceLanguage = RHI::ShaderLanguage::HLSL;
    pipelineDesc.fragmentShader.entryPoint = L"MainPS";
    pipelineDesc.fragmentShader.profile = L"ps_5_0";

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

    return true;
}

void SceneViewportInfiniteGridPass::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;
}

} // namespace Editor
} // namespace XCEngine