#include "Rendering/Pipelines/BuiltinForwardPipeline.h"

#include "Debug/Logger.h"
#include "Rendering/Graph/RenderGraph.h"
#include "Rendering/Pipelines/Internal/BuiltinForwardSceneSetup.h"
#include "RHI/RHICommandList.h"
#include "Rendering/Internal/RenderSurfacePipelineUtils.h"
#include "Rendering/RenderSurface.h"

#include <memory>
#include <string>

namespace XCEngine {
namespace Rendering {
namespace Pipelines {
namespace {

bool IsDepthFormat(RHI::Format format) {
    return format == RHI::Format::D24_UNorm_S8_UInt ||
           format == RHI::Format::D32_Float;
}

void TransitionMainDirectionalShadowForSampling(
    const RenderContext& context,
    const RenderSceneData& sceneData) {
    context.commandList->TransitionBarrier(
        sceneData.lighting.mainDirectionalShadow.shadowMap,
        RHI::ResourceStates::DepthWrite,
        RHI::ResourceStates::PixelShaderResource);
}

void RestoreMainDirectionalShadowAfterSampling(
    const RenderContext& context,
    const RenderSceneData& sceneData) {
    context.commandList->TransitionBarrier(
        sceneData.lighting.mainDirectionalShadow.shadowMap,
        RHI::ResourceStates::PixelShaderResource,
        RHI::ResourceStates::DepthWrite);
}

std::vector<RHI::RHIResourceView*> CollectSurfaceColorAttachments(const RenderSurface& surface) {
    std::vector<RHI::RHIResourceView*> renderTargets;
    const Core::uint32 colorAttachmentCount =
        ::XCEngine::Rendering::Internal::ResolveSurfaceColorAttachmentCount(surface);
    renderTargets.reserve(colorAttachmentCount);
    for (Core::uint32 attachmentIndex = 0; attachmentIndex < colorAttachmentCount; ++attachmentIndex) {
        RHI::RHIResourceView* renderTarget = surface.GetColorAttachments()[attachmentIndex];
        if (renderTarget == nullptr) {
            break;
        }

        renderTargets.push_back(renderTarget);
    }

    return renderTargets;
}

RenderSurface BuildGraphManagedForwardSceneSurface(const RenderSurface& templateSurface) {
    RenderSurface surface = templateSurface;
    surface.SetAutoTransitionEnabled(false);
    return surface;
}

struct ForwardSceneGraphExecutionState {
    bool initialized = false;
};

Containers::String BuildForwardScenePhasePassName(
    const Containers::String& baseName,
    ScenePhase scenePhase) {
    std::string name = baseName.CStr();
    name += '.';
    name += ToString(scenePhase);
    return Containers::String(name.c_str());
}

bool ScenePhaseSamplesMainDirectionalShadow(ScenePhase scenePhase) {
    return scenePhase == ScenePhase::Opaque ||
           scenePhase == ScenePhase::Transparent;
}

} // namespace

bool BuiltinForwardPipeline::ShouldSampleMainDirectionalShadowMap(const RenderSceneData& sceneData) {
    return sceneData.lighting.HasMainDirectionalShadow() &&
        sceneData.lighting.mainDirectionalShadow.shadowMap != nullptr &&
        IsDepthFormat(sceneData.lighting.mainDirectionalShadow.shadowMap->GetFormat());
}

bool BuiltinForwardPipeline::SupportsMainSceneRenderGraph() const {
    return true;
}

bool BuiltinForwardPipeline::RecordMainSceneRenderGraph(
    const RenderPipelineMainSceneRenderGraphContext& context) {
    const Containers::String passName = context.passName;
    const RenderContext renderContext = context.renderContext;
    const std::shared_ptr<const RenderSceneData> sceneData =
        std::make_shared<RenderSceneData>(context.sceneData);
    const RenderSurface surfaceTemplate =
        BuildGraphManagedForwardSceneSurface(context.surfaceTemplate);
    const bool hasSourceSurface = context.sourceSurface != nullptr;
    const RenderSurface sourceSurface =
        hasSourceSurface ? *context.sourceSurface : RenderSurface();
    RHI::RHIResourceView* const sourceColorView = context.sourceColorView;
    const RHI::ResourceStates sourceColorState = context.sourceColorState;
    const std::vector<RenderGraphTextureHandle> colorTargets = context.colorTargets;
    const RenderGraphTextureHandle depthTarget = context.depthTarget;
    const RenderGraphTextureHandle mainDirectionalShadowTexture =
        context.mainDirectionalShadowTexture;
    bool* const executionSucceeded = context.executionSucceeded;
    const std::shared_ptr<ForwardSceneGraphExecutionState> graphExecutionState =
        std::make_shared<ForwardSceneGraphExecutionState>();
    const SceneRenderFeaturePassBeginCallback beginRecordedPass =
        [this,
         graphExecutionState,
         renderContext,
         sceneData,
         hasSourceSurface,
         sourceSurface,
         sourceColorView,
         sourceColorState,
         executionSucceeded](
            const RenderPassContext& passContext,
            bool clearAttachments) -> bool {
            if (executionSucceeded != nullptr && !(*executionSucceeded)) {
                return false;
            }

            if (!graphExecutionState->initialized) {
                if (!Initialize(renderContext)) {
                    Debug::Logger::Get().Error(
                        Debug::LogCategory::Rendering,
                        "BuiltinForwardPipeline::RecordMainSceneRenderGraph failed during execution: Initialize returned false");
                    if (executionSucceeded != nullptr) {
                        *executionSucceeded = false;
                    }
                    return false;
                }

                const FrameExecutionContext executionContext(
                    renderContext,
                    passContext.surface,
                    *sceneData,
                    hasSourceSurface ? &sourceSurface : nullptr,
                    sourceColorView,
                    sourceColorState);
                if (!m_forwardSceneFeatureHost.Prepare(executionContext)) {
                    Debug::Logger::Get().Error(
                        Debug::LogCategory::Rendering,
                        "BuiltinForwardPipeline::RecordMainSceneRenderGraph failed during execution: SceneRenderFeatureHost::Prepare returned false");
                    if (executionSucceeded != nullptr) {
                        *executionSucceeded = false;
                    }
                    return false;
                }

                graphExecutionState->initialized = true;
            }

            if (!BeginForwardScenePass(passContext, clearAttachments)) {
                Debug::Logger::Get().Error(
                    Debug::LogCategory::Rendering,
                    "BuiltinForwardPipeline::RecordMainSceneRenderGraph failed during execution: BeginForwardScenePass returned false");
                if (executionSucceeded != nullptr) {
                    *executionSucceeded = false;
                }
                return false;
            }

            return true;
        };
    const SceneRenderFeaturePassEndCallback endRecordedPass =
        [this](const RenderPassContext& passContext) {
            EndForwardScenePass(passContext);
        };

    bool clearAttachments = true;
    for (const Pipelines::Internal::ForwardSceneStep& step : Pipelines::Internal::GetBuiltinForwardSceneSteps()) {
        if (step.type == Pipelines::Internal::ForwardSceneStepType::InjectionPoint) {
            const SceneRenderFeaturePassRenderGraphContext featureContext = {
                context.graphBuilder,
                passName,
                renderContext,
                *sceneData,
                surfaceTemplate,
                hasSourceSurface ? &sourceSurface : nullptr,
                sourceColorView,
                sourceColorState,
                {},
                colorTargets,
                depthTarget,
                mainDirectionalShadowTexture,
                clearAttachments,
                executionSucceeded,
                beginRecordedPass,
                endRecordedPass,
                context.blackboard
            };
            bool recordedAnyPass = false;
            if (!m_forwardSceneFeatureHost.Record(
                    featureContext,
                    step.injectionPoint,
                    &recordedAnyPass)) {
                return false;
            }

            if (recordedAnyPass) {
                clearAttachments = false;
            }
            continue;
        }

        const Containers::String phasePassName =
            BuildForwardScenePhasePassName(passName, step.scenePhase);
        context.graphBuilder.AddRasterPass(
            phasePassName,
            [this,
             surfaceTemplate,
             renderContext,
             sceneData,
             hasSourceSurface,
             sourceSurface,
             sourceColorView,
             sourceColorState,
             colorTargets,
             depthTarget,
             mainDirectionalShadowTexture,
             executionSucceeded,
             beginRecordedPass,
             endRecordedPass,
             clearAttachments,
             scenePhase = step.scenePhase](
                RenderGraphPassBuilder& passBuilder) {
                for (RenderGraphTextureHandle colorTarget : colorTargets) {
                    if (colorTarget.IsValid()) {
                        passBuilder.WriteTexture(colorTarget);
                    }
                }

                if (depthTarget.IsValid()) {
                    passBuilder.WriteDepthTexture(depthTarget);
                }
                if (ScenePhaseSamplesMainDirectionalShadow(scenePhase) &&
                    mainDirectionalShadowTexture.IsValid()) {
                    passBuilder.ReadTexture(mainDirectionalShadowTexture);
                }

                passBuilder.SetExecuteCallback(
                    [this,
                     surfaceTemplate,
                     renderContext,
                     sceneData,
                     hasSourceSurface,
                     sourceSurface,
                     sourceColorView,
                     sourceColorState,
                     executionSucceeded,
                     beginRecordedPass,
                     endRecordedPass,
                     clearAttachments,
                     scenePhase](
                        const RenderGraphExecutionContext&) {
                        if (executionSucceeded != nullptr && !(*executionSucceeded)) {
                            return;
                        }

                        const FrameExecutionContext executionContext(
                            renderContext,
                            surfaceTemplate,
                            *sceneData,
                            hasSourceSurface ? &sourceSurface : nullptr,
                            sourceColorView,
                            sourceColorState);
                        const RenderPassContext passContext =
                            BuildRenderPassContext(executionContext);
                        if (!beginRecordedPass(passContext, clearAttachments)) {
                            return;
                        }

                        const ScenePhaseExecutionContext scenePhaseExecutionContext =
                            BuildScenePhaseExecutionContext(executionContext, scenePhase);
                        const bool renderResult =
                            ExecuteBuiltinScenePhase(scenePhaseExecutionContext);
                        endRecordedPass(passContext);
                        if (executionSucceeded != nullptr) {
                            *executionSucceeded = renderResult;
                        }
                    });
            });
        clearAttachments = false;
    }

    return true;
}

bool BuiltinForwardPipeline::Render(
    const FrameExecutionContext& executionContext) {
    return ExecuteForwardSceneFrame(executionContext, true);
}

bool BuiltinForwardPipeline::ExecuteForwardSceneFrame(
    const FrameExecutionContext& executionContext,
    bool manageMainDirectionalShadowTransitions) {
    if (!Initialize(executionContext.renderContext)) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinForwardPipeline::Render failed: Initialize returned false");
        return false;
    }

