#include "Rendering/Execution/Internal/CameraFrameRenderGraphSurfaceUtils.h"

#include "Rendering/Execution/CameraFrameRenderGraphStagePolicy.h"
#include "Rendering/Graph/RenderGraph.h"
#include "Rendering/RenderPass.h"

#include <algorithm>
#include <string>

namespace XCEngine {
namespace Rendering {

namespace {

constexpr RHI::Format kRenderGraphImportedColorFormat = RHI::Format::R8G8B8A8_UNorm;
constexpr RHI::Format kRenderGraphImportedDepthFormat = RHI::Format::D24_UNorm_S8_UInt;

Containers::String BuildRenderGraphResourceName(
    const Containers::String& surfaceName,
    const char* slotName,
    size_t index = 0u,
    bool indexed = false) {
    std::string name = surfaceName.CStr();
    name += '.';
    name += slotName;
    if (indexed) {
        name += std::to_string(index);
    }

    return Containers::String(name.c_str());
}

RenderGraphTextureDesc BuildImportedTextureDesc(
    const RenderSurface& surface,
    RHI::Format format) {
    RenderGraphTextureDesc desc = {};
    desc.width = surface.GetWidth() > 0u ? surface.GetWidth() : surface.GetRenderAreaWidth();
    desc.height = surface.GetHeight() > 0u ? surface.GetHeight() : surface.GetRenderAreaHeight();
    desc.format = static_cast<Core::uint32>(format);
    desc.textureType = static_cast<Core::uint32>(RHI::TextureType::Texture2D);
    desc.sampleCount = std::max(surface.GetSampleCount(), 1u);
    desc.sampleQuality = surface.GetSampleQuality();
    return desc;
}

RenderGraphImportedTextureOptions BuildImportedTextureOptions(
    const RenderSurface& surface,
    bool isDepth,
    RenderGraphSurfaceImportUsage usage,
    bool graphOwnsTransitions) {
    const RHI::ResourceStates beforeState =
        isDepth ? surface.GetDepthStateBefore() : surface.GetColorStateBefore();
    const RHI::ResourceStates afterState =
        isDepth ? surface.GetDepthStateAfter() : surface.GetColorStateAfter();

    RenderGraphImportedTextureOptions options = {};
    options.initialState =
        usage == RenderGraphSurfaceImportUsage::Output
            ? beforeState
            : afterState;
    options.finalState =
        usage == RenderGraphSurfaceImportUsage::Output
            ? afterState
            : options.initialState;
    options.graphOwnsTransitions = graphOwnsTransitions;
    return options;
}

RenderGraphTextureHandle ImportRenderGraphTexture(
    RenderGraphBuilder& builder,
    RenderGraphImportedTextureRegistry& registry,
    const Containers::String& name,
    const RenderGraphTextureDesc& desc,
    RHI::RHIResourceView* view,
    const RenderGraphImportedTextureOptions& importedOptions) {
    if (view == nullptr) {
        return {};
    }

    const auto existing = registry.find(view);
    if (existing != registry.end()) {
        builder.MergeImportedTextureOptions(existing->second, importedOptions);
        return existing->second;
    }

    const RenderGraphTextureHandle handle =
        builder.ImportTexture(name, desc, view, importedOptions);
    registry.emplace(view, handle);
    return handle;
}

} // namespace

Containers::String BuildRenderGraphSequencePassName(
    const Containers::String& stageName,
    size_t passIndex) {
    std::string name = stageName.CStr();
    name += ".Pass";
    name += std::to_string(passIndex);
    return Containers::String(name.c_str());
}

RenderGraphTextureDesc BuildFullscreenTransientTextureDesc(
    const RenderSurface& surface) {
    RenderGraphTextureDesc desc = {};
    desc.width = surface.GetWidth() > 0u ? surface.GetWidth() : surface.GetRenderAreaWidth();
    desc.height = surface.GetHeight() > 0u ? surface.GetHeight() : surface.GetRenderAreaHeight();
    desc.format = static_cast<Core::uint32>(kRenderGraphImportedColorFormat);
    desc.textureType = static_cast<Core::uint32>(RHI::TextureType::Texture2D);
    desc.sampleCount = 1u;
    desc.sampleQuality = 0u;
    return desc;
}

RenderGraphImportedSurface ImportRenderGraphSurface(
    RenderGraphBuilder& builder,
    RenderGraphImportedTextureRegistry& registry,
    const Containers::String& surfaceName,
    const RenderSurface* surface,
    RenderGraphSurfaceImportUsage usage,
    bool graphOwnsColorTransitions,
    bool graphOwnsDepthTransitions) {
    RenderGraphImportedSurface importedSurface = {};
    if (surface == nullptr) {
        return importedSurface;
    }

    const RenderGraphTextureDesc colorDesc =
        BuildImportedTextureDesc(*surface, kRenderGraphImportedColorFormat);
    const std::vector<RHI::RHIResourceView*>& colorAttachments = surface->GetColorAttachments();
    importedSurface.colorTextures.reserve(colorAttachments.size());
    for (size_t colorIndex = 0; colorIndex < colorAttachments.size(); ++colorIndex) {
        RHI::RHIResourceView* colorAttachment = colorAttachments[colorIndex];
        if (colorAttachment == nullptr) {
            continue;
        }

        importedSurface.colorTextures.push_back(
            ImportRenderGraphTexture(
                builder,
                registry,
                BuildRenderGraphResourceName(surfaceName, "Color", colorIndex, true),
                colorDesc,
                colorAttachment,
                BuildImportedTextureOptions(
                    *surface,
                    false,
                    usage,
                    graphOwnsColorTransitions)));
    }

    if (RHI::RHIResourceView* depthAttachment = surface->GetDepthAttachment();
        depthAttachment != nullptr) {
        importedSurface.depthTexture =
            ImportRenderGraphTexture(
                builder,
                registry,
                BuildRenderGraphResourceName(surfaceName, "Depth"),
                BuildImportedTextureDesc(*surface, kRenderGraphImportedDepthFormat),
                depthAttachment,
                BuildImportedTextureOptions(
                    *surface,
                    true,
                    usage,
                    graphOwnsDepthTransitions));
    }

    return importedSurface;
}

RenderGraphTextureHandle GetPrimaryColorTexture(
    const RenderGraphImportedSurface& surface) {
    for (RenderGraphTextureHandle texture : surface.colorTextures) {
        if (texture.IsValid()) {
            return texture;
        }
    }

    return {};
}

void ReadRenderGraphSurface(
    RenderGraphPassBuilder& passBuilder,
    const RenderGraphImportedSurface& surface) {
    for (RenderGraphTextureHandle texture : surface.colorTextures) {
        if (texture.IsValid()) {
            passBuilder.ReadTexture(texture);
        }
    }

    if (surface.depthTexture.IsValid()) {
        passBuilder.ReadDepthTexture(surface.depthTexture);
    }
}

void ReadRenderGraphColorSurface(
    RenderGraphPassBuilder& passBuilder,
    const RenderGraphImportedSurface& surface) {
    for (RenderGraphTextureHandle texture : surface.colorTextures) {
        if (texture.IsValid()) {
            passBuilder.ReadTexture(texture);
        }
    }
}

void WriteRenderGraphSurface(
    RenderGraphPassBuilder& passBuilder,
    const RenderGraphImportedSurface& surface) {
    for (RenderGraphTextureHandle texture : surface.colorTextures) {
        if (texture.IsValid()) {
            passBuilder.WriteTexture(texture);
        }
    }

    if (surface.depthTexture.IsValid()) {
        passBuilder.WriteDepthTexture(surface.depthTexture);
    }
}

void WriteRenderGraphColorSurface(
    RenderGraphPassBuilder& passBuilder,
    const RenderGraphImportedSurface& surface) {
    for (RenderGraphTextureHandle texture : surface.colorTextures) {
        if (texture.IsValid()) {
            passBuilder.WriteTexture(texture);
        }
    }
}

RenderSurface BuildGraphManagedImportedSurface(
    const RenderSurface& templateSurface,
    RHI::ResourceStates stateBefore,
    RHI::ResourceStates stateAfter) {
    RenderSurface surface = templateSurface;
    surface.SetAutoTransitionEnabled(false);
    surface.SetColorStateBefore(stateBefore);
    surface.SetColorStateAfter(stateAfter);
    return surface;
}

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

bool CanUseGraphManagedImportedSurface(
    const RenderGraphImportedSurface& surface,
    const RenderGraphExecutionContext& graphContext) {
    bool hasAnyTexture = false;

    for (RenderGraphTextureHandle texture : surface.colorTextures) {
        if (!texture.IsValid()) {
            continue;
        }

        hasAnyTexture = true;
        if (!graphContext.OwnsTextureTransitions(texture)) {
            return false;
        }
    }

    if (surface.depthTexture.IsValid()) {
        hasAnyTexture = true;
        if (!graphContext.OwnsTextureTransitions(surface.depthTexture)) {
            return false;
        }
    }

    return hasAnyTexture;
}

RenderGraphTextureHandle ResolveStageOutputColorHandle(
    CameraFrameStage stage,
    const CameraFramePlan& plan,
    const Containers::String& stageName,
    const RenderPassContext& stagePassContext,
    const RenderGraphImportedSurface& outputSurface,
    RenderGraphBuilder& graphBuilder) {
    if (UsesCameraFrameStageGraphManagedOutputColor(plan, stage)) {
        switch (stage) {
        case CameraFrameStage::MainScene:
            return graphBuilder.CreateTransientTexture(
                stageName + ".Color",
                BuildImportedTextureDesc(
                    stagePassContext.surface,
                    kRenderGraphImportedColorFormat));
        case CameraFrameStage::PostProcess:
            return graphBuilder.CreateTransientTexture(
                stageName + ".Color",
                BuildFullscreenTransientTextureDesc(stagePassContext.surface));
        default:
            break;
        }
    }

    return GetPrimaryColorTexture(outputSurface);
}

} // namespace Rendering
} // namespace XCEngine