XCEngine/new_editor/app/Rendering/Viewport/ViewportObjectIdPicker.h

#pragma once

#include <XCEngine/RHI/RHICommandQueue.h>
#include <XCEngine/RHI/RHIEnums.h>
#include <XCEngine/RHI/RHITexture.h>
#include <XCEngine/Rendering/Picking/ObjectIdCodec.h>
#include <XCEngine/UI/Types.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <utility>

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

struct ViewportObjectIdPickContext {
    ::XCEngine::RHI::RHICommandQueue* commandQueue = nullptr;
    ::XCEngine::RHI::RHITexture* texture = nullptr;
    ::XCEngine::RHI::ResourceStates textureState =
        ::XCEngine::RHI::ResourceStates::Common;
    std::uint32_t textureWidth = 0;
    std::uint32_t textureHeight = 0;
    bool hasValidFrame = false;
    std::uint64_t frameSerial = 0u;
    ::XCEngine::UI::UISize viewportSize = {};
    ::XCEngine::UI::UIPoint viewportMousePosition = {};
};

struct ViewportObjectIdReadbackRequest {
    ::XCEngine::RHI::RHICommandQueue* commandQueue = nullptr;
    ::XCEngine::RHI::RHITexture* texture = nullptr;
    ::XCEngine::RHI::ResourceStates textureState =
        ::XCEngine::RHI::ResourceStates::Common;
    std::uint32_t pixelX = 0;
    std::uint32_t pixelY = 0;
};

enum class ViewportObjectIdPickStatus : std::uint8_t {
    Unavailable = 0,
    Success,
    ReadbackFailed
};

struct ViewportObjectIdPickResult {
    ViewportObjectIdPickStatus status = ViewportObjectIdPickStatus::Unavailable;
    std::uint64_t frameSerial = 0u;
    ::XCEngine::Rendering::RenderObjectId renderObjectId =
        ::XCEngine::Rendering::kInvalidRenderObjectId;
    std::uint64_t resolvedEntityId = 0u;

    bool HasResolvedSample() const {
        return status == ViewportObjectIdPickStatus::Success;
    }

    bool HasResolvedEntity() const {
        return HasResolvedSample() && resolvedEntityId != 0u;
    }
};

inline bool CanPickViewportObjectId(const ViewportObjectIdPickContext& context) {
    return context.commandQueue != nullptr &&
        context.texture != nullptr &&
        context.textureWidth > 0u &&
        context.textureHeight > 0u &&
        context.hasValidFrame &&
        context.frameSerial != 0u &&
        context.viewportSize.width > 1.0f &&
        context.viewportSize.height > 1.0f &&
        context.viewportMousePosition.x >= 0.0f &&
        context.viewportMousePosition.y >= 0.0f &&
        context.viewportMousePosition.x <= context.viewportSize.width &&
        context.viewportMousePosition.y <= context.viewportSize.height;
}

inline std::uint32_t ClampViewportObjectIdPixelCoordinate(
    float value,
    std::uint32_t extent) {
    if (extent == 0u) {
        return 0u;
    }

    const float maxCoordinate = static_cast<float>(extent - 1u);
    const float clamped = (std::max)(0.0f, (std::min)(value, maxCoordinate));
    return static_cast<std::uint32_t>(std::floor(clamped));
}

inline std::uint32_t ResolveViewportObjectIdPixelCoordinate(
    float viewportCoordinate,
    float viewportExtent,
    std::uint32_t textureExtent) {
    if (viewportExtent <= 0.0f || textureExtent == 0u) {
        return 0u;
    }

    const float normalized = viewportCoordinate / viewportExtent;
    const float textureCoordinate =
        normalized * static_cast<float>(textureExtent);
    return ClampViewportObjectIdPixelCoordinate(textureCoordinate, textureExtent);
}

inline bool BuildViewportObjectIdReadbackRequest(
    const ViewportObjectIdPickContext& context,
    ViewportObjectIdReadbackRequest& outRequest) {
    outRequest = {};
    if (!CanPickViewportObjectId(context)) {
        return false;
    }

    outRequest.commandQueue = context.commandQueue;
    outRequest.texture = context.texture;
    outRequest.textureState = context.textureState;
    outRequest.pixelX = ResolveViewportObjectIdPixelCoordinate(
        context.viewportMousePosition.x,
        context.viewportSize.width,
        context.textureWidth);
    outRequest.pixelY = ResolveViewportObjectIdPixelCoordinate(
        context.viewportMousePosition.y,
        context.viewportSize.height,
        context.textureHeight);
    return true;
}

template <typename ReadPixelFn>
ViewportObjectIdPickResult PickViewportRenderObjectId(
    const ViewportObjectIdPickContext& context,
    ReadPixelFn&& readPixel) {
    ViewportObjectIdPickResult result = {};

    ViewportObjectIdReadbackRequest request = {};
    if (!BuildViewportObjectIdReadbackRequest(context, request)) {
        return result;
    }

    result.frameSerial = context.frameSerial;
    std::array<std::uint8_t, 4> rgba = {};
    if (!readPixel(request, rgba)) {
        result.status = ViewportObjectIdPickStatus::ReadbackFailed;
        return result;
    }

    result.status = ViewportObjectIdPickStatus::Success;
    result.renderObjectId = ::XCEngine::Rendering::DecodeRenderObjectIdFromColor(
        rgba[0],
        rgba[1],
        rgba[2],
        rgba[3]);
    return result;
}

template <typename ReadPixelFn>
bool TryPickViewportRenderObjectId(
    const ViewportObjectIdPickContext& context,
    ReadPixelFn&& readPixel,
    ::XCEngine::Rendering::RenderObjectId& outRenderObjectId) {
    const ViewportObjectIdPickResult result =
        PickViewportRenderObjectId(context, std::forward<ReadPixelFn>(readPixel));
    outRenderObjectId = result.renderObjectId;
    return result.HasResolvedSample();
}

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