    if (!m_forwardSceneFeatureHost.Prepare(executionContext)) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinForwardPipeline::Render failed: SceneRenderFeatureHost::Prepare returned false");
        return false;
    }

    const RenderPassContext passContext = BuildRenderPassContext(executionContext);

    if (!BeginForwardScenePass(passContext)) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinForwardPipeline::Render failed: BeginForwardScenePass returned false");
        return false;
    }

    const bool sampledDirectionalShadow =
        manageMainDirectionalShadowTransitions &&
        ShouldSampleMainDirectionalShadowMap(executionContext.sceneData);
    if (sampledDirectionalShadow) {
        TransitionMainDirectionalShadowForSampling(
            executionContext.renderContext,
            executionContext.sceneData);
    }

    const bool renderResult = ExecuteForwardScene(executionContext);

    if (sampledDirectionalShadow) {
        RestoreMainDirectionalShadowAfterSampling(
            executionContext.renderContext,
            executionContext.sceneData);
    }
    EndForwardScenePass(passContext);

    return renderResult;
}

bool BuiltinForwardPipeline::BeginForwardScenePass(
    const RenderPassContext& passContext,
    bool clearAttachments) {
    const RenderContext& context = passContext.renderContext;
    const RenderSurface& surface = passContext.surface;
    const RenderSceneData& sceneData = passContext.sceneData;
    RHI::RHIResourceView* depthAttachment = surface.GetDepthAttachment();

    std::vector<RHI::RHIResourceView*> renderTargets = CollectSurfaceColorAttachments(surface);
    if (renderTargets.empty()) {
        return false;
    }

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

    RHI::RHICommandList* commandList = context.commandList;
    if (surface.IsAutoTransitionEnabled()) {
        for (RHI::RHIResourceView* renderTarget : renderTargets) {
            if (renderTarget != nullptr) {
                commandList->TransitionBarrier(
                    renderTarget,
                    surface.GetColorStateBefore(),
                    RHI::ResourceStates::RenderTarget);
            }
        }

        if (depthAttachment != nullptr) {
            commandList->TransitionBarrier(
                depthAttachment,
                surface.GetDepthStateBefore(),
                RHI::ResourceStates::DepthWrite);
        }
    }

    commandList->SetRenderTargets(
        static_cast<uint32_t>(renderTargets.size()),
        renderTargets.data(),
        depthAttachment);

    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
    };
    const RHI::Rect clearRects[] = { scissorRect };
    commandList->SetViewport(viewport);
    commandList->SetScissorRect(scissorRect);

    const Math::Color clearColor = surface.HasClearColorOverride()
        ? surface.GetClearColorOverride()
        : sceneData.cameraData.clearColor;
    const float clearValues[4] = { clearColor.r, clearColor.g, clearColor.b, clearColor.a };
    if (clearAttachments &&
        HasRenderClearFlag(sceneData.cameraData.clearFlags, RenderClearFlags::Color) &&
        !HasSkybox(sceneData)) {
        for (RHI::RHIResourceView* renderTarget : renderTargets) {
            if (renderTarget != nullptr) {
                commandList->ClearRenderTarget(renderTarget, clearValues, 1, clearRects);
            }
        }
    }
    if (clearAttachments &&
        depthAttachment != nullptr &&
        HasRenderClearFlag(sceneData.cameraData.clearFlags, RenderClearFlags::Depth)) {
        commandList->ClearDepthStencil(depthAttachment, 1.0f, 0, 1, clearRects);
    }

    commandList->SetPrimitiveTopology(RHI::PrimitiveTopology::TriangleList);
    return true;
}

