#include "TransformManipulator.h"
#include "../utils/Logger.h"
#include "../core/CommandHistory.h"
#include "../core/commands/TransformCommand.h"

#include <pxr/usd/usd/prim.h>
#include <pxr/usd/usd/editContext.h>
#include <pxr/usd/usdGeom/xformCommonAPI.h>
#include <pxr/usd/usdGeom/xformable.h>
#include <pxr/usd/usdGeom/xformCache.h>
#include <pxr/usd/usdGeom/xformOp.h>
#include <pxr/base/gf/matrix4d.h>
#include <pxr/base/gf/matrix4f.h>
#include <pxr/base/gf/vec3d.h>
#include <pxr/base/gf/rotation.h>
#include <pxr/base/gf/math.h>
#include <pxr/base/gf/line.h>
#include <pxr/base/gf/lineSeg.h>
#include <pxr/base/gf/plane.h>
#include <pxr/base/gf/ray.h>
#include <pxr/base/vt/value.h>

#include <cmath>
#include <algorithm>
#include <vector>
#include <memory>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

namespace UsdLayerManager {

// ---------------------------------------------------------------------------
// Colours / constants
// ---------------------------------------------------------------------------
static const ImU32 kColX      = IM_COL32(214,  38,  38, 255);
static const ImU32 kColY      = IM_COL32( 38, 179,  38, 255);
static const ImU32 kColZ      = IM_COL32( 38,  90, 220, 255);
static const ImU32 kColHover  = IM_COL32(255, 128,  16, 255);
static const ImU32 kColCenter = IM_COL32(255, 255, 255, 220);
static const ImU32 kAxisColors[3] = { kColX, kColY, kColZ };

static constexpr float kTranslationLineThick = 3.0f;
static constexpr float kRotationLineThick    = 2.0f;
static constexpr float kScaleLineThick       = 3.0f;
static constexpr float kScaleCircleRadius    = 5.0f;
static constexpr float kCenterCircleRadius   = 5.0f;

// ============================================================================
// SetupCompatibleXform
//
// Fast path  : prim already has XformCommonAPI-compatible ops -> read T/R/S.
// Fallback   : decompose the composed local matrix and insert a compatible
//              op-stack (with !resetXformStack!) in the current edit layer.
//              This handles prims from references that use separate rotateX/Y/Z
//              or matrix ops without touching the reference layer.
//
// Mirrors usdtweak's per-prim check: if (xformAPI) { ... } else { ... }
// but rather than falling back to a single matrix op, we create compatible
// standard TRS ops (better for layer-based override workflows).
// ============================================================================
static bool SetupCompatibleXform(
    const pxr::UsdStageRefPtr&                 stage,
    const pxr::SdfPath&                        primPath,
    pxr::GfVec3d&                              outT,
    pxr::GfVec3f&                              outR,
    pxr::GfVec3f&                              outS,
    pxr::UsdGeomXformCommonAPI::RotationOrder& outRotOrder)
{
    pxr::UsdPrim prim = stage->GetPrimAtPath(primPath);
    if (!prim) return false;

    pxr::UsdGeomXformCommonAPI api(prim);
    pxr::GfVec3f pivot3f;

    if (api.GetXformVectors(&outT, &outR, &outS, &pivot3f,
                            &outRotOrder, pxr::UsdTimeCode::Default()))
        return true;

    // Fallback -----------------------------------------------------------------
    LOG_WARNING("Prim " + primPath.GetString() +
                " has an incompatible xform stack. Creating a compatible "
                "override in the edit layer.");

    pxr::UsdGeomXformCache xc(pxr::UsdTimeCode::Default());
    bool reset = false;
    pxr::GfMatrix4d localMat = xc.GetLocalTransformation(prim, &reset);

    pxr::GfMatrix4d rotMat, shearMat, projMat;
    pxr::GfVec3d    scale3d, translate3d;
    if (!localMat.Factor(&rotMat, &scale3d, &shearMat, &translate3d, &projMat)) {
        rotMat = pxr::GfMatrix4d(1); scale3d = {1,1,1}; translate3d = {0,0,0};
    }

    double tw = 0, fb = 0, lr = 0, sw = 0;
    pxr::GfRotation::DecomposeRotation(
        rotMat, {1,0,0}, {0,1,0}, {0,0,1}, 1.0, &tw, &fb, &lr, &sw, false);

    outT = translate3d;
    outR = { float(pxr::GfRadiansToDegrees(tw)),
             float(pxr::GfRadiansToDegrees(fb)),
             float(pxr::GfRadiansToDegrees(lr)) };
    outS = { float(scale3d[0]), float(scale3d[1]), float(scale3d[2]) };
    outRotOrder = pxr::UsdGeomXformCommonAPI::RotationOrderXYZ;

    {
        pxr::UsdEditContext ec(stage, stage->GetEditTarget());
        pxr::UsdGeomXformable xf(prim);
        xf.SetResetXformStack(true);
        pxr::UsdGeomXformCommonAPI fresh(prim);
        fresh.SetTranslate(outT, pxr::UsdTimeCode::Default());
        fresh.SetRotate   (outR, outRotOrder, pxr::UsdTimeCode::Default());
        fresh.SetScale    (outS, pxr::UsdTimeCode::Default());
    }
    return true;
}

// ============================================================================
// Stage / selection
// ============================================================================
void TransformManipulator::SetStage(pxr::UsdStageRefPtr stage)
{
    m_stage = stage; m_primPath = pxr::SdfPath(); m_isDragging = false;
}

void TransformManipulator::SetSelectedPrim(const pxr::SdfPath& path)
{
    m_primPath = path; m_isDragging = false;
}

// ============================================================================
// GetGizmoAxes
// World  mode: fixed {1,0,0}, {0,1,0}, {0,0,1}
// Object mode: rows 0..2 of local-to-world matrix (normalised world-space
//              local axis directions — mirrors usdtweak's
//              ComputeManipulatorToWorldTransform().GetRow3(i))
// ============================================================================
void TransformManipulator::GetGizmoAxes(pxr::GfVec3d axes[3]) const
{
    axes[0] = {1,0,0}; axes[1] = {0,1,0}; axes[2] = {0,0,1};
    if (m_transformSpace == TransformSpace::World) return;
    if (!m_stage || m_primPath.IsEmpty()) return;
    pxr::UsdPrim prim = m_stage->GetPrimAtPath(m_primPath);
    if (!prim) return;
    pxr::UsdGeomXformCache xc(pxr::UsdTimeCode::Default());
    pxr::GfMatrix4d l2w = xc.GetLocalToWorldTransform(prim);
    for (int i = 0; i < 3; ++i) {
        pxr::GfVec3d row(l2w[i][0], l2w[i][1], l2w[i][2]);
        double n = row.GetLength();
        if (n > 1e-9) axes[i] = row / n;
    }
}

// ============================================================================
// WorldToScreen   p_clip = (p_world, 1) * VP   (USD row-vector convention)
// ============================================================================
bool TransformManipulator::WorldToScreen(const pxr::GfVec3d& world,
                                          const pxr::GfMatrix4d& vp,
                                          int vW, int vH,
                                          const ImVec2& imgPos, ImVec2& out)
{
    double cx = vp[0][0]*world[0]+vp[1][0]*world[1]+vp[2][0]*world[2]+vp[3][0];
    double cy = vp[0][1]*world[0]+vp[1][1]*world[1]+vp[2][1]*world[2]+vp[3][1];
    double cw = vp[0][3]*world[0]+vp[1][3]*world[1]+vp[2][3]*world[2]+vp[3][3];
    if (cw <= 0.0) return false;
    double iw = 1.0/cw;
    out = ImVec2(imgPos.x + float((cx*iw+1.0)*0.5*vW),
                 imgPos.y + float((1.0-cy*iw)*0.5*vH));
    return true;
}

// ============================================================================
// ComputeMouseRay   p_world = p_clip * VP^{-1}
// ============================================================================
pxr::GfRay TransformManipulator::ComputeMouseRay(const pxr::GfMatrix4d& vp,
                                                   const ImVec2& mouse,
                                                   const ImVec2& imgPos,
                                                   int vW, int vH)
{
    const pxr::GfRay kFB(pxr::GfVec3d(0), pxr::GfVec3d(0,0,-1));
    float nx =  2.f*(mouse.x-imgPos.x)/float(vW)-1.f;
    float ny =  1.f-2.f*(mouse.y-imgPos.y)/float(vH);
    double det = 0;
    pxr::GfMatrix4d vi = vp.GetInverse(&det);
    if (std::abs(det) < 1e-12) return kFB;
    auto up = [&](double z) -> pxr::GfVec3d {
        double wx=nx*vi[0][0]+ny*vi[1][0]+z*vi[2][0]+vi[3][0];
        double wy=nx*vi[0][1]+ny*vi[1][1]+z*vi[2][1]+vi[3][1];
        double wz=nx*vi[0][2]+ny*vi[1][2]+z*vi[2][2]+vi[3][2];
        double ww=nx*vi[0][3]+ny*vi[1][3]+z*vi[2][3]+vi[3][3];
        if (std::abs(ww)<1e-12) return pxr::GfVec3d(0);
        return {wx/ww, wy/ww, wz/ww};
    };
    pxr::GfVec3d dir = up(1.0)-up(-1.0);
    double n = dir.GetLength();
    if (n < 1e-9) return kFB;
    return pxr::GfRay(up(-1.0), dir/n);
}

// ============================================================================
// ComputeScreenFactor  (ImGuizmo constant-apparent-size formula)
// ============================================================================
float TransformManipulator::ComputeScreenFactor(const pxr::GfMatrix4d& vp,
                                                 const pxr::GfVec3d& pivot,
                                                 int vW, int vH, float frac)
{
    double pw=vp[0][3]*pivot[0]+vp[1][3]*pivot[1]+vp[2][3]*pivot[2]+vp[3][3];
    if (pw<=0) return 1.f;
    double iw=1/pw;
    double px=(vp[0][0]*pivot[0]+vp[1][0]*pivot[1]+vp[2][0]*pivot[2]+vp[3][0])*iw;
    double py=(vp[0][1]*pivot[0]+vp[1][1]*pivot[1]+vp[2][1]*pivot[2]+vp[3][1])*iw;
    float dr=(float)vW/(float)std::max(vH,1), mx=0;
    const pxr::GfVec3d kA[3]={{1,0,0},{0,1,0},{0,0,1}};
    for (auto& ax : kA) {
        pxr::GfVec3d t=pivot+ax;
        double tw=vp[0][3]*t[0]+vp[1][3]*t[1]+vp[2][3]*t[2]+vp[3][3];
        if (tw<=0) continue;
        double ti=1/tw;
        float dx=float((vp[0][0]*t[0]+vp[1][0]*t[1]+vp[2][0]*t[2]+vp[3][0])*ti-px);
        float dy=float((vp[0][1]*t[0]+vp[1][1]*t[1]+vp[2][1]*t[2]+vp[3][1])*ti-py);
        if (dr<1) dx*=dr; else dy/=dr;
        mx=std::max(mx,std::sqrt(dx*dx+dy*dy));
    }
    return (mx<1e-6f)?1.f:frac/mx;
}

// ============================================================================
// PointToSegmentDist
// ============================================================================
float TransformManipulator::PointToSegmentDist(ImVec2 p, ImVec2 a, ImVec2 b)
{
    float dx=b.x-a.x,dy=b.y-a.y,lsq=dx*dx+dy*dy;
    if (lsq<1e-6f){float ex=p.x-a.x,ey=p.y-a.y;return std::sqrt(ex*ex+ey*ey);}
    float t=std::max(0.f,std::min(1.f,((p.x-a.x)*dx+(p.y-a.y)*dy)/lsq));
    float cx=a.x+t*dx-p.x,cy=a.y+t*dy-p.y;
    return std::sqrt(cx*cx+cy*cy);
}

// ============================================================================
// HitTestAxes  (Move / Scale)
// ============================================================================
int TransformManipulator::HitTestAxes(const pxr::GfMatrix4d& vp,
                                       const pxr::GfVec3d& pivot, float sf,
                                       const ImVec2& imgPos, int vW, int vH,
                                       const ImVec2& mouse,
                                       const pxr::GfVec3d axes[3]) const
{
    ImVec2 pivSS;
    if (!WorldToScreen(pivot, vp, vW, vH, imgPos, pivSS)) return -1;
    float best=10.f; int hit=-1;
    for (int i=0;i<3;++i) {
        ImVec2 tipSS;
        if (!WorldToScreen(pivot+axes[i]*sf,vp,vW,vH,imgPos,tipSS)) continue;
        float d=PointToSegmentDist(mouse,pivSS,tipSS);
        if (d<best){best=d;hit=i;}
    }
    return hit;
}

// ============================================================================
// HitTestRotateRings  (front-facing half-arc proximity)
// ============================================================================
int TransformManipulator::HitTestRotateRings(const pxr::GfMatrix4d& vp,
                                              const pxr::GfVec3d& pivot, float sf,
                                              const pxr::GfVec3d& camEye,
                                              const ImVec2& imgPos, int vW, int vH,
                                              const ImVec2& mouse,
                                              const pxr::GfVec3d axes[3]) const
{
    static constexpr int   kSeg=32;
    static constexpr float kDF=1.2f;
    float r=sf*kDF;
    pxr::GfVec3d c2s=pivot-camEye;
    double l=c2s.GetLength();
    if (l<1e-9) c2s={0,0,-1}; else c2s/=l;
    float best=10.f; int hit=-1;
    for (int axis=0;axis<3;++axis) {
        const pxr::GfVec3d& u=(axis==0)?axes[1]:axes[0];
        const pxr::GfVec3d& v=(axis<2) ?axes[2]:axes[1];
        float ap=float(c2s[0]*u[0]+c2s[1]*u[1]+c2s[2]*u[2]);
        float bp=float(c2s[0]*v[0]+c2s[1]*v[1]+c2s[2]*v[2]);
        float as=std::atan2(bp,ap)+float(M_PI)*0.5f;
        ImVec2 prev; bool hp=false;
        for (int s=0;s<=kSeg;++s) {
            float a=as+float(M_PI)*(float(s)/float(kSeg));
            pxr::GfVec3d p=pivot+u*(r*std::cos(a))+v*(r*std::sin(a));
            ImVec2 ss;
            if (!WorldToScreen(p,vp,vW,vH,imgPos,ss)){hp=false;continue;}
            if (hp){float d=PointToSegmentDist(mouse,prev,ss);if(d<best){best=d;hit=axis;}}
            prev=ss; hp=true;
        }
    }
    return hit;
}

// ============================================================================
// Draw helpers
// ============================================================================
void TransformManipulator::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])
{
    ImVec2 pivSS;
    if (!WorldToScreen(pivot,vp,vW,vH,imgPos,pivSS)) return;
    for (int i=0;i<3;++i) {
        ImU32 col=(i==m_dragAxis||i==m_hoveredAxis)?kColHover:kAxisColors[i];
        ImVec2 shSS,tipSS;
        if (!WorldToScreen(pivot+axes[i]*sf*0.78f,vp,vW,vH,imgPos,shSS)) continue;
        if (!WorldToScreen(pivot+axes[i]*sf,       vp,vW,vH,imgPos,tipSS)) continue;
        dl->AddLine(pivSS,shSS,col,kTranslationLineThick);
        float adx=tipSS.x-shSS.x,ady=tipSS.y-shSS.y,al=std::sqrt(adx*adx+ady*ady);
        if (al<1.f) continue;
        float px=-ady/al,py=adx/al;
        float tl=std::sqrt((tipSS.x-pivSS.x)*(tipSS.x-pivSS.x)+(tipSS.y-pivSS.y)*(tipSS.y-pivSS.y));
        float hw=tl*0.12f;
        dl->AddTriangleFilled(tipSS,{shSS.x+px*hw,shSS.y+py*hw},{shSS.x-px*hw,shSS.y-py*hw},col);
    }
    dl->AddCircleFilled(pivSS,kCenterCircleRadius,kColCenter,16);
}

