XCEngine/editor/src/Viewport/SceneViewportTransformGizmoFrameBuilder.h

#pragma once

#include "Core/IEditorContext.h"
#include "Core/ISceneManager.h"
#include "Core/ISelectionManager.h"
#include "SceneViewportEditorOverlayData.h"
#include "SceneViewportMoveGizmo.h"
#include "SceneViewportRotateGizmo.h"
#include "SceneViewportScaleGizmo.h"

#include <XCEngine/Components/GameObject.h>
#include <XCEngine/Components/MeshFilterComponent.h>

#include <cstdint>
#include <vector>

namespace XCEngine {
namespace Editor {

enum class SceneViewportActiveGizmoKind : uint8_t {
    None = 0,
    Move,
    Rotate,
    Scale
};

struct SceneViewportSelectionGizmoState {
    Components::GameObject* primaryObject = nullptr;
    std::vector<Components::GameObject*> selectedObjects = {};
    Math::Vector3 pivotWorldPosition = Math::Vector3::Zero();
    Math::Quaternion primaryWorldRotation = Math::Quaternion::Identity();
};

struct SceneViewportTransformGizmoFrameState {
    SceneViewportOverlayData overlay = {};
    SceneViewportSelectionGizmoState selectionState = {};
    SceneViewportMoveGizmoContext moveContext = {};
    SceneViewportRotateGizmoContext rotateContext = {};
    SceneViewportScaleGizmoContext scaleContext = {};
    SceneViewportActiveGizmoKind activeGizmoKind = SceneViewportActiveGizmoKind::None;
};

inline SceneViewportActiveGizmoKind GetActiveSceneViewportGizmoKind(
    const SceneViewportMoveGizmo& moveGizmo,
    const SceneViewportRotateGizmo& rotateGizmo,
    const SceneViewportScaleGizmo& scaleGizmo) {
    if (moveGizmo.IsActive()) {
        return SceneViewportActiveGizmoKind::Move;
    }
    if (rotateGizmo.IsActive()) {
        return SceneViewportActiveGizmoKind::Rotate;
    }
    if (scaleGizmo.IsActive()) {
        return SceneViewportActiveGizmoKind::Scale;
    }

    return SceneViewportActiveGizmoKind::None;
}

inline Math::Quaternion ComputeStableWorldRotation(const Components::GameObject* gameObject) {
    if (gameObject == nullptr || gameObject->GetTransform() == nullptr) {
        return Math::Quaternion::Identity();
    }

    const auto* transform = gameObject->GetTransform();
    Math::Quaternion worldRotation = transform->GetLocalRotation();
    for (const auto* parent = transform->GetParent();
         parent != nullptr;
         parent = parent->GetParent()) {
        worldRotation = parent->GetLocalRotation() * worldRotation;
    }

    return worldRotation.Normalized();
}

inline Math::Vector3 GetGameObjectPivotWorldPosition(const Components::GameObject* gameObject) {
    if (gameObject == nullptr || gameObject->GetTransform() == nullptr) {
        return Math::Vector3::Zero();
    }

    return gameObject->GetTransform()->GetPosition();
}

inline Math::Vector3 GetGameObjectCenterWorldPosition(const Components::GameObject* gameObject) {
    if (gameObject == nullptr || gameObject->GetTransform() == nullptr) {
        return Math::Vector3::Zero();
    }

    if (auto* meshFilter = gameObject->GetComponent<Components::MeshFilterComponent>()) {
        if (Resources::Mesh* mesh = meshFilter->GetMesh();
            mesh != nullptr && mesh->IsValid()) {
            return gameObject->GetTransform()->TransformPoint(mesh->GetBounds().center);
        }
    }

    return gameObject->GetTransform()->GetPosition();
}

inline SceneViewportSelectionGizmoState BuildSceneViewportSelectionGizmoState(
    IEditorContext& context,
    bool useCenterPivot) {
    SceneViewportSelectionGizmoState state = {};
    const uint64_t primaryEntityId = context.GetSelectionManager().GetSelectedEntity();
    if (primaryEntityId != 0) {
        state.primaryObject = context.GetSceneManager().GetEntity(primaryEntityId);
    }

    const std::vector<uint64_t>& selectedEntities = context.GetSelectionManager().GetSelectedEntities();
    state.selectedObjects.reserve(selectedEntities.size());
    for (uint64_t entityId : selectedEntities) {
        if (entityId == 0) {
            continue;
        }

        if (auto* gameObject = context.GetSceneManager().GetEntity(entityId)) {
            state.selectedObjects.push_back(gameObject);
        }
    }

    if (state.primaryObject == nullptr && !state.selectedObjects.empty()) {
        state.primaryObject = state.selectedObjects.back();
    }
    if (state.primaryObject != nullptr && state.selectedObjects.empty()) {
        state.selectedObjects.push_back(state.primaryObject);
    }
    if (state.primaryObject != nullptr) {
        state.primaryWorldRotation = ComputeStableWorldRotation(state.primaryObject);
    }
    if (state.selectedObjects.empty()) {
        return state;
    }

    if (useCenterPivot) {
        Math::Vector3 centerSum = Math::Vector3::Zero();
        for (const auto* gameObject : state.selectedObjects) {
            centerSum += GetGameObjectCenterWorldPosition(gameObject);
        }
        state.pivotWorldPosition = centerSum / static_cast<float>(state.selectedObjects.size());
    } else {
        state.pivotWorldPosition = GetGameObjectPivotWorldPosition(state.primaryObject);
    }

    return state;
}

inline SceneViewportMoveGizmoContext BuildMoveGizmoContext(
    const SceneViewportSelectionGizmoState& selectionState,
    const SceneViewportOverlayData& overlay,
    const Math::Vector2& viewportSize,
    const Math::Vector2& mousePosition,
    bool localSpace) {
    SceneViewportMoveGizmoContext gizmoContext = {};
    gizmoContext.overlay = overlay;
    gizmoContext.viewportSize = viewportSize;
    gizmoContext.mousePosition = mousePosition;
    gizmoContext.selectedObject = selectionState.primaryObject;
    gizmoContext.selectedObjects = selectionState.selectedObjects;
    gizmoContext.pivotWorldPosition = selectionState.pivotWorldPosition;
    gizmoContext.axisOrientation = localSpace
        ? selectionState.primaryWorldRotation
        : Math::Quaternion::Identity();

    return gizmoContext;
}

inline SceneViewportRotateGizmoContext BuildRotateGizmoContext(
    const SceneViewportSelectionGizmoState& selectionState,
    const SceneViewportOverlayData& overlay,
    const Math::Vector2& viewportSize,
    const Math::Vector2& mousePosition,
    bool localSpace,
    bool rotateAroundSharedPivot) {
    SceneViewportRotateGizmoContext gizmoContext = {};
    gizmoContext.overlay = overlay;
    gizmoContext.viewportSize = viewportSize;
    gizmoContext.mousePosition = mousePosition;
    gizmoContext.selectedObject = selectionState.primaryObject;
    gizmoContext.selectedObjects = selectionState.selectedObjects;
    gizmoContext.pivotWorldPosition = selectionState.pivotWorldPosition;
    gizmoContext.axisOrientation = localSpace
        ? selectionState.primaryWorldRotation
        : Math::Quaternion::Identity();
    gizmoContext.localSpace = localSpace;
    gizmoContext.rotateAroundSharedPivot = rotateAroundSharedPivot;

    return gizmoContext;
}

inline SceneViewportScaleGizmoContext BuildScaleGizmoContext(
    const SceneViewportSelectionGizmoState& selectionState,
    const SceneViewportOverlayData& overlay,
    const Math::Vector2& viewportSize,
    const Math::Vector2& mousePosition,
    bool localSpace) {
    SceneViewportScaleGizmoContext gizmoContext = {};
    gizmoContext.overlay = overlay;
    gizmoContext.viewportSize = viewportSize;
    gizmoContext.mousePosition = mousePosition;
    gizmoContext.selectedObject = selectionState.primaryObject;
    gizmoContext.pivotWorldPosition = selectionState.pivotWorldPosition;
    gizmoContext.axisOrientation = localSpace
        ? selectionState.primaryWorldRotation
        : Math::Quaternion::Identity();

    return gizmoContext;
}

inline void CancelSceneViewportTransformGizmoDrags(
    IEditorContext& context,
    SceneViewportMoveGizmo& moveGizmo,
    SceneViewportRotateGizmo& rotateGizmo,
    SceneViewportScaleGizmo& scaleGizmo) {
    if (moveGizmo.IsActive()) {
        moveGizmo.CancelDrag(&context.GetUndoManager());
    }
    if (rotateGizmo.IsActive()) {
        rotateGizmo.CancelDrag(&context.GetUndoManager());
    }
    if (scaleGizmo.IsActive()) {
        scaleGizmo.CancelDrag(&context.GetUndoManager());
    }
}

inline SceneViewportTransformGizmoFrameState RefreshSceneViewportTransformGizmos(
    IEditorContext& context,
    const SceneViewportOverlayData& overlay,
    const Math::Vector2& viewportSize,
    const Math::Vector2& mousePosition,
    bool useCenterPivot,
    bool localSpace,
    bool usingTransformTool,
    bool showingMoveGizmo,
    SceneViewportMoveGizmo& moveGizmo,
    bool showingRotateGizmo,
    SceneViewportRotateGizmo& rotateGizmo,
    bool showingScaleGizmo,
    SceneViewportScaleGizmo& scaleGizmo) {
    SceneViewportTransformGizmoFrameState state = {};
    state.overlay = overlay;
    state.selectionState = BuildSceneViewportSelectionGizmoState(context, useCenterPivot);

    if (showingMoveGizmo) {
        state.moveContext = BuildMoveGizmoContext(
            state.selectionState,
            overlay,
            viewportSize,
            mousePosition,
            localSpace);
        if (moveGizmo.IsActive() &&
            (state.moveContext.selectedObject == nullptr ||
             context.GetSelectionManager().GetSelectedEntity() != moveGizmo.GetActiveEntityId())) {
            moveGizmo.CancelDrag(&context.GetUndoManager());
        }
    } else if (moveGizmo.IsActive()) {
        moveGizmo.CancelDrag(&context.GetUndoManager());
    }

    if (showingRotateGizmo) {
        state.rotateContext = BuildRotateGizmoContext(
            state.selectionState,
            overlay,
            viewportSize,
            mousePosition,
            localSpace,
            useCenterPivot);
        if (rotateGizmo.IsActive() &&
            (state.rotateContext.selectedObject == nullptr ||
             context.GetSelectionManager().GetSelectedEntity() != rotateGizmo.GetActiveEntityId())) {
            rotateGizmo.CancelDrag(&context.GetUndoManager());
        }
    } else if (rotateGizmo.IsActive()) {
        rotateGizmo.CancelDrag(&context.GetUndoManager());
    }

    if (showingScaleGizmo) {
        state.scaleContext = BuildScaleGizmoContext(
            state.selectionState,
            overlay,
            viewportSize,
            mousePosition,
            localSpace);
        state.scaleContext.uniformOnly = usingTransformTool;
        if (scaleGizmo.IsActive() &&
            (state.scaleContext.selectedObject == nullptr ||
             context.GetSelectionManager().GetSelectedEntity() != scaleGizmo.GetActiveEntityId())) {
            scaleGizmo.CancelDrag(&context.GetUndoManager());
        }
    } else if (scaleGizmo.IsActive()) {
        scaleGizmo.CancelDrag(&context.GetUndoManager());
    }

    state.activeGizmoKind = GetActiveSceneViewportGizmoKind(moveGizmo, rotateGizmo, scaleGizmo);

    if (showingMoveGizmo) {
        SceneViewportMoveGizmoContext updateContext = state.moveContext;
        if (state.activeGizmoKind != SceneViewportActiveGizmoKind::None &&
            state.activeGizmoKind != SceneViewportActiveGizmoKind::Move) {
            updateContext.mousePosition = Math::Vector2(-1.0f, -1.0f);
        }
        moveGizmo.Update(updateContext);
    }
    if (showingRotateGizmo) {
        SceneViewportRotateGizmoContext updateContext = state.rotateContext;
        if (state.activeGizmoKind != SceneViewportActiveGizmoKind::None &&
            state.activeGizmoKind != SceneViewportActiveGizmoKind::Rotate) {
            updateContext.mousePosition = Math::Vector2(-1.0f, -1.0f);
        }
        rotateGizmo.Update(updateContext);
    }
    if (showingScaleGizmo) {
        SceneViewportScaleGizmoContext updateContext = state.scaleContext;
        if (state.activeGizmoKind != SceneViewportActiveGizmoKind::None &&
            state.activeGizmoKind != SceneViewportActiveGizmoKind::Scale) {
            updateContext.mousePosition = Math::Vector2(-1.0f, -1.0f);
        }
        scaleGizmo.Update(updateContext);
    }

    return state;
}

} // namespace Editor
} // namespace XCEngine