bool BuiltinForwardPipeline::Render(
    const RenderContext& context,
    const RenderSurface& surface,
    const RenderSceneData& sceneData) {
    return Render(FrameExecutionContext(context, surface, sceneData));
}

void BuiltinForwardPipeline::EndForwardScenePass(const RenderPassContext& passContext) {
    const RenderContext& context = passContext.renderContext;
    const RenderSurface& surface = passContext.surface;
    RHI::RHIResourceView* depthAttachment = surface.GetDepthAttachment();
    std::vector<RHI::RHIResourceView*> renderTargets = CollectSurfaceColorAttachments(surface);
    RHI::RHICommandList* commandList = context.commandList;

    commandList->EndRenderPass();
    if (!surface.IsAutoTransitionEnabled()) {
        return;
    }

    for (RHI::RHIResourceView* renderTarget : renderTargets) {
        if (renderTarget != nullptr) {
            commandList->TransitionBarrier(
                renderTarget,
                RHI::ResourceStates::RenderTarget,
                surface.GetColorStateAfter());
        }
    }

    if (depthAttachment != nullptr) {
        commandList->TransitionBarrier(
            depthAttachment,
            RHI::ResourceStates::DepthWrite,
            surface.GetDepthStateAfter());
    }
}

} // namespace Pipelines
} // namespace Rendering
} // namespace XCEngine