void TransformManipulator::DrawRotateGizmo(ImDrawList* dl, const pxr::GfMatrix4d& vp,
                                            const pxr::GfVec3d& pivot, float sf,
                                            const pxr::GfVec3d& camEye,
                                            const ImVec2& imgPos, int vW, int vH,
                                            const pxr::GfVec3d axes[3])
{
    static constexpr int kSeg=64; static constexpr float kDF=1.2f;
    float r=sf*kDF;
    pxr::GfVec3d c2s=pivot-camEye;
    double cl=c2s.GetLength();
    if (cl<1e-9) c2s={0,0,-1}; else c2s/=cl;
    for (int axis=0;axis<3;++axis) {
        ImU32 col=(axis==m_dragAxis||axis==m_hoveredAxis)?kColHover:kAxisColors[axis];
        float lw=(axis==m_dragAxis||axis==m_hoveredAxis)?kRotationLineThick+1.5f:kRotationLineThick;
        const pxr::GfVec3d& u=(axis==0)?axes[1]:axes[0];
        const pxr::GfVec3d& v=(axis<2) ?axes[2]:axes[1];
        float ap=float(c2s[0]*u[0]+c2s[1]*u[1]+c2s[2]*u[2]);
        float bp=float(c2s[0]*v[0]+c2s[1]*v[1]+c2s[2]*v[2]);
        float as=std::atan2(bp,ap)+float(M_PI)*0.5f;
        std::vector<ImVec2> pts; pts.reserve(kSeg+1);
        for (int s=0;s<=kSeg;++s) {
            float a=as+float(M_PI)*(float(s)/float(kSeg));
            pxr::GfVec3d p=pivot+u*(r*std::cos(a))+v*(r*std::sin(a));
            ImVec2 ss;
            if (WorldToScreen(p,vp,vW,vH,imgPos,ss)) pts.push_back(ss);
        }
        if (pts.size()>1)
            dl->AddPolyline(pts.data(),int(pts.size()),col,ImDrawFlags_None,lw);
    }
}

