#include "ViewportHostService.h"

#include "Host/ViewportRenderHost.h"
#include <Rendering/D3D12/D3D12ShaderResourceDescriptorAllocator.h>

#include <XCEngine/Components/CameraComponent.h>
#include <XCEngine/RHI/RHICommandList.h>
#include <XCEngine/Rendering/Execution/SceneRenderer.h>

#include <utility>

namespace XCEngine::UI::Editor::App {

namespace {

using ::XCEngine::RHI::ResourceStates;

void SetViewportStatusIfEmpty(
    std::string& statusText,
    std::string_view message) {
    if (statusText.empty()) {
        statusText = std::string(message);
    }
}

} // namespace

ViewportHostService::ViewportHostService()
    : m_textureDescriptorAllocator(
        std::make_unique<Host::D3D12ShaderResourceDescriptorAllocator>()) {
}

ViewportHostService::~ViewportHostService() = default;

void ViewportHostService::AttachWindowRenderer(
    Host::ViewportRenderHost& windowRenderer) {
    if (m_windowRenderer == &windowRenderer) {
        m_device = windowRenderer.GetRHIDevice();
        if (m_device != nullptr &&
            !m_textureDescriptorAllocator->IsInitialized()) {
            m_textureDescriptorAllocator->Initialize(*m_device);
        }
        return;
    }

    Shutdown();
    m_windowRenderer = &windowRenderer;
    m_device = windowRenderer.GetRHIDevice();
    if (m_device != nullptr) {
        m_textureDescriptorAllocator->Initialize(*m_device);
    }
}

void ViewportHostService::DetachWindowRenderer() {
    Shutdown();
}

void ViewportHostService::SetSurfacePresentationEnabled(bool enabled) {
    m_surfacePresentationEnabled = enabled;
}

void ViewportHostService::SetSceneViewportRenderRequest(
    SceneViewportRenderRequest request) {
    m_sceneViewportRenderRequest = request;
}

void ViewportHostService::Shutdown() {
    for (ViewportEntry& entry : m_entries) {
        DestroyViewportEntry(entry);
    }

    m_sceneViewportRenderPassBundle.Shutdown();
    m_textureDescriptorAllocator->Shutdown();
    m_windowRenderer = nullptr;
    m_device = nullptr;
    m_surfacePresentationEnabled = false;
    m_sceneViewportRenderRequest = {};
    m_sceneRenderer.reset();
    m_sceneViewportLastRenderContext = {};
}

void ViewportHostService::BeginFrame() {
    for (ViewportEntry& entry : m_entries) {
        entry.requestedWidth = 0;
        entry.requestedHeight = 0;
        entry.requestedThisFrame = false;
        entry.renderedThisFrame = false;
        entry.kind = (&entry == &m_entries[0]) ? ViewportKind::Scene : ViewportKind::Game;
    }
}

ViewportHostService::ViewportEntry& ViewportHostService::GetEntry(
    ViewportKind kind) {
    const std::size_t index = kind == ViewportKind::Scene ? 0u : 1u;
    ViewportEntry& entry = m_entries[index];
    entry.kind = kind;
    return entry;
}

const ViewportHostService::ViewportEntry& ViewportHostService::GetEntry(
    ViewportKind kind) const {
    const std::size_t index = kind == ViewportKind::Scene ? 0u : 1u;
    return m_entries[index];
}

void ViewportHostService::DestroyViewportEntry(ViewportEntry& entry) {
    m_renderTargetManager.DestroyTargets(
        m_textureDescriptorAllocator.get(),
        entry.renderTargets);
    entry = {};
}

void ViewportHostService::EnsureSceneRenderer() {
    if (!m_sceneRenderer) {
        m_sceneRenderer = std::make_unique<::XCEngine::Rendering::SceneRenderer>();
    }
}

ViewportFrame ViewportHostService::RequestViewport(
    ViewportKind kind,
    const ::XCEngine::UI::UISize& requestedSize) {
    ViewportEntry& entry = GetEntry(kind);
    entry.requestedThisFrame = requestedSize.width > 1.0f && requestedSize.height > 1.0f;
    entry.requestedWidth = entry.requestedThisFrame
        ? static_cast<std::uint32_t>(requestedSize.width)
        : 0u;
    entry.requestedHeight = entry.requestedThisFrame
        ? static_cast<std::uint32_t>(requestedSize.height)
        : 0u;

    if (!entry.requestedThisFrame) {
        return BuildFrame(entry, requestedSize);
    }

    if (m_windowRenderer == nullptr || m_device == nullptr) {
        return BuildFrame(entry, requestedSize);
    }

    if (!EnsureViewportResources(entry)) {
        return BuildFrame(entry, requestedSize);
    }

    return BuildFrame(entry, requestedSize);
}

ViewportObjectIdPickResult ViewportHostService::PickSceneViewportObject(
    const ::XCEngine::UI::UISize& viewportSize,
    const ::XCEngine::UI::UIPoint& viewportMousePosition) {
    if (!m_sceneViewportRenderRequest.IsValid()) {
        return {};
    }

    ViewportEntry& entry = GetEntry(ViewportKind::Scene);
    const ViewportObjectIdPickResult objectIdPick =
        PickSceneViewportObjectWithObjectId(
            entry,
            viewportSize,
            viewportMousePosition);
    if (objectIdPick.status == ViewportObjectIdPickStatus::ReadbackFailed) {
        SetViewportStatusIfEmpty(
            entry.statusText,
            "Scene object id readback failed");
    }
    return objectIdPick;
}

void ViewportHostService::RenderRequestedViewports(
    const ::XCEngine::Rendering::RenderContext& renderContext) {
    if (m_windowRenderer == nullptr || m_device == nullptr || !renderContext.IsValid()) {
        return;
    }

    m_sceneViewportLastRenderContext = renderContext;

    for (ViewportEntry& entry : m_entries) {
        if (!entry.requestedThisFrame || !EnsureViewportResources(entry)) {
            continue;
        }

        if (entry.kind == ViewportKind::Scene) {
            RenderSceneViewport(entry, renderContext);
        } else {
            entry.statusText.clear();
            ClearViewport(entry, renderContext, 0.09f, 0.09f, 0.09f, 1.0f);
        }

        entry.renderedThisFrame = true;
    }
}

bool ViewportHostService::EnsureViewportResources(ViewportEntry& entry) {
    const ViewportResourceReuseQuery reuseQuery =
        BuildViewportRenderTargetsReuseQuery(
            entry.kind,
            entry.renderTargets,
            entry.requestedWidth,
            entry.requestedHeight);
    if (CanReuseViewportResources(reuseQuery)) {
        return true;
    }

    if (m_windowRenderer == nullptr ||
        m_device == nullptr ||
        entry.requestedWidth == 0u ||
        entry.requestedHeight == 0u) {
        return false;
    }

    return m_renderTargetManager.EnsureTargets(
        entry.kind,
        entry.requestedWidth,
        entry.requestedHeight,
        *m_device,
        *m_textureDescriptorAllocator,
        entry.renderTargets);
}

bool ViewportHostService::RenderSceneViewport(
    ViewportEntry& entry,
    const ::XCEngine::Rendering::RenderContext& renderContext) {
    if (m_sceneViewportRenderRequest.camera == nullptr) {
        ApplySceneViewportFallback(
            entry,
            renderContext,
            "Scene view camera is unavailable",
            0.18f,
            0.07f,
            0.07f,
            1.0f);
        return false;
    }

    if (m_sceneViewportRenderRequest.scene == nullptr) {
        ApplySceneViewportFallback(
            entry,
            renderContext,
            "No active scene",
            0.07f,
            0.08f,
            0.10f,
            1.0f);
        return false;
    }

    EnsureSceneRenderer();
    entry.statusText.clear();

    ::XCEngine::Rendering::RenderSurface surface =
        BuildViewportColorSurface(entry.renderTargets);
    std::vector<::XCEngine::Rendering::CameraFramePlan> plans =
        m_sceneRenderer->BuildFramePlans(
            *m_sceneViewportRenderRequest.scene,
            m_sceneViewportRenderRequest.camera,
            renderContext,
            surface);
    if (plans.empty()) {
        ApplySceneViewportFallback(
            entry,
            renderContext,
            "Scene renderer failed",
            0.18f,
            0.07f,
            0.07f,
            1.0f);
        return false;
    }

    SceneViewportRenderPlanBuildResult renderPlan =
        m_sceneViewportRenderPassBundle.BuildRenderPlan(
            entry.renderTargets,
            m_sceneViewportRenderRequest);
    ApplySceneViewportRenderPlan(
        entry.renderTargets,
        renderPlan.plan,
        plans.front());
    if (renderPlan.warningStatusText != nullptr) {
        SetViewportStatusIfEmpty(entry.statusText, renderPlan.warningStatusText);
    }

    if (!m_sceneRenderer->Render(plans)) {
        ApplySceneViewportFallback(
            entry,
            renderContext,
            "Scene renderer failed",
            0.18f,
            0.07f,
            0.07f,
            1.0f);
        return false;
    }

    MarkSceneViewportRenderSuccess(
        entry.renderTargets,
        renderPlan.plan,
        plans.front());
    return true;
}

ViewportObjectIdPickResult ViewportHostService::PickSceneViewportObjectWithObjectId(
    ViewportEntry& entry,
    const ::XCEngine::UI::UISize& viewportSize,
    const ::XCEngine::UI::UIPoint& viewportMousePosition) {
    if (m_device == nullptr) {
        return {};
    }

    ViewportObjectIdPickContext pickContext = {};
    pickContext.commandQueue = m_sceneViewportLastRenderContext.commandQueue;
    pickContext.texture = entry.renderTargets.objectIdTexture;
    pickContext.textureState = entry.renderTargets.objectIdState;
    pickContext.textureWidth = entry.renderTargets.width;
    pickContext.textureHeight = entry.renderTargets.height;
    pickContext.hasValidFrame = entry.renderTargets.hasValidObjectIdFrame;
    pickContext.viewportSize = viewportSize;
    pickContext.viewportMousePosition = viewportMousePosition;

    return PickViewportObjectIdEntity(
        pickContext,
        [this](
            const ViewportObjectIdReadbackRequest& request,
            std::array<std::uint8_t, 4>& outRgba) {
            return m_device != nullptr &&
                m_device->ReadTexturePixelRGBA8(
                    request.commandQueue,
                    request.texture,
                    request.textureState,
                    request.pixelX,
                    request.pixelY,
                    outRgba);
        });
}

void ViewportHostService::ApplySceneViewportFallback(
    ViewportEntry& entry,
    const ::XCEngine::Rendering::RenderContext& renderContext,
    std::string statusText,
    float r,
    float g,
    float b,
    float a) {
    entry.statusText = std::move(statusText);
    entry.renderTargets.hasValidObjectIdFrame = false;
    ClearViewport(entry, renderContext, r, g, b, a);
}

void ViewportHostService::ClearViewport(
    ViewportEntry& entry,
    const ::XCEngine::Rendering::RenderContext& renderContext,
    float r,
    float g,
    float b,
    float a) {
    if (renderContext.commandList == nullptr ||
        entry.renderTargets.colorView == nullptr ||
        entry.renderTargets.depthView == nullptr) {
        return;
    }

    auto* commandList = renderContext.commandList;
    auto* colorView = entry.renderTargets.colorView;
    const float clearColor[4] = { r, g, b, a };
    commandList->TransitionBarrier(
        colorView,
        entry.renderTargets.colorState,
        ResourceStates::RenderTarget);
    commandList->SetRenderTargets(1, &colorView, entry.renderTargets.depthView);
    commandList->ClearRenderTarget(colorView, clearColor);
    commandList->ClearDepthStencil(entry.renderTargets.depthView, 1.0f, 0);
    commandList->TransitionBarrier(
        colorView,
        ResourceStates::RenderTarget,
        ResourceStates::PixelShaderResource);
    entry.renderTargets.colorState = ResourceStates::PixelShaderResource;
    entry.renderTargets.objectIdState = ResourceStates::Common;
    entry.renderTargets.selectionMaskState = ResourceStates::Common;
    entry.renderTargets.hasValidObjectIdFrame = false;
}

ViewportFrame ViewportHostService::BuildFrame(
    const ViewportEntry& entry,
    const ::XCEngine::UI::UISize& requestedSize) const {
    ViewportFrame frame = {};
    frame.requestedSize = requestedSize;
    frame.renderSize = ::XCEngine::UI::UISize(
        static_cast<float>(entry.renderTargets.width),
        static_cast<float>(entry.renderTargets.height));
    frame.wasRequested = entry.requestedThisFrame;
    frame.statusText = entry.statusText;

    if (m_surfacePresentationEnabled &&
        entry.renderTargets.textureHandle.IsValid()) {
        frame.texture = entry.renderTargets.textureHandle;
        frame.hasTexture = true;
    }

    return frame;
}

} // namespace XCEngine::UI::Editor::App