UsdLayerManager/src/ui/TransformManipulator.h

#pragma once

#include <imgui.h>

#include <pxr/usd/usd/stage.h>
#include <pxr/usd/usd/timeCode.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/base/gf/vec3d.h>
#include <pxr/base/gf/vec3f.h>
#include <pxr/base/gf/matrix4d.h>
#include <pxr/base/gf/line.h>
#include <pxr/base/gf/plane.h>
#include <pxr/base/gf/ray.h>
#include <pxr/usd/usdGeom/xformCommonAPI.h>

#include "../core/commands/TransformCommand.h"

namespace UsdLayerManager {

class CommandHistory;

enum class ManipulatorMode { Select, Move, Rotate, Scale };
enum class TransformSpace  { Object, World };

class TransformManipulator {
public:
    TransformManipulator()  = default;
    ~TransformManipulator() = default;

    void            SetMode(ManipulatorMode m) { m_mode = m; }
    ManipulatorMode GetMode()            const { return m_mode; }
    void           SetTransformSpace(TransformSpace s) { m_transformSpace = s; }
    TransformSpace GetTransformSpace()             const { return m_transformSpace; }

    void SetStage(pxr::UsdStageRefPtr stage);
    void SetSelectedPrim(const pxr::SdfPath& path);
    void SetCommandHistory(CommandHistory* h) { m_commandHistory = h; }

    /// displayTime — frame used to read/evaluate transforms (timeline frame).
    /// editTime    — frame writes are authored at (Default unless Auto-Key).
    void SetTimeCodes(pxr::UsdTimeCode displayTime, pxr::UsdTimeCode editTime) {
        m_displayTime = displayTime;
        m_editTime    = editTime;
    }

    void Render(ImDrawList* dl, const pxr::GfMatrix4d& viewProj,
                const pxr::GfVec3d& pivot, const pxr::GfVec3d& cameraEye,
                const ImVec2& imagePos, int viewW, int viewH);

    bool HandleInput(const pxr::GfMatrix4d& viewProj,
                     const pxr::GfVec3d& pivot, const pxr::GfVec3d& cameraEye,
                     const ImVec2& imagePos, int viewW, int viewH,
                     bool viewportHovered);

    bool IsDragging() const { return m_isDragging; }

private:
    static float ComputeScreenFactor(const pxr::GfMatrix4d& vp,
                                     const pxr::GfVec3d& pivot,
                                     int vW, int vH, float frac = 0.15f);
    static bool  WorldToScreen(const pxr::GfVec3d& world,
                                const pxr::GfMatrix4d& vp,
                                int vW, int vH,
                                const ImVec2& imgPos, ImVec2& out);
    static float PointToSegmentDist(ImVec2 p, ImVec2 a, ImVec2 b);
    static pxr::GfRay ComputeMouseRay(const pxr::GfMatrix4d& vp,
                                       const ImVec2& mouse,
                                       const ImVec2& imgPos, int vW, int vH);

    void DrawMoveGizmo  (ImDrawList*, const pxr::GfMatrix4d&, const pxr::GfVec3d&,
                         float, const ImVec2&, int, int, const pxr::GfVec3d[3]);
    void DrawRotateGizmo(ImDrawList*, const pxr::GfMatrix4d&, const pxr::GfVec3d&,
                         float, const pxr::GfVec3d&, const ImVec2&, int, int,
                         const pxr::GfVec3d[3]);
    void DrawScaleGizmo (ImDrawList*, const pxr::GfMatrix4d&, const pxr::GfVec3d&,
                         float, const ImVec2&, int, int, const pxr::GfVec3d[3]);

    int HitTestAxes(const pxr::GfMatrix4d&, const pxr::GfVec3d&, float,
                    const ImVec2&, int, int, const ImVec2&,
                    const pxr::GfVec3d[3]) const;
    int HitTestRotateRings(const pxr::GfMatrix4d&, const pxr::GfVec3d&, float,
                           const pxr::GfVec3d&, const ImVec2&, int, int,
                           const ImVec2&, const pxr::GfVec3d[3]) const;

    void ApplyTranslate(const pxr::GfVec3d&);
    void ApplyRotate(const pxr::GfVec3f&, pxr::UsdGeomXformCommonAPI::RotationOrder);
    void ApplyScale(const pxr::GfVec3f&);
    void GetGizmoAxes(pxr::GfVec3d[3]) const;

    // -----------------------------------------------------------------------
    ManipulatorMode     m_mode           = ManipulatorMode::Select;
    TransformSpace      m_transformSpace = TransformSpace::Object;
    pxr::UsdStageRefPtr m_stage;
    pxr::SdfPath        m_primPath;
    CommandHistory*     m_commandHistory = nullptr;
    pxr::UsdTimeCode    m_displayTime = pxr::UsdTimeCode::Default();
    pxr::UsdTimeCode    m_editTime    = pxr::UsdTimeCode::Default();

    bool   m_isDragging  = false;
    int    m_dragAxis    = -1;
    int    m_hoveredAxis = -1;

    // ---- Translate drag ---------------------------------------------------
    pxr::GfLine  m_dragMoveAxisLine;
    pxr::GfVec3d m_dragMoveOriginOnAxis;

    // ---- Rotate drag (usdtweak clock-hand, extended with space-switch) ----
    //
    // Ring plane: world-space ring normal (for mouse-ray intersection + sign).
    //
    // Delta-rotation axis is the canonical axis e_a for the dragged ring in
    // BOTH spaces; only the composition order differs:
    //   Object => SetRotate(delta) * initRot
    //             delta applied in the prim's own local frame, so the dragged
    //             local axis stays fixed -> rotates purely about its own axis.
    //   World  => initRot * SetRotate(delta)
    //             initRot maps local->parent, then delta about the world axis.
    pxr::GfVec3d    m_dragRotatePlaneNormal;     // world-space
    pxr::GfVec3d    m_dragRotateDeltaAxis;       // canonical axis e_a for ring
    pxr::GfVec3d    m_dragRotateFrom;            // clock-hand at drag start
    pxr::GfMatrix4d m_dragRotateInitialRotMat;   // local->parent rot at drag start
    bool            m_dragRotateObjectSpace = true;

    // ---- Scale drag -------------------------------------------------------
    pxr::GfLine  m_dragScaleAxisLine;
    pxr::GfVec3d m_dragScaleOriginOnAxis;

    // ---- Snapshot ---------------------------------------------------------
    pxr::GfVec3d m_dragOriginalTranslate = {0,0,0};
    pxr::GfVec3f m_dragOriginalRotate    = {0,0,0};
    pxr::GfVec3f m_dragOriginalScale     = {1,1,1};
    pxr::UsdGeomXformCommonAPI::RotationOrder m_dragOriginalRotOrder =
        pxr::UsdGeomXformCommonAPI::RotationOrderXYZ;

    pxr::GfVec3d m_dragStartTranslate = {0,0,0};
    pxr::GfVec3f m_dragStartRotate    = {0,0,0};
    pxr::GfVec3f m_dragStartScale     = {1,1,1};

    // Sample presence at m_editTime, captured at drag start (before the
    // first live write) so the end-of-drag TransformCommand can undo
    // keyframe creation correctly.
    TransformCommand::SamplePresence m_dragPreEditSamples;
};

} // namespace UsdLayerManager