void TransformManipulator::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])
{
    ImVec2 pivSS;
    if (!WorldToScreen(pivot,vp,vW,vH,imgPos,pivSS)) return;
    for (int i=0;i<3;++i) {
        ImU32 col=(i==m_dragAxis||i==m_hoveredAxis)?kColHover:kAxisColors[i];
        ImVec2 tipSS;
        if (!WorldToScreen(pivot+axes[i]*sf,vp,vW,vH,imgPos,tipSS)) continue;
        dl->AddLine(pivSS,tipSS,col,kScaleLineThick);
        dl->AddCircleFilled(tipSS,kScaleCircleRadius,col,16);
    }
    dl->AddCircleFilled(pivSS,kCenterCircleRadius+1.f,kColCenter,16);
}

// ============================================================================
// Render
// ============================================================================
void TransformManipulator::Render(ImDrawList* dl, const pxr::GfMatrix4d& vp,
                                   const pxr::GfVec3d& pivot,
                                   const pxr::GfVec3d& camEye,
                                   const ImVec2& imgPos, int vW, int vH)
{
    if (m_mode==ManipulatorMode::Select) return;
    if (!m_stage||m_primPath.IsEmpty()||!dl||vW<=0||vH<=0) return;
    float sf=ComputeScreenFactor(vp,pivot,vW,vH,0.15f);
    pxr::GfVec3d axes[3]; GetGizmoAxes(axes);
    switch (m_mode) {
        case ManipulatorMode::Move:   DrawMoveGizmo  (dl,vp,pivot,sf,imgPos,vW,vH,axes); break;
        case ManipulatorMode::Rotate: DrawRotateGizmo(dl,vp,pivot,sf,camEye,imgPos,vW,vH,axes); break;
        case ManipulatorMode::Scale:  DrawScaleGizmo (dl,vp,pivot,sf,imgPos,vW,vH,axes); break;
        default: break;
    }
}

