#include <XCEngine/Rendering/RenderPassGraphContract.h>

#include "Rendering/Execution/Internal/CameraFrameGraph/SurfaceUtils.h"
#include "Rendering/FrameData/RenderSceneData.h"
#include "Rendering/Graph/RenderGraph.h"

#include <algorithm>
#include <memory>

namespace XCEngine {
namespace Rendering {
namespace {

bool ResolveGraphManagedSourceSurface(
    const RenderSurface* sourceSurfaceTemplate,
    RHI::RHIResourceView* sourceColorView,
    RHI::ResourceStates sourceColorState,
    RenderGraphTextureHandle sourceColorTexture,
    const RenderGraphExecutionContext& graphContext,
    const RenderPassGraphIO& io,
    const RenderSurface*& outSourceSurface,
    RHI::RHIResourceView*& outSourceColorView,
    RHI::ResourceStates& outSourceColorState,
    RenderSurface& outGraphManagedSourceSurface) {
    outSourceSurface = sourceSurfaceTemplate;
    outSourceColorView = sourceColorView;
    outSourceColorState = sourceColorState;

    if (!io.readSourceColor ||
        !sourceColorTexture.IsValid() ||
        !graphContext.OwnsTextureTransitions(sourceColorTexture)) {
        return true;
    }

    if (graphContext.IsTransientTexture(sourceColorTexture)) {
        if (sourceSurfaceTemplate == nullptr) {
            return false;
        }

        RHI::RHIResourceView* renderTargetView =
            graphContext.ResolveTextureView(
                sourceColorTexture,
                RenderGraphTextureViewType::RenderTarget);
        RHI::RHIResourceView* shaderResourceView =
            graphContext.ResolveTextureView(
                sourceColorTexture,
                RenderGraphTextureViewType::ShaderResource);
        if (renderTargetView == nullptr || shaderResourceView == nullptr) {
            return false;
        }

        outGraphManagedSourceSurface = *sourceSurfaceTemplate;
        outGraphManagedSourceSurface.SetColorAttachment(renderTargetView);
        outGraphManagedSourceSurface.SetAutoTransitionEnabled(false);
        outGraphManagedSourceSurface.SetColorStateBefore(RHI::ResourceStates::PixelShaderResource);
        outGraphManagedSourceSurface.SetColorStateAfter(RHI::ResourceStates::PixelShaderResource);

        outSourceSurface = &outGraphManagedSourceSurface;
        outSourceColorView = shaderResourceView;
        outSourceColorState = RHI::ResourceStates::PixelShaderResource;
        return true;
    }

    if (sourceSurfaceTemplate != nullptr) {
        outGraphManagedSourceSurface =
            BuildGraphManagedImportedSurface(
                *sourceSurfaceTemplate,
                RHI::ResourceStates::PixelShaderResource,
                RHI::ResourceStates::PixelShaderResource);
        outSourceSurface = &outGraphManagedSourceSurface;
    }
    outSourceColorState = RHI::ResourceStates::PixelShaderResource;
    return true;
}

bool ResolveGraphManagedOutputSurface(
    const RenderSurface& surfaceTemplate,
    const std::vector<RenderGraphTextureHandle>& colorTargets,
    RenderGraphTextureHandle depthTarget,
    const RenderGraphExecutionContext& graphContext,
    const RenderPassGraphIO& io,
    const RenderSurface*& outSurface,
    RenderSurface& outGraphManagedSurface) {
    outSurface = &surfaceTemplate;

    const bool ownsAnyColorTransitions =
        io.writeColor &&
        std::any_of(
            colorTargets.begin(),
            colorTargets.end(),
            [&](RenderGraphTextureHandle texture) {
                return texture.IsValid() &&
                       graphContext.OwnsTextureTransitions(texture);
            });
    const bool ownsDepthTransitions =
        io.writeDepth &&
        depthTarget.IsValid() &&
        graphContext.OwnsTextureTransitions(depthTarget);
    if (!ownsAnyColorTransitions && !ownsDepthTransitions) {
        return true;
    }

    outGraphManagedSurface = BuildGraphManagedPassSurface(surfaceTemplate);

    if (ownsAnyColorTransitions) {
        std::vector<RHI::RHIResourceView*> colorAttachments =
            surfaceTemplate.GetColorAttachments();
        if (colorAttachments.size() < colorTargets.size()) {
            colorAttachments.resize(colorTargets.size(), nullptr);
        }

        for (size_t colorIndex = 0u; colorIndex < colorTargets.size(); ++colorIndex) {
            const RenderGraphTextureHandle colorTarget = colorTargets[colorIndex];
            if (!colorTarget.IsValid() ||
                !graphContext.OwnsTextureTransitions(colorTarget) ||
                !graphContext.IsTransientTexture(colorTarget)) {
                continue;
            }

            colorAttachments[colorIndex] =
                graphContext.ResolveTextureView(
                    colorTarget,
                    RenderGraphTextureViewType::RenderTarget);
            if (colorAttachments[colorIndex] == nullptr) {
                return false;
            }
        }

        outGraphManagedSurface.SetColorAttachments(colorAttachments);
    }

    if (ownsDepthTransitions &&
        graphContext.IsTransientTexture(depthTarget)) {
        RHI::RHIResourceView* depthAttachment =
            graphContext.ResolveTextureView(
                depthTarget,
                RenderGraphTextureViewType::DepthStencil);
        if (depthAttachment == nullptr) {
            return false;
        }

        outGraphManagedSurface.SetDepthAttachment(depthAttachment);
    }

    outSurface = &outGraphManagedSurface;
    return true;
}

} // namespace

bool RecordCallbackRasterRenderPass(
    const RenderPassRenderGraphContext& context,
    const RenderPassGraphIO& io,
    RenderPassGraphExecutePassCallback executePassCallback,
    std::vector<RenderGraphTextureHandle> additionalReadTextures) {
    if (!executePassCallback) {
        return false;
    }

    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 surface = context.surface;
    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 RenderGraphTextureHandle sourceColorTexture = context.sourceColorTexture;
    const std::vector<RenderGraphTextureHandle> colorTargets = context.colorTargets;
    const RenderGraphTextureHandle depthTarget = context.depthTarget;
    bool* const executionSucceeded = context.executionSucceeded;
    const RenderPassGraphBeginCallback beginPassCallback = context.beginPassCallback;
    const RenderPassGraphEndCallback endPassCallback = context.endPassCallback;

    context.graphBuilder.AddRasterPass(
        passName,
        [renderContext,
         sceneData,
         surface,
         hasSourceSurface,
         sourceSurface,
         sourceColorView,
         sourceColorState,
         sourceColorTexture,
         colorTargets,
         depthTarget,
         executionSucceeded,
         beginPassCallback,
         endPassCallback,
         executePassCallback,
         additionalReadTextures,
         io](
            RenderGraphPassBuilder& passBuilder) {
            if (io.readSourceColor && sourceColorTexture.IsValid()) {
                passBuilder.ReadTexture(sourceColorTexture);
            }

            for (RenderGraphTextureHandle readTexture : additionalReadTextures) {
                if (readTexture.IsValid()) {
                    passBuilder.ReadTexture(readTexture);
                }
            }

            if (io.writeColor) {
                for (RenderGraphTextureHandle colorTarget : colorTargets) {
                    if (colorTarget.IsValid()) {
                        passBuilder.WriteTexture(colorTarget);
                    }
                }
            }

            if (io.writeDepth && depthTarget.IsValid()) {
                passBuilder.WriteDepthTexture(depthTarget);
            }

            passBuilder.SetExecuteCallback(
                [renderContext,
                 sceneData,
                 surface,
                 hasSourceSurface,
                 sourceSurface,
                 sourceColorView,
                 sourceColorState,
                 sourceColorTexture,
                 colorTargets,
                 depthTarget,
                 executionSucceeded,
                 beginPassCallback,
                 endPassCallback,
                 executePassCallback,
                 io](
                    const RenderGraphExecutionContext& executionContext) {
                    if (executionSucceeded != nullptr && !(*executionSucceeded)) {
                        return;
                    }

                    const RenderSurface* resolvedSourceSurface =
                        hasSourceSurface ? &sourceSurface : nullptr;
                    RHI::RHIResourceView* resolvedSourceColorView = sourceColorView;
                    RHI::ResourceStates resolvedSourceColorState = sourceColorState;
                    RenderSurface graphManagedSourceSurface = {};
                    if (!ResolveGraphManagedSourceSurface(
                            hasSourceSurface ? &sourceSurface : nullptr,
                            sourceColorView,
                            sourceColorState,
                            sourceColorTexture,
                            executionContext,
                            io,
                            resolvedSourceSurface,
                            resolvedSourceColorView,
                            resolvedSourceColorState,
                            graphManagedSourceSurface)) {
                        if (executionSucceeded != nullptr) {
                            *executionSucceeded = false;
                        }
                        return;
                    }

                    const RenderSurface* resolvedSurface = &surface;
                    RenderSurface graphManagedSurface = {};
                    if (!ResolveGraphManagedOutputSurface(
                            surface,
                            colorTargets,
                            depthTarget,
                            executionContext,
                            io,
                            resolvedSurface,
                            graphManagedSurface)) {
                        if (executionSucceeded != nullptr) {
                            *executionSucceeded = false;
                        }
                        return;
                    }

                    const RenderPassContext passContext = {
                        renderContext,
                        *resolvedSurface,
                        *sceneData,
                        resolvedSourceSurface,
                        resolvedSourceColorView,
                        resolvedSourceColorState
                    };
                    if (beginPassCallback &&
                        !beginPassCallback(passContext)) {
                        if (executionSucceeded != nullptr) {
                            *executionSucceeded = false;
                        }
                        return;
                    }

                    const bool executeResult = executePassCallback(passContext);
                    if (endPassCallback) {
                        endPassCallback(passContext);
                    }
                    if (executionSucceeded != nullptr) {
                        *executionSucceeded = executeResult;
                    }
                });
        });
    return true;
}

bool RecordRasterRenderPass(
    RenderPass& pass,
    const RenderPassRenderGraphContext& context,
    const RenderPassGraphIO& io) {
    return RecordCallbackRasterRenderPass(
        context,
        io,
        [&pass](const RenderPassContext& passContext) {
            return pass.Execute(passContext);
        });
}

} // namespace Rendering
} // namespace XCEngine