#pragma once

#include <imgui.h>

#include <pxr/usd/usd/stage.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/usd/usdGeom/xformCommonAPI.h>

namespace UsdLayerManager {

class CommandHistory;

/// Active transform tool mode — mirrors Maya Q/W/E/R convention.
enum class ManipulatorMode {
    Select, ///< Q — no gizmo, normal click-to-select
    Move,   ///< W — translate along axis arrows
    Rotate, ///< E — rotate around axis rings
    Scale   ///< R — scale along axis handles
};

/// Coordinate space in which the gizmo axes are expressed.
enum class TransformSpace {
    Object, ///< Gizmo axes align with the selected prim's local axes (default)
    World,  ///< Gizmo axes are fixed world-space X/Y/Z
};

/// Maya-style interactive transform gizmo.
///
/// Rendering is done with ImGui DrawList (2-D screen-space overlay), drawn
/// AFTER ImGui::Image() for the viewport — exactly the same approach used by
/// ImGuizmo.  No OpenGL resources are needed; the FBO bind/unbind dance is
/// entirely eliminated.
///
/// Gizmo world-space size is computed each frame using ImGuizmo's screen-
/// factor formula: project a camera-aligned unit vector to clip space and
/// derive the world size that spans a fixed fraction of the screen.  This
/// gives constant apparent size regardless of camera distance or FOV.
class TransformManipulator {
public:
    TransformManipulator()  = default;
    ~TransformManipulator() = default;

    // -----------------------------------------------------------------------
    // Mode
    // -----------------------------------------------------------------------
    void            SetMode(ManipulatorMode mode) { m_mode = mode; }
    ManipulatorMode GetMode()               const { return m_mode; }

    // -----------------------------------------------------------------------
    // Transform space
    // -----------------------------------------------------------------------
    void           SetTransformSpace(TransformSpace space) { m_transformSpace = space; }
    TransformSpace GetTransformSpace()                const { return m_transformSpace; }

    // -----------------------------------------------------------------------
    // Stage / selection
    // -----------------------------------------------------------------------
    void SetStage(pxr::UsdStageRefPtr stage);
    void SetSelectedPrim(const pxr::SdfPath& path);
    void SetCommandHistory(CommandHistory* history) { m_commandHistory = history; }

    // -----------------------------------------------------------------------
    // Per-frame API (called from ViewportPanel::Render)
    // -----------------------------------------------------------------------

    /// Draw the gizmo as a 2-D overlay onto @p dl.
    /// Call AFTER ImGui::Image() so the overlay appears on top of the scene.
    /// @param dl         ImGui::GetWindowDrawList() of the Viewport window.
    /// @param viewProj   Combined view × projection matrix (USD row-major).
    /// @param pivot      World-space pivot (bounding-box centre of selection).
    /// @param cameraEye  World-space camera eye position (for half-arc orientation).
    /// @param imagePos   Screen-space top-left corner of the rendered image.
    /// @param viewW/H    Viewport pixel dimensions.
    void Render(ImDrawList*           dl,
                const pxr::GfMatrix4d& viewProj,
                const pxr::GfVec3d&   pivot,
                const pxr::GfVec3d&   cameraEye,
                const ImVec2&         imagePos,
                int viewW, int viewH);

    /// Process mouse input.  Must be called BEFORE camera-drag / prim-pick
    /// logic in ViewportPanel so the gizmo can consume LMB clicks first.
    /// @return true if the gizmo consumed the event.
    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:
    // -----------------------------------------------------------------------
    // ImGuizmo-style screen-factor computation
    // -----------------------------------------------------------------------

    /// Compute the world-space gizmo size so that the gizmo spans
    /// @p desiredFraction of the smaller viewport dimension in NDC.
    ///
    /// Algorithm (from ImGuizmo):
    ///   1. Project @p pivot to clip space.
    ///   2. Project @p pivot + each world axis unit vector to clip space.
    ///   3. Measure clip-space length (aspect-ratio corrected).
    ///   4. screenFactor = desiredFraction / maxClipLen.
    static float ComputeScreenFactor(const pxr::GfMatrix4d& viewProj,
                                     const pxr::GfVec3d&    pivot,
                                     int viewW, int viewH,
                                     float desiredFraction = 0.15f);

    // -----------------------------------------------------------------------
    // Screen-space helpers
    // -----------------------------------------------------------------------

    /// Project a world-space point to absolute screen coordinates.
    /// Returns false if the point is behind the camera (w ≤ 0).
    static bool WorldToScreen(const pxr::GfVec3d&    world,
                               const pxr::GfMatrix4d& viewProj,
                               int viewW, int viewH,
                               const ImVec2&          imagePos,
                               ImVec2&                outScreen);

    /// Distance from point @p p to line segment @p a – @p b (2-D).
    static float PointToSegmentDist(ImVec2 p, ImVec2 a, ImVec2 b);

    // -----------------------------------------------------------------------
    // Per-mode drawing  (all ImGui DrawList, screen-space)
    // -----------------------------------------------------------------------
    void DrawMoveGizmo  (ImDrawList* dl, const pxr::GfMatrix4d& vp,
                         const pxr::GfVec3d& pivot, float sf,
                         const ImVec2& imgPos, int vW, int vH,
                         const pxr::GfVec3d axes[3]);
    void DrawRotateGizmo(ImDrawList* dl, const pxr::GfMatrix4d& vp,
                         const pxr::GfVec3d& pivot, float sf,
                         const pxr::GfVec3d& cameraEye,
                         const ImVec2& imgPos, int vW, int vH,
                         const pxr::GfVec3d axes[3]);
    void DrawScaleGizmo (ImDrawList* dl, const pxr::GfMatrix4d& vp,
                         const pxr::GfVec3d& pivot, float sf,
                         const ImVec2& imgPos, int vW, int vH,
                         const pxr::GfVec3d axes[3]);

    // -----------------------------------------------------------------------
    // Hit-testing
    // -----------------------------------------------------------------------

    /// Returns axis index 0=X 1=Y 2=Z, or -1 if nothing hit (move / scale).
    int HitTestAxes(const pxr::GfMatrix4d& vp,
                    const pxr::GfVec3d& pivot, float sf,
                    const ImVec2& imgPos, int vW, int vH,
                    const ImVec2& mousePosAbsolute,
                    const pxr::GfVec3d axes[3]) const;

    /// Returns axis index 0=X 1=Y 2=Z, or -1 if nothing hit (rotate rings).
    /// Uses proximity to the VISIBLE front-facing half-arc only.
    int HitTestRotateRings(const pxr::GfMatrix4d& vp,
                           const pxr::GfVec3d& pivot, float sf,
                           const pxr::GfVec3d& cameraEye,
                           const ImVec2& imgPos, int vW, int vH,
                           const ImVec2& mousePosAbsolute,
                           const pxr::GfVec3d axes[3]) const;

    // -----------------------------------------------------------------------
    // USD transform write helpers
    // -----------------------------------------------------------------------
    void ApplyMoveDelta  (const pxr::GfVec3d& worldDelta);
    void ApplyRotateDelta(int axisIndex, float angleDeg);
    void ApplyScaleDelta (int axisIndex, float factor);

    /// Fills @p outAxes[3] with the gizmo X/Y/Z axis directions in world space.
    /// In World space: fixed unit vectors.
    /// In Object space: the prim's local axes extracted from its local-to-world matrix.
    void GetGizmoAxes(pxr::GfVec3d outAxes[3]) const;

    // -----------------------------------------------------------------------
    // State
    // -----------------------------------------------------------------------
    ManipulatorMode     m_mode = ManipulatorMode::Select;
    TransformSpace      m_transformSpace = TransformSpace::Object;
    pxr::UsdStageRefPtr m_stage;
    pxr::SdfPath        m_primPath;
    CommandHistory*     m_commandHistory = nullptr;

    // Drag state
    bool   m_isDragging  = false;
    int    m_dragAxis    = -1;
    ImVec2 m_dragLastPos = {0.f, 0.f};

    // For rotation drag: screen-angle around projected pivot center
    float  m_dragRotateLastAngle = 0.f;   ///< atan2 angle of mouse around pivot (radians)

    // Saved xform at drag START (never mutated during drag — used for undo)
    pxr::GfVec3d m_dragOriginalTranslate = {0.0, 0.0, 0.0};
    pxr::GfVec3f m_dragOriginalRotate    = {0.f, 0.f, 0.f};
    pxr::GfVec3f m_dragOriginalScale     = {1.f, 1.f, 1.f};
    pxr::UsdGeomXformCommonAPI::RotationOrder m_dragOriginalRotOrder =
        pxr::UsdGeomXformCommonAPI::RotationOrderXYZ;

    // Working accumulator for the current drag (updated each frame)
    pxr::GfVec3d m_dragStartTranslate = {0.0, 0.0, 0.0};
    pxr::GfVec3f m_dragStartRotate    = {0.f, 0.f, 0.f};
    pxr::GfVec3f m_dragStartScale     = {1.f, 1.f, 1.f};

    // Hover highlight
    int m_hoveredAxis = -1;
};

} // namespace UsdLayerManager