// ============================================================================
// HandleInput
//
// Full port of usdtweak's three manipulators with World/Object space switch.
//
// TRANSLATE  (PositionManipulator::OnUpdate)
//   GfFindClosestPoints(mouseRay, axisLine) -> signed scalar delta along axis.
//   Applied ABSOLUTE each frame: newTranslate = origTranslate + parentSpaceDelta.
//   Parent-space conversion: w2p.TransformDir(worldDeltaVec).
//
// ROTATE  (RotationManipulator::OnUpdate + space switch)
//   Clock-hand: intersect mouse ray with ring plane (GfPlane) -> rotateTo vector.
//   worldRotation = GfRotation(dragStartClockHand, currentClockHand).
//   axisSign      = dot(planeNormal, worldRotation.GetAxis()) > 0 ? 1 : -1.
//   localDeltaAxis:
//     Object => initRot.GetRow3(dragAxis)   [usdtweak exact: local axis in parent space]
//     World  => world unit vector
//   deltaRotation = GfRotation(localDeltaAxis * axisSign, angle).
//   resultingRotation:
//     Object => GfMatrix4d(1).SetRotate(delta) * initRot   [usdtweak exact]
//     World  => initRot * GfMatrix4d(1).SetRotate(delta)   [USD row-vector world delta]
//   DecomposeRotation with FIXED initRot rows as reference axes + current hints.
//
// SCALE  (ScaleManipulator::OnUpdate)
//   Same axis-line projection as Translate.
//   ratio = curT / originT (axis-line parameter, robust for off-origin pivots).
//   newScale[axis] = origScale[axis] * ratio  (absolute from drag start).
// ============================================================================
bool TransformManipulator::HandleInput(const pxr::GfMatrix4d& vp,
                                        const pxr::GfVec3d& pivot,
                                        const pxr::GfVec3d& camEye,
                                        const ImVec2& imgPos,
                                        int vW, int vH, bool hovered)
{
    if (m_mode==ManipulatorMode::Select) return false;
    if (!m_stage||m_primPath.IsEmpty())  return false;

    ImGuiIO& io=ImGui::GetIO();
    ImVec2 mouse=io.MousePos;
    float sf=ComputeScreenFactor(vp,pivot,vW,vH,0.15f);
    pxr::GfVec3d axes[3]; GetGizmoAxes(axes);

    // Hover update
    if (!m_isDragging && hovered) {
        m_hoveredAxis=(m_mode==ManipulatorMode::Rotate)
            ?HitTestRotateRings(vp,pivot,sf,camEye,imgPos,vW,vH,mouse,axes)
            :HitTestAxes       (vp,pivot,sf,        imgPos,vW,vH,mouse,axes);
    }

    bool consumed=false;

    // =========================================================================
    // BEGIN drag
    // =========================================================================
    if (hovered && !m_isDragging &&
        ImGui::IsMouseClicked(ImGuiMouseButton_Left) && !io.KeyAlt)
    {
        int hit=(m_mode==ManipulatorMode::Rotate)
            ?HitTestRotateRings(vp,pivot,sf,camEye,imgPos,vW,vH,mouse,axes)
            :HitTestAxes       (vp,pivot,sf,        imgPos,vW,vH,mouse,axes);

        if (hit>=0) {
            m_isDragging=true; m_dragAxis=hit; consumed=true;

            pxr::GfVec3d trans; pxr::GfVec3f rot,scale;
            pxr::UsdGeomXformCommonAPI::RotationOrder rotOrder;
            if (!SetupCompatibleXform(m_stage,m_primPath,trans,rot,scale,rotOrder)) {
                m_isDragging=false; return false;
            }
            m_dragOriginalTranslate=trans; m_dragOriginalRotate=rot;
            m_dragOriginalScale=scale;     m_dragOriginalRotOrder=rotOrder;
            m_dragStartTranslate=trans;    m_dragStartRotate=rot;
            m_dragStartScale=scale;

            pxr::GfRay ray=ComputeMouseRay(vp,mouse,imgPos,vW,vH);

            if (m_mode==ManipulatorMode::Move) {
                // usdtweak: _axisLine = GfLine(pivot, objectTransform.GetRow3(axis))
                m_dragMoveAxisLine=pxr::GfLine(pivot,axes[m_dragAxis]);
                pxr::GfVec3d rpt;
                pxr::GfFindClosestPoints(ray,m_dragMoveAxisLine,&rpt,&m_dragMoveOriginOnAxis);
            }
            else if (m_mode==ManipulatorMode::Rotate) {
                // Ring plane in world space.
                // usdtweak: _planeNormal3d = manipCoords.GetRow3(_selectedAxis)
                m_dragRotatePlaneNormal=axes[m_dragAxis];
                pxr::GfPlane plane(m_dragRotatePlaneNormal,pivot);
                double dist=0;
                if (ray.Intersect(plane,&dist) && dist>0)
                    m_dragRotateFrom=pivot-ray.GetPoint(dist);
                else
                    m_dragRotateFrom=axes[(m_dragAxis+1)%3];

                // Capture initial rotation matrix (local->parent, fixed for whole drag).
                // usdtweak: _rotateMatrixOnBegin = GetOpTransform(opType, VtValue(rotation))
                pxr::UsdGeomXformOp::Type opType=
                    pxr::UsdGeomXformCommonAPI::ConvertRotationOrderToOpType(rotOrder);
                m_dragRotateInitialRotMat=
                    pxr::UsdGeomXformOp::GetOpTransform(opType,pxr::VtValue(rot));

                // Capture delta axis and space mode at drag start.
                // Object: initRot.GetRow3(axis) -- local axis in parent space.
                // World : world unit vector.
                if (m_transformSpace==TransformSpace::Object) {
                    m_dragRotateDeltaAxis=m_dragRotateInitialRotMat.GetRow3(m_dragAxis);
                    double n=m_dragRotateDeltaAxis.GetLength();
                    if (n>1e-9) m_dragRotateDeltaAxis/=n;
                } else {
                    static const pxr::GfVec3d kW[3]={{1,0,0},{0,1,0},{0,0,1}};
                    m_dragRotateDeltaAxis=kW[m_dragAxis];
                }
                m_dragRotateObjectSpace=(m_transformSpace==TransformSpace::Object);
            }
            else if (m_mode==ManipulatorMode::Scale) {
                // usdtweak: _axisLine = GfLine(pivot, objectTransform.GetRow3(axis))
                m_dragScaleAxisLine=pxr::GfLine(pivot,axes[m_dragAxis]);
                pxr::GfVec3d rpt;
                pxr::GfFindClosestPoints(ray,m_dragScaleAxisLine,&rpt,&m_dragScaleOriginOnAxis);
            }
        }
    }

    // =========================================================================
    // DRAG update
    // =========================================================================
    if (m_isDragging) {
        consumed=true;
        if (ImGui::IsMouseDown(ImGuiMouseButton_Left)) {
            if (m_dragAxis<0||m_dragAxis>2) { m_isDragging=false;m_dragAxis=-1;return consumed; }

            pxr::GfRay ray=ComputeMouseRay(vp,mouse,imgPos,vW,vH);

            // -----------------------------------------------------------------
            // TRANSLATE
            // usdtweak: translation[_selectedAxis] += delta
            //           where delta = signed distance along axis line.
            // We extend to support both Object and World space via parent-space
            // conversion of the world-space delta vector.
            // -----------------------------------------------------------------
            if (m_mode==ManipulatorMode::Move) {
                pxr::GfVec3d rpt,cur;
                pxr::GfFindClosestPoints(ray,m_dragMoveAxisLine,&rpt,&cur);
                double oT=0,cT=0;
                m_dragMoveAxisLine.FindClosestPoint(m_dragMoveOriginOnAxis,&oT);
                m_dragMoveAxisLine.FindClosestPoint(cur,&cT);
                double delta=cT-oT;

                // World-space delta vector along the gizmo axis.
                pxr::GfVec3d worldDelta=axes[m_dragAxis]*delta;

                // Convert to parent space.
                // (For Object mode: axes[i] is the local axis in world space.
                //  w2p.TransformDir maps it to the local axis in parent space,
                //  which equals initRot.GetRow3(i). Adding this to the stored
                //  parent-space translation is geometrically correct.)
                pxr::GfVec3d parentDelta=worldDelta;
                pxr::UsdPrim prim=m_stage->GetPrimAtPath(m_primPath);
                if (prim) {
                    pxr::UsdGeomXformCache xc(pxr::UsdTimeCode::Default());
                    pxr::UsdPrim parent=prim.GetParent();
                    if (parent && !parent.IsPseudoRoot()) {
                        pxr::GfMatrix4d p2w=xc.GetLocalToWorldTransform(parent);
                        double det=0; pxr::GfMatrix4d w2p=p2w.GetInverse(&det);
                        if (std::abs(det)>1e-9) parentDelta=w2p.TransformDir(worldDelta);
                    }
                }

                ApplyTranslate(m_dragOriginalTranslate+parentDelta);
            }

            // -----------------------------------------------------------------
            // ROTATE
            // Direct port of RotationManipulator::OnUpdate with space switch.
            // -----------------------------------------------------------------
            else if (m_mode==ManipulatorMode::Rotate) {
                pxr::GfPlane plane(m_dragRotatePlaneNormal,pivot);
                double dist=0;
                if (!ray.Intersect(plane,&dist)||dist<=0) return consumed;
                pxr::GfVec3d rotateTo=pivot-ray.GetPoint(dist);
                if (m_dragRotateFrom.GetLength()<1e-9||rotateTo.GetLength()<1e-9)
                    return consumed;

                // usdtweak: const GfRotation worldRotation(_rotateFrom, rotateTo)
                pxr::GfRotation worldRotation(m_dragRotateFrom,rotateTo);

                // usdtweak: axisSign = _planeNormal3d * worldRotation.GetAxis() > 0
                double axisSign=
                    (pxr::GfDot(m_dragRotatePlaneNormal,worldRotation.GetAxis())>0.0)
                    ?1.0:-1.0;

                // usdtweak: const GfRotation deltaRotation(localPlaneNormal*axisSign, angle)
                pxr::GfRotation deltaRot(m_dragRotateDeltaAxis*axisSign,
                                         worldRotation.GetAngle());
                pxr::GfMatrix4d deltaM=pxr::GfMatrix4d(1.0).SetRotate(deltaRot);

                // usdtweak: resultingRotation = GfMatrix4d(1).SetRotate(delta)*initRot
                //   Object (usdtweak exact): deltaM * initRot
                //   World (USD row-vector) : initRot * deltaM
                pxr::GfMatrix4d result=
                    m_dragRotateObjectSpace
                    ? deltaM * m_dragRotateInitialRotMat
                    : m_dragRotateInitialRotMat * deltaM;

                // usdtweak: GfRotation::DecomposeRotation(resultingRotation,
                //             pxAxis, pyAxis, pzAxis, 1.0, &tw, &fb, &lr, &sw, true)
                const pxr::GfVec3d pxA=m_dragRotateInitialRotMat.GetRow3(0);
                const pxr::GfVec3d pyA=m_dragRotateInitialRotMat.GetRow3(1);
                const pxr::GfVec3d pzA=m_dragRotateInitialRotMat.GetRow3(2);

                double tw=pxr::GfDegreesToRadians(double(m_dragStartRotate[0]));
                double fb=pxr::GfDegreesToRadians(double(m_dragStartRotate[1]));
                double lr=pxr::GfDegreesToRadians(double(m_dragStartRotate[2]));
                double sw=0;

                pxr::GfRotation::DecomposeRotation(
                    result,pxA,pyA,pzA,1.0,&tw,&fb,&lr,&sw,/*useHint=*/true);

                pxr::GfVec3f newRot(float(pxr::GfRadiansToDegrees(tw)),
                                    float(pxr::GfRadiansToDegrees(fb)),
                                    float(pxr::GfRadiansToDegrees(lr)));
                ApplyRotate(newRot,m_dragOriginalRotOrder);
                m_dragStartRotate=newRot;
            }

            // -----------------------------------------------------------------
            // SCALE
            // usdtweak: scale[_selectedAxis] = _scalesOnBegin[i][axis] * ratio
            //           ratio = mouseOnAxis.GetLength() / originOnAxis.GetLength()
            // We use axis-line T parameters (distance from pivot) which is more
            // robust when the pivot is not at the world origin.
            // -----------------------------------------------------------------
            else if (m_mode==ManipulatorMode::Scale) {
                pxr::GfVec3d rpt,cur;
                pxr::GfFindClosestPoints(ray,m_dragScaleAxisLine,&rpt,&cur);
                double oT=0,cT=0;
                m_dragScaleAxisLine.FindClosestPoint(m_dragScaleOriginOnAxis,&oT);
                m_dragScaleAxisLine.FindClosestPoint(cur,&cT);
                double ratio=(std::abs(oT)>1e-6)?cT/oT:1.0;
                pxr::GfVec3f ns=m_dragOriginalScale;
                ns[m_dragAxis]=std::max(0.001f,float(double(m_dragOriginalScale[m_dragAxis])*ratio));
                ApplyScale(ns); m_dragStartScale=ns;
            }

        } else {
            // Mouse released.
            bool moved=
                (m_dragStartTranslate!=m_dragOriginalTranslate)||
                (m_dragStartRotate   !=m_dragOriginalRotate   )||
                (m_dragStartScale    !=m_dragOriginalScale    );
            if (moved && m_commandHistory && m_stage && !m_primPath.IsEmpty()) {
                auto cmd=std::make_unique<TransformCommand>(
                    m_stage, m_primPath,
                    m_stage->GetEditTarget().GetLayer(),
                    m_dragOriginalTranslate,m_dragOriginalRotate,m_dragOriginalScale,
                    m_dragStartTranslate,   m_dragStartRotate,   m_dragStartScale,
                    m_dragOriginalRotOrder,
                    "Transform "+m_primPath.GetName());
                m_commandHistory->Push(std::move(cmd));
            }
            m_isDragging=false; m_dragAxis=-1;
        }
    }
    return consumed;
}

