UsdLayerManager/src/ui/CurveEditorPanel.h

#pragma once

#include "../core/CommandHistory.h"
#include <pxr/usd/usd/stage.h>
#include <pxr/usd/usd/attribute.h>
#include <pxr/usd/usd/timeCode.h>
#include <pxr/usd/usd/notice.h>
#include <pxr/base/tf/weakBase.h>
#include <pxr/base/tf/notice.h>
#include <imgui.h>
#include <functional>
#include <string>
#include <vector>

namespace UsdLayerManager {

class CurveEditorPanel : public pxr::TfWeakBase {
public:
    void SetStage(pxr::UsdStageRefPtr stage);
    void SetCommandHistory(CommandHistory* history) { m_commandHistory = history; }
    void SetSelectedPrimPath(const std::string& path);
    void SetTimeCodes(pxr::UsdTimeCode displayTime, pxr::UsdTimeCode editTime);
    void Render();

    // Fires when the user drags the time cursor; drives the Timeline panel.
    std::function<void(double frame)> OnScrub;

private:
    // ── Internal data ────────────────────────────────────────────────────────

    struct BezierKey {
        double time  = 0.0;
        double value = 0.0;
        // Tangent handles as (dt, dv) offsets from the key.
        // in-handle is left of the key (dt < 0), out-handle is right (dt > 0).
        double inTangentDt  = -1.0, inTangentDv  = 0.0;
        double outTangentDt =  1.0, outTangentDv = 0.0;
        bool   tangentsBroken = false;
    };

    struct CurveChannel {
        std::string          displayName;
        pxr::UsdAttribute    attr;
        int                  component = -1; // -1=scalar, 0/1/2/3 = X/Y/Z/W
        ImVec4               color     = {1,1,1,1};
        bool                 visible   = true;
        bool                 dirty     = false;
        std::vector<double>  origTimes;
        std::vector<double>  origValues;
        std::vector<BezierKey> keys;
    };

    enum class HandlePart { Key, InTangent, OutTangent };
    struct SelectedHandle { int chanIdx; int keyIdx; HandlePart part; };

    // ── Methods ──────────────────────────────────────────────────────────────

    void RefreshFromStage();
    void FitCurveFromSamples(CurveChannel& ch);
    void BakeChannelToUsd(CurveChannel& ch);

    // Bezier math
    double EvaluateBezier(const std::vector<BezierKey>& keys, double T) const;
    void   DrawBezierSegment(ImDrawList* dl, const BezierKey& k0, const BezierKey& k1,
                             ImVec4 color, ImVec2 canvasMin, ImVec2 canvasMax) const;

    // Coordinate helpers
    float  TimeToX(double t,  ImVec2 cMin, ImVec2 cMax) const;
    float  ValToY (double v,  ImVec2 cMin, ImVec2 cMax) const;
    double XToTime(float  x,  ImVec2 cMin, ImVec2 cMax) const;
    double YToVal (float  y,  ImVec2 cMin, ImVec2 cMax) const;
    ImVec2 KeyToScreen(const BezierKey& k, ImVec2 cMin, ImVec2 cMax) const;

    // USD channel read/write helpers
    double ReadChannelValue (const pxr::UsdAttribute& attr, int comp, double time) const;
    void   WriteChannelValue(const pxr::UsdAttribute& attr, int comp,
                             double time, double val) const;
    void   ClearAllTimeSamples(const pxr::UsdAttribute& attr) const;

    // Rendering sub-sections
    void RenderToolbar();
    void RenderAttrList(float listWidth);
    void RenderCanvas(ImVec2 canvasMin, ImVec2 canvasMax);
    void RenderGrid(ImDrawList* dl, ImVec2 cMin, ImVec2 cMax);
    void RenderTimeCursor(ImDrawList* dl, ImVec2 cMin, ImVec2 cMax);
    void RenderCurve(ImDrawList* dl, CurveChannel& ch, ImVec2 cMin, ImVec2 cMax);
    void HandleCanvasInput(ImVec2 cMin, ImVec2 cMax);

    // Selection helpers
    bool IsSelected(int chanIdx, int keyIdx, HandlePart part) const;
    void Select(int chanIdx, int keyIdx, HandlePart part, bool additive);
    void ClearSelection();

    // Interaction helpers
    void FrameAll();
    void ApplyAllDirty();
    void RevertAll();
    void DeleteSelectedKeys();
    void AddKeyAt(int chanIdx, double time, double value);
    void RefitNeighborTangents(CurveChannel& ch, int keyIdx);

    // Bezier metadata persistence (stored in prim customData["curveEditor"])
    static std::string ChannelMetadataKey(const CurveChannel& ch);
    bool SaveBezierToMetadata(const pxr::UsdPrim& prim, const CurveChannel& ch) const;
    bool SaveBezierToMetadata(const pxr::UsdPrim& prim, const std::string& channelKey,
                              const std::vector<BezierKey>& keys) const;
    bool LoadBezierFromMetadata(const pxr::UsdPrim& prim, CurveChannel& ch) const;

    // ── State ────────────────────────────────────────────────────────────────

    pxr::UsdStageRefPtr  m_stage;
    CommandHistory*      m_commandHistory = nullptr;
    std::string          m_selectedPrimPath;
    pxr::UsdTimeCode     m_displayTime = pxr::UsdTimeCode::Default();
    pxr::UsdTimeCode     m_editTime    = pxr::UsdTimeCode::Default();

    std::vector<CurveChannel>  m_channels;
    std::vector<SelectedHandle> m_selection;
    int                        m_activeChannel = -1; // for double-click-to-add

    bool m_needsRefresh    = true;
    bool m_suppressNotice  = false; // set during bake writes to silence OnObjectsChanged

    // USD change notification
    pxr::TfNotice::Key m_stageChangeKey;
    void OnObjectsChanged(const pxr::UsdNotice::ObjectsChanged& notice,
                          const pxr::UsdStageWeakPtr& sender);

    // View range
    double m_viewMinTime = 0.0,  m_viewMaxTime = 100.0;
    double m_viewMinVal  = -1.0, m_viewMaxVal  =   1.0;

    // Pan/zoom drag state
    bool   m_isPanning    = false;
    ImVec2 m_panLastMouse = {0,0};
    double m_panMinTimeStart = 0.0, m_panMaxTimeStart = 0.0;
    double m_panMinValStart  = 0.0, m_panMaxValStart  = 0.0;

    // Drag-move state
    bool   m_isDragging    = false;
    ImVec2 m_dragStartMouse= {0,0};
    double m_dragStartTime = 0.0, m_dragStartValue = 0.0;
    struct KeySnapshot { int chanIdx; int keyIdx; double time; double value;
                         double inDt, inDv, outDt, outDv; };
    std::vector<KeySnapshot> m_dragSnapshot;

    // Selection box drag
    bool   m_isBoxSelecting = false;
    ImVec2 m_boxSelectStart = {0,0};

    // Scrub cursor drag
    bool m_isDraggingCursor = false;

    // Options
    bool   m_simplifyOnBake  = false;
    float  m_simplifyTol     = 0.01f;
    bool   m_snapToFrame     = true;
};

} // namespace UsdLayerManager