#include "Rendering/Execution/SceneRenderer.h"

#include "Components/CameraComponent.h"
#include "Debug/Logger.h"
#include "Rendering/Caches/FullscreenPassSurfaceCache.h"
#include "Rendering/Planning/CameraPostProcessPassFactory.h"
#include "Rendering/Planning/FinalColorPassFactory.h"
#include "Rendering/Planning/SceneRenderRequestUtils.h"
#include "Rendering/RenderPipelineAsset.h"

namespace XCEngine {
namespace Rendering {

namespace {

bool CompareCameraFramePlans(
    const CameraFramePlan& lhs,
    const CameraFramePlan& rhs) {
    return SceneRenderRequestUtils::CompareCameraRenderRequests(
        lhs.request,
        rhs.request);
}

RenderSurface ConfigureFullscreenStageSurface(
    const FullscreenPassSurfaceCache::SurfaceEntry& entry,
    const RenderSurface& templateSurface,
    bool copyDepthAttachment) {
    RenderSurface surface = entry.surface;
    if (copyDepthAttachment) {
        surface.SetDepthAttachment(templateSurface.GetDepthAttachment());
        surface.SetDepthStateBefore(templateSurface.GetDepthStateBefore());
        surface.SetDepthStateAfter(templateSurface.GetDepthStateAfter());
        if (templateSurface.HasClearColorOverride()) {
            surface.SetClearColorOverride(templateSurface.GetClearColorOverride());
        }
    }

    if (templateSurface.HasCustomRenderArea()) {
        surface.SetRenderArea(templateSurface.GetRenderArea());
    } else {
        surface.ResetRenderArea();
    }

    surface.SetColorStateBefore(entry.currentColorState);
    surface.SetColorStateAfter(RHI::ResourceStates::PixelShaderResource);
    return surface;
}

void UpdateTrackedFullscreenSurfaceState(
    std::vector<std::unique_ptr<FullscreenPassSurfaceCache>>& surfaceCaches,
    const RenderSurface* surface) {
    if (surface == nullptr || surface->GetColorAttachments().empty()) {
        return;
    }

    RHI::RHIResourceView* colorAttachment = surface->GetColorAttachments()[0];
    if (colorAttachment == nullptr) {
        return;
    }

    for (const std::unique_ptr<FullscreenPassSurfaceCache>& cache : surfaceCaches) {
        if (cache == nullptr) {
            continue;
        }

        for (size_t entryIndex = 0; entryIndex < cache->GetSurfaceCount(); ++entryIndex) {
            FullscreenPassSurfaceCache::SurfaceEntry* entry = cache->GetSurfaceEntry(entryIndex);
            if (entry == nullptr || entry->renderTargetView != colorAttachment) {
                continue;
            }

            entry->currentColorState = surface->GetColorStateAfter();
            return;
        }
    }
}

} // namespace

SceneRenderer::SceneRenderer() = default;

SceneRenderer::SceneRenderer(std::unique_ptr<RenderPipeline> pipeline)
    : m_cameraRenderer(std::move(pipeline)) {
}

SceneRenderer::SceneRenderer(std::shared_ptr<const RenderPipelineAsset> pipelineAsset)
    : m_cameraRenderer(std::move(pipelineAsset)) {
}

SceneRenderer::~SceneRenderer() = default;

void SceneRenderer::SetPipeline(std::unique_ptr<RenderPipeline> pipeline) {
    m_cameraRenderer.SetPipeline(std::move(pipeline));
}

void SceneRenderer::SetPipelineAsset(std::shared_ptr<const RenderPipelineAsset> pipelineAsset) {
    m_cameraRenderer.SetPipelineAsset(std::move(pipelineAsset));
}

std::vector<CameraRenderRequest> SceneRenderer::BuildRenderRequests(
    const Components::Scene& scene,
    Components::CameraComponent* overrideCamera,
    const RenderContext& context,
    const RenderSurface& surface) {
    const std::vector<CameraFramePlan> plans =
        BuildFramePlans(scene, overrideCamera, context, surface);
    std::vector<CameraRenderRequest> requests = {};
    requests.reserve(plans.size());
    for (const CameraFramePlan& plan : plans) {
        requests.push_back(BuildLegacyCameraRenderRequest(plan));
    }
    return requests;
}

std::vector<CameraFramePlan> SceneRenderer::BuildFramePlans(
    const Components::Scene& scene,
    Components::CameraComponent* overrideCamera,
    const RenderContext& context,
    const RenderSurface& surface) {
    const std::vector<CameraRenderRequest> requests =
        m_requestPlanner.BuildRequests(scene, overrideCamera, context, surface);
    std::vector<CameraFramePlan> plans = CreateFramePlansFromLegacyRequests(requests);
    ResolveCameraFinalColorPolicies(plans);
    AttachFullscreenStageRequests(context, plans);
    return plans;
}

bool SceneRenderer::Render(const CameraRenderRequest& request) {
    return Render(CameraFramePlan::FromRequest(request));
}

bool SceneRenderer::Render(const std::vector<CameraRenderRequest>& requests) {
    std::vector<CameraFramePlan> plans = CreateFramePlansFromLegacyRequests(requests);
    return Render(plans);
}

bool SceneRenderer::Render(const CameraFramePlan& plan) {
    return m_cameraRenderer.Render(plan);
}

bool SceneRenderer::Render(const std::vector<CameraFramePlan>& plans) {
    if (plans.empty()) {
        return false;
    }

    for (const CameraFramePlan& plan : plans) {
        if (!plan.IsValid()) {
            return false;
        }
    }

    std::vector<CameraFramePlan> sortedPlans = plans;
    std::stable_sort(
        sortedPlans.begin(),
        sortedPlans.end(),
        CompareCameraFramePlans);

    bool rendered = false;
    for (const CameraFramePlan& plan : sortedPlans) {
        if (!m_cameraRenderer.Render(plan)) {
            return false;
        }

        UpdateTrackedFullscreenSurfaceState(
            m_ownedFullscreenStageSurfaces,
            &plan.GetMainSceneSurface());
        if (plan.postProcess.IsRequested()) {
            UpdateTrackedFullscreenSurfaceState(
                m_ownedFullscreenStageSurfaces,
                &plan.postProcess.destinationSurface);
        }

        rendered = true;
    }

    return rendered;
}

bool SceneRenderer::Render(
    const Components::Scene& scene,
    Components::CameraComponent* overrideCamera,
    const RenderContext& context,
    const RenderSurface& surface) {
    return Render(BuildFramePlans(scene, overrideCamera, context, surface));
}

std::vector<CameraFramePlan> SceneRenderer::CreateFramePlansFromLegacyRequests(
    const std::vector<CameraRenderRequest>& requests) const {
    std::vector<CameraFramePlan> plans = {};
    plans.reserve(requests.size());
    for (const CameraRenderRequest& request : requests) {
        plans.push_back(CameraFramePlan::FromRequest(request));
    }

    return plans;
}

void SceneRenderer::PrepareOwnedFullscreenStageState(size_t requestCount) {
    m_ownedPostProcessSequences.clear();
    m_ownedPostProcessSequences.resize(requestCount);
    m_ownedFinalOutputSequences.clear();
    m_ownedFinalOutputSequences.resize(requestCount);

    if (m_ownedFullscreenStageSurfaces.size() < requestCount) {
        m_ownedFullscreenStageSurfaces.resize(requestCount);
    }

    for (size_t index = 0; index < requestCount; ++index) {
        if (m_ownedFullscreenStageSurfaces[index] == nullptr) {
            m_ownedFullscreenStageSurfaces[index] = std::make_unique<FullscreenPassSurfaceCache>();
        }
    }
}

void SceneRenderer::ResolveCameraFinalColorPolicies(
    std::vector<CameraFramePlan>& plans) const {
    const RenderPipelineAsset* pipelineAsset = GetPipelineAsset();
    const FinalColorSettings pipelineDefaults =
        pipelineAsset != nullptr ? pipelineAsset->GetDefaultFinalColorSettings() : FinalColorSettings();

    for (CameraFramePlan& plan : plans) {
        if (plan.request.camera == nullptr) {
            continue;
        }

        plan.finalColorPolicy = ResolveFinalColorPolicy(
            pipelineDefaults,
            &plan.request.camera->GetFinalColorOverrides());
    }
}

void SceneRenderer::AttachFullscreenStageRequests(
    const RenderContext& context,
    std::vector<CameraFramePlan>& plans) {
    PrepareOwnedFullscreenStageState(plans.size());

    for (size_t index = 0; index < plans.size(); ++index) {
        CameraFramePlan& plan = plans[index];
        if (plan.request.camera == nullptr ||
            plan.request.context.device == nullptr ||
            !HasValidColorTarget(plan.request.surface)) {
            continue;
        }

        std::unique_ptr<RenderPassSequence> postProcessSequence =
            BuildCameraPostProcessPassSequence(plan.request.camera->GetPostProcessPasses());
        std::unique_ptr<RenderPassSequence> finalOutputSequence =
            BuildFinalColorPassSequence(plan.finalColorPolicy);

        const bool hasPostProcess =
            postProcessSequence != nullptr && postProcessSequence->GetPassCount() > 0u;
        const bool hasFinalOutput =
            finalOutputSequence != nullptr && finalOutputSequence->GetPassCount() > 0u;
        if (!hasPostProcess && !hasFinalOutput) {
            continue;
        }

        if (plan.request.surface.GetSampleCount() > 1u) {
            Debug::Logger::Get().Error(
                Debug::LogCategory::Rendering,
                "SceneRenderer fullscreen post-process/final-output chain currently requires a single-sample main scene surface");
            continue;
        }

        const std::vector<RHI::RHIResourceView*>& colorAttachments = plan.request.surface.GetColorAttachments();
        const RHI::Format colorFormat = colorAttachments[0]->GetFormat();
        if (colorFormat == RHI::Format::Unknown) {
            continue;
        }

        const size_t fullscreenSurfaceCount = hasPostProcess && hasFinalOutput ? 2u : 1u;
        FullscreenPassSurfaceCache* surfaceCache = m_ownedFullscreenStageSurfaces[index].get();
        if (surfaceCache == nullptr ||
            !surfaceCache->EnsureSurfaces(
                context,
                plan.request.surface.GetWidth(),
                plan.request.surface.GetHeight(),
                colorFormat,
                fullscreenSurfaceCount)) {
            continue;
        }

        const FullscreenPassSurfaceCache::SurfaceEntry* sceneColorEntry = surfaceCache->GetSurfaceEntry(0u);
        if (sceneColorEntry == nullptr || sceneColorEntry->shaderResourceView == nullptr) {
            continue;
        }

        const FullscreenPassSurfaceCache::SurfaceEntry* postProcessOutputEntry =
            hasPostProcess && hasFinalOutput ? surfaceCache->GetSurfaceEntry(1u) : nullptr;
        if (hasPostProcess && hasFinalOutput &&
            (postProcessOutputEntry == nullptr || postProcessOutputEntry->shaderResourceView == nullptr)) {
            continue;
        }

        if (hasPostProcess) {
            plan.postProcess.sourceSurface =
                ConfigureFullscreenStageSurface(*sceneColorEntry, plan.request.surface, true);
            plan.postProcess.sourceColorView = sceneColorEntry->shaderResourceView;
            plan.postProcess.sourceColorState = plan.postProcess.sourceSurface.GetColorStateAfter();
            plan.postProcess.destinationSurface =
                hasFinalOutput
                    ? ConfigureFullscreenStageSurface(*postProcessOutputEntry, plan.request.surface, false)
                    : plan.request.surface;
            m_ownedPostProcessSequences[index] = std::move(postProcessSequence);
            plan.postProcess.passes = m_ownedPostProcessSequences[index].get();
        }

        if (hasFinalOutput) {
            const FullscreenPassSurfaceCache::SurfaceEntry* finalOutputSourceEntry =
                hasPostProcess ? postProcessOutputEntry : sceneColorEntry;
            plan.finalOutput.sourceSurface =
                hasPostProcess
                    ? plan.postProcess.destinationSurface
                    : ConfigureFullscreenStageSurface(*sceneColorEntry, plan.request.surface, true);
            plan.finalOutput.sourceColorView = finalOutputSourceEntry->shaderResourceView;
            plan.finalOutput.sourceColorState = plan.finalOutput.sourceSurface.GetColorStateAfter();
            plan.finalOutput.destinationSurface = plan.request.surface;
            m_ownedFinalOutputSequences[index] = std::move(finalOutputSequence);
            plan.finalOutput.passes = m_ownedFinalOutputSequences[index].get();
        }
    }
}

} // namespace Rendering
} // namespace XCEngine