// ============================================================================
// USD write helpers — write ABSOLUTE values
// ============================================================================
void TransformManipulator::ApplyTranslate(const pxr::GfVec3d& t)
{
    if (!m_stage||m_primPath.IsEmpty()) return;
    pxr::UsdPrim p=m_stage->GetPrimAtPath(m_primPath); if (!p) return;
    pxr::UsdEditContext ec(m_stage,m_stage->GetEditTarget());
    pxr::UsdGeomXformCommonAPI(p).SetTranslate(t,pxr::UsdTimeCode::Default());
    m_dragStartTranslate=t;
}

void TransformManipulator::ApplyRotate(const pxr::GfVec3f& r,
                                        pxr::UsdGeomXformCommonAPI::RotationOrder ro)
{
    if (!m_stage||m_primPath.IsEmpty()) return;
    pxr::UsdPrim p=m_stage->GetPrimAtPath(m_primPath); if (!p) return;
    pxr::UsdEditContext ec(m_stage,m_stage->GetEditTarget());
    pxr::UsdGeomXformCommonAPI(p).SetRotate(r,ro,pxr::UsdTimeCode::Default());
    m_dragStartRotate=r;
}

void TransformManipulator::ApplyScale(const pxr::GfVec3f& s)
{
    if (!m_stage||m_primPath.IsEmpty()) return;
    pxr::UsdPrim p=m_stage->GetPrimAtPath(m_primPath); if (!p) return;
    pxr::UsdEditContext ec(m_stage,m_stage->GetEditTarget());
    pxr::UsdGeomXformCommonAPI(p).SetScale(s,pxr::UsdTimeCode::Default());
    m_dragStartScale=s;
}

} // namespace UsdLayerManager