#include "CurveEditorPanel.h"
#include "../core/commands/AttributeSetCommand.h"
#include "../utils/Logger.h"

#include <pxr/usd/sdf/path.h>
#include <pxr/base/vt/value.h>
#include <pxr/base/gf/vec2f.h>
#include <pxr/base/gf/vec2d.h>
#include <pxr/base/gf/vec3f.h>
#include <pxr/base/gf/vec3d.h>
#include <pxr/base/gf/vec4f.h>
#include <pxr/base/gf/vec4d.h>
#include <pxr/usd/sdf/valueTypeName.h>
#include <pxr/usd/usd/notice.h>
#include <pxr/usd/usd/prim.h>
#include <pxr/base/tf/notice.h>
#include <pxr/base/tf/weakPtr.h>
#include <pxr/base/vt/dictionary.h>
#include <pxr/base/vt/array.h>

#include <imgui.h>
#include <imgui_internal.h>

#include <algorithm>
#include <cmath>
#include <cstring>
#include <memory>

PXR_NAMESPACE_USING_DIRECTIVE

namespace UsdLayerManager {

// ── Colour palette for channels ──────────────────────────────────────────────

static const ImVec4 kComponentColors[4] = {
    {0.96f, 0.36f, 0.36f, 1.f},  // X – red
    {0.36f, 0.86f, 0.36f, 1.f},  // Y – green
    {0.36f, 0.56f, 0.96f, 1.f},  // Z – blue
    {0.96f, 0.86f, 0.36f, 1.f},  // W – yellow
};
static const ImVec4 kScalarPalette[] = {
    {0.9f, 0.7f, 0.2f, 1.f},
    {0.5f, 0.9f, 0.8f, 1.f},
    {0.9f, 0.5f, 0.8f, 1.f},
    {0.7f, 0.9f, 0.4f, 1.f},
    {0.4f, 0.7f, 0.9f, 1.f},
};
static int s_scalarColorIdx = 0;

// ── USD helpers ──────────────────────────────────────────────────────────────

double CurveEditorPanel::ReadChannelValue(const UsdAttribute& attr, int comp,
                                          double time) const
{
    VtValue v;
    if (!attr.Get(&v, UsdTimeCode(time))) return 0.0;

    auto tn = attr.GetTypeName();
    if (tn == SdfValueTypeNames->Float   || tn == SdfValueTypeNames->Double  ||
        tn == SdfValueTypeNames->Half    || tn == SdfValueTypeNames->Int     ||
        tn == SdfValueTypeNames->Bool    || tn == SdfValueTypeNames->Int64   ||
        tn == SdfValueTypeNames->UInt   || tn == SdfValueTypeNames->UInt64) {
        if (v.IsHolding<float>())  return double(v.UncheckedGet<float>());
        if (v.IsHolding<double>()) return v.UncheckedGet<double>();
        if (v.IsHolding<int>())    return double(v.UncheckedGet<int>());
        if (v.IsHolding<bool>())   return v.UncheckedGet<bool>() ? 1.0 : 0.0;
        if (v.IsHolding<GfHalf>()) return double(float(v.UncheckedGet<GfHalf>()));
        return 0.0;
    }
    if (comp < 0) comp = 0;
    if (v.IsHolding<GfVec2f>()) { auto& u=v.UncheckedGet<GfVec2f>(); return comp<2?double(u[comp]):0.0; }
    if (v.IsHolding<GfVec2d>()) { auto& u=v.UncheckedGet<GfVec2d>(); return comp<2?u[comp]:0.0; }
    if (v.IsHolding<GfVec3f>()) { auto& u=v.UncheckedGet<GfVec3f>(); return comp<3?double(u[comp]):0.0; }
    if (v.IsHolding<GfVec3d>()) { auto& u=v.UncheckedGet<GfVec3d>(); return comp<3?u[comp]:0.0; }
    if (v.IsHolding<GfVec4f>()) { auto& u=v.UncheckedGet<GfVec4f>(); return comp<4?double(u[comp]):0.0; }
    if (v.IsHolding<GfVec4d>()) { auto& u=v.UncheckedGet<GfVec4d>(); return comp<4?u[comp]:0.0; }
    return 0.0;
}

void CurveEditorPanel::WriteChannelValue(const UsdAttribute& attr, int comp,
                                         double time, double val) const
{
    UsdTimeCode tc(time);
    auto tn = attr.GetTypeName();

    if (comp < 0) {
        // Scalar
        if (tn == SdfValueTypeNames->Float)  { attr.Set(float(val),  tc); return; }
        if (tn == SdfValueTypeNames->Double) { attr.Set(val,          tc); return; }
        if (tn == SdfValueTypeNames->Int)    { attr.Set(int(std::round(val)), tc); return; }
        if (tn == SdfValueTypeNames->Bool)   { attr.Set(val != 0.0,  tc); return; }
        if (tn == SdfValueTypeNames->Half)   { attr.Set(GfHalf(float(val)), tc); return; }
        return;
    }

    // Vector — read full value, replace component, write back
    VtValue existing;
    attr.Get(&existing, tc);

    if (tn == SdfValueTypeNames->Float3  || tn == SdfValueTypeNames->Vector3f ||
        tn == SdfValueTypeNames->Point3f || tn == SdfValueTypeNames->Normal3f ||
        tn == SdfValueTypeNames->Color3f) {
        GfVec3f u = existing.IsHolding<GfVec3f>() ? existing.UncheckedGet<GfVec3f>() : GfVec3f(0);
        if (comp < 3) u[comp] = float(val);
        attr.Set(u, tc);
        return;
    }
    if (tn == SdfValueTypeNames->Double3 || tn == SdfValueTypeNames->Vector3d ||
        tn == SdfValueTypeNames->Point3d || tn == SdfValueTypeNames->Normal3d) {
        GfVec3d u = existing.IsHolding<GfVec3d>() ? existing.UncheckedGet<GfVec3d>() : GfVec3d(0);
        if (comp < 3) u[comp] = val;
        attr.Set(u, tc);
        return;
    }
    if (tn == SdfValueTypeNames->Float2) {
        GfVec2f u = existing.IsHolding<GfVec2f>() ? existing.UncheckedGet<GfVec2f>() : GfVec2f(0);
        if (comp < 2) u[comp] = float(val);
        attr.Set(u, tc);
        return;
    }
    if (tn == SdfValueTypeNames->Double2) {
        GfVec2d u = existing.IsHolding<GfVec2d>() ? existing.UncheckedGet<GfVec2d>() : GfVec2d(0);
        if (comp < 2) u[comp] = val;
        attr.Set(u, tc);
        return;
    }
    if (tn == SdfValueTypeNames->Float4 || tn == SdfValueTypeNames->Color4f) {
        GfVec4f u = existing.IsHolding<GfVec4f>() ? existing.UncheckedGet<GfVec4f>() : GfVec4f(0);
        if (comp < 4) u[comp] = float(val);
        attr.Set(u, tc);
        return;
    }
    if (tn == SdfValueTypeNames->Double4) {
        GfVec4d u = existing.IsHolding<GfVec4d>() ? existing.UncheckedGet<GfVec4d>() : GfVec4d(0);
        if (comp < 4) u[comp] = val;
        attr.Set(u, tc);
        return;
    }
}

void CurveEditorPanel::ClearAllTimeSamples(const UsdAttribute& attr) const
{
    std::vector<double> times;
    attr.GetTimeSamples(&times);
    for (double t : times)
        attr.ClearAtTime(UsdTimeCode(t));
}

// ── Stage / prim refresh ──────────────────────────────────────────────────────

void CurveEditorPanel::SetStage(UsdStageRefPtr stage)
{
    TfNotice::Revoke(m_stageChangeKey);
    m_stage = stage;
    m_channels.clear();
    m_selection.clear();
    m_activeChannel = -1;
    m_needsRefresh = true;
    if (stage) {
        m_stageChangeKey = TfNotice::Register(
            TfCreateWeakPtr(this),
            &CurveEditorPanel::OnObjectsChanged,
            stage);
    }
}

void CurveEditorPanel::OnObjectsChanged(
    const UsdNotice::ObjectsChanged& notice,
    const UsdStageWeakPtr& /*sender*/)
{
    if (m_suppressNotice) return;
    if (m_selectedPrimPath.empty()) return;
    SdfPath primPath(m_selectedPrimPath);

    // Structural change (prim added/removed/retyped) → full refresh
    for (const SdfPath& p : notice.GetResyncedPaths()) {
        if (primPath.HasPrefix(p) || p.HasPrefix(primPath)) {
            m_needsRefresh = true;
            return;
        }
    }

    // Value-only change → re-fit clean channels; ignore prim-level metadata changes
    for (const SdfPath& changed : notice.GetChangedInfoOnlyPaths()) {
        // Prim-level path means a metadata change (e.g. customData) — not an attr write
        if (changed == primPath) continue;

        if (changed.GetPrimPath() != primPath) continue;

        bool knownChannel = false;
        for (auto& ch : m_channels) {
            if (changed == ch.attr.GetPath()) {
                knownChannel = true;
                if (!ch.dirty) FitCurveFromSamples(ch); // preserve unsaved edits
                break;
            }
        }

        // Changed attribute not in our list → may be a newly animated attr → full refresh
        if (!knownChannel)
            m_needsRefresh = true;
    }
}

void CurveEditorPanel::SetSelectedPrimPath(const std::string& path)
{
    if (m_selectedPrimPath == path) return;
    m_selectedPrimPath = path;
    m_channels.clear();
    m_selection.clear();
    m_activeChannel = -1;
    m_needsRefresh = true;
}

void CurveEditorPanel::SetTimeCodes(UsdTimeCode displayTime, UsdTimeCode editTime)
{
    m_displayTime = displayTime;
    m_editTime    = editTime;
}

void CurveEditorPanel::RefreshFromStage()
{
    m_needsRefresh = false;
    m_channels.clear();
    m_selection.clear();
    m_activeChannel = -1;
    s_scalarColorIdx = 0;

    if (!m_stage || m_selectedPrimPath.empty()) return;

    UsdPrim prim = m_stage->GetPrimAtPath(SdfPath(m_selectedPrimPath));
    if (!prim.IsValid()) return;

    static const char* kComponents[] = {"X","Y","Z","W"};

    for (const UsdAttribute& attr : prim.GetAuthoredAttributes()) {
        std::vector<double> times;
        if (!attr.GetTimeSamples(&times) || times.empty()) continue;

        auto tn = attr.GetTypeName();
        std::string baseName = attr.GetName().GetString();

        // Determine component count
        int numComp = 1;
        bool isScalar = true;
        if (tn == SdfValueTypeNames->Float2  || tn == SdfValueTypeNames->Double2)
            { numComp = 2; isScalar = false; }
        else if (tn == SdfValueTypeNames->Float3  || tn == SdfValueTypeNames->Double3 ||
                 tn == SdfValueTypeNames->Vector3f || tn == SdfValueTypeNames->Vector3d ||
                 tn == SdfValueTypeNames->Point3f  || tn == SdfValueTypeNames->Point3d  ||
                 tn == SdfValueTypeNames->Normal3f || tn == SdfValueTypeNames->Normal3d ||
                 tn == SdfValueTypeNames->Color3f)
            { numComp = 3; isScalar = false; }
        else if (tn == SdfValueTypeNames->Float4  || tn == SdfValueTypeNames->Double4 ||
                 tn == SdfValueTypeNames->Color4f)
            { numComp = 4; isScalar = false; }
        else if (tn != SdfValueTypeNames->Float   && tn != SdfValueTypeNames->Double &&
                 tn != SdfValueTypeNames->Half    && tn != SdfValueTypeNames->Int    &&
                 tn != SdfValueTypeNames->Bool    && tn != SdfValueTypeNames->Int64  &&
                 tn != SdfValueTypeNames->UInt    && tn != SdfValueTypeNames->UInt64) {
            continue; // unsupported type
        }

        for (int ci = 0; ci < numComp; ++ci) {
            CurveChannel ch;
            ch.attr      = attr;
            ch.component = isScalar ? -1 : ci;
            if (isScalar) {
                ch.displayName = baseName;
                ch.color = kScalarPalette[s_scalarColorIdx++ % 5];
            } else {
                ch.displayName = baseName + " [" + kComponents[ci] + "]";
                ch.color = kComponentColors[ci];
            }
            FitCurveFromSamples(ch);
            m_channels.push_back(std::move(ch));
        }
    }

    if (!m_channels.empty()) FrameAll();
}

// ── Bezier metadata persistence ───────────────────────────────────────────────
//
// Stored under prim.customData["curveEditor"]["channels"][channelKey].
// The channel key uses '@' to separate attribute name from component index
// (avoiding the ':' path separator used by SetCustomDataByKey).
//
// Saved USDA example:
//   customData = {
//       dictionary curveEditor = {
//           dictionary channels = {
//               dictionary "xformOp:translate@0" = {
//                   double[] times  = [1, 12, 24]
//                   double[] values = [0, 3, 0]
//                   double[] inDt   = [-4, -4, -4]
//                   double[] inDv   = [0, -1, 0]
//                   double[] outDt  = [4, 4, 4]
//                   double[] outDv  = [0, 1, 0]
//                   int[]    broken = [0, 0, 0]
//               }
//           }
//       }
//   }

/*static*/
std::string CurveEditorPanel::ChannelMetadataKey(const CurveChannel& ch)
{
    return ch.attr.GetName().GetString() + "@" + std::to_string(ch.component);
}

bool CurveEditorPanel::SaveBezierToMetadata(const UsdPrim& prim,
                                            const std::string& channelKey,
                                            const std::vector<BezierKey>& keys) const
{
    if (!prim || keys.empty()) return false;

    VtArray<double> times, values, inDt, inDv, outDt, outDv;
    VtArray<int>    broken;
    times.reserve(keys.size());
    for (const BezierKey& k : keys) {
        times .push_back(k.time);
        values.push_back(k.value);
        inDt  .push_back(k.inTangentDt);
        inDv  .push_back(k.inTangentDv);
        outDt .push_back(k.outTangentDt);
        outDv .push_back(k.outTangentDv);
        broken.push_back(k.tangentsBroken ? 1 : 0);
    }

    VtDictionary keyData;
    keyData["times"]  = VtValue(times);
    keyData["values"] = VtValue(values);
    keyData["inDt"]   = VtValue(inDt);
    keyData["inDv"]   = VtValue(inDv);
    keyData["outDt"]  = VtValue(outDt);
    keyData["outDv"]  = VtValue(outDv);
    keyData["broken"] = VtValue(broken);

    VtDictionary ceDict;
    VtValue existing = prim.GetCustomDataByKey(TfToken("curveEditor"));
    if (existing.IsHolding<VtDictionary>())
        ceDict = existing.UncheckedGet<VtDictionary>();

    VtDictionary channels;
    {
        auto it = ceDict.find("channels");
        if (it != ceDict.end() && it->second.IsHolding<VtDictionary>())
            channels = it->second.UncheckedGet<VtDictionary>();
    }

    channels[channelKey] = VtValue(keyData);
    ceDict["channels"] = VtValue(channels);
    prim.SetCustomDataByKey(TfToken("curveEditor"), VtValue(ceDict));
    return true;
}

bool CurveEditorPanel::SaveBezierToMetadata(const UsdPrim& prim,
                                            const CurveChannel& ch) const
{
    return SaveBezierToMetadata(prim, ChannelMetadataKey(ch), ch.keys);
}

bool CurveEditorPanel::LoadBezierFromMetadata(const UsdPrim& prim,
                                              CurveChannel& ch) const
{
    if (!prim) return false;

    VtValue ceVal = prim.GetCustomDataByKey(TfToken("curveEditor"));
    if (!ceVal.IsHolding<VtDictionary>()) return false;
    const VtDictionary& ceDict = ceVal.UncheckedGet<VtDictionary>();

    auto chIt = ceDict.find("channels");
    if (chIt == ceDict.end() || !chIt->second.IsHolding<VtDictionary>()) return false;
    const VtDictionary& channels = chIt->second.UncheckedGet<VtDictionary>();

    auto kIt = channels.find(ChannelMetadataKey(ch));
    if (kIt == channels.end() || !kIt->second.IsHolding<VtDictionary>()) return false;
    const VtDictionary& kd = kIt->second.UncheckedGet<VtDictionary>();

    auto getDoubleArr = [&](const char* name) -> VtArray<double> {
        auto it = kd.find(name);
        if (it != kd.end() && it->second.IsHolding<VtArray<double>>())
            return it->second.UncheckedGet<VtArray<double>>();
        return {};
    };
    auto getIntArr = [&](const char* name) -> VtArray<int> {
        auto it = kd.find(name);
        if (it != kd.end() && it->second.IsHolding<VtArray<int>>())
            return it->second.UncheckedGet<VtArray<int>>();
        return {};
    };

    VtArray<double> times  = getDoubleArr("times");
    VtArray<double> values = getDoubleArr("values");
    VtArray<double> inDt   = getDoubleArr("inDt");
    VtArray<double> inDv   = getDoubleArr("inDv");
    VtArray<double> outDt  = getDoubleArr("outDt");
    VtArray<double> outDv  = getDoubleArr("outDv");
    VtArray<int>    broken = getIntArr("broken");

    size_t n = times.size();
    if (n == 0 || values.size() != n) return false;

    ch.keys.clear();
    ch.keys.reserve(n);
    for (size_t i = 0; i < n; ++i) {
        BezierKey k;
        k.time           = times[i];
        k.value          = values[i];
        k.inTangentDt    = (inDt.size()  > i) ? inDt[i]  : -1.0;
        k.inTangentDv    = (inDv.size()  > i) ? inDv[i]  : 0.0;
        k.outTangentDt   = (outDt.size() > i) ? outDt[i] : 1.0;
        k.outTangentDv   = (outDv.size() > i) ? outDv[i] : 0.0;
        k.tangentsBroken = (broken.size() > i) ? broken[i] != 0 : false;
        ch.keys.push_back(k);
    }
    return true;
}

// ── Catmull-Rom fit ───────────────────────────────────────────────────────────

void CurveEditorPanel::FitCurveFromSamples(CurveChannel& ch)
{
    std::vector<double> times;
    ch.attr.GetTimeSamples(&times);

    ch.keys.clear();
    ch.origTimes.clear();
    ch.origValues.clear();

    if (times.empty()) { ch.dirty = false; return; }

    // Read values
    std::vector<double> vals;
    vals.reserve(times.size());
    for (double t : times)
        vals.push_back(ReadChannelValue(ch.attr, ch.component, t));

    ch.origTimes  = times;
    ch.origValues = vals;

    // Prefer saved Bezier tangents over Catmull-Rom re-fit
    if (m_stage && !m_selectedPrimPath.empty()) {
        UsdPrim prim = m_stage->GetPrimAtPath(SdfPath(m_selectedPrimPath));
        if (prim && LoadBezierFromMetadata(prim, ch)) {
            ch.dirty = false;
            return;
        }
    }

    int n = int(times.size());
    for (int i = 0; i < n; ++i) {
        BezierKey k;
        k.time  = times[i];
        k.value = vals[i];

        // Catmull-Rom slope estimate
        double slopePrev = (i > 0)
            ? (vals[i] - vals[i-1]) / (times[i] - times[i-1]) : 0.0;
        double slopeNext = (i < n-1)
            ? (vals[i+1] - vals[i]) / (times[i+1] - times[i]) : 0.0;

        double slope;
        if (i == 0)   slope = slopeNext;
        else if (i == n-1) slope = slopePrev;
        else          slope = (slopePrev + slopeNext) * 0.5;

        double dtPrev = (i > 0)   ? (times[i] - times[i-1]) / 3.0 : 1.0;
        double dtNext = (i < n-1) ? (times[i+1] - times[i]) / 3.0 : 1.0;

        k.inTangentDt  = -dtPrev;  k.inTangentDv  = -dtPrev  * slope;
        k.outTangentDt =  dtNext;  k.outTangentDv =  dtNext   * slope;

        ch.keys.push_back(k);
    }
    ch.dirty = false;
}

void CurveEditorPanel::RefitNeighborTangents(CurveChannel& ch, int keyIdx)
{
    int n = int(ch.keys.size());
    auto refit = [&](int i) {
        if (i < 0 || i >= n) return;
        BezierKey& k = ch.keys[i];
        if (k.tangentsBroken) return;

        double slopePrev = (i > 0)
            ? (k.value - ch.keys[i-1].value) / (k.time - ch.keys[i-1].time) : 0.0;
        double slopeNext = (i < n-1)
            ? (ch.keys[i+1].value - k.value) / (ch.keys[i+1].time - k.time) : 0.0;

        double slope;
        if (i == 0)   slope = slopeNext;
        else if (i == n-1) slope = slopePrev;
        else          slope = (slopePrev + slopeNext) * 0.5;

        double dtPrev = (i > 0)   ? (k.time - ch.keys[i-1].time) / 3.0 : 1.0;
        double dtNext = (i < n-1) ? (ch.keys[i+1].time - k.time) / 3.0 : 1.0;

        k.inTangentDt  = -dtPrev; k.inTangentDv  = -dtPrev  * slope;
        k.outTangentDt =  dtNext; k.outTangentDv =  dtNext   * slope;
    };
    refit(keyIdx - 1);
    refit(keyIdx);
    refit(keyIdx + 1);
}

// ── Bezier evaluation ────────────────────────────────────────────────────────

// Evaluate cubic Bezier in 1D.
static double CubicBez1D(double p0, double p1, double p2, double p3, double u)
{
    double inv = 1.0 - u;
    return inv*inv*inv*p0 + 3.0*inv*inv*u*p1 + 3.0*inv*u*u*p2 + u*u*u*p3;
}

double CurveEditorPanel::EvaluateBezier(const std::vector<BezierKey>& keys, double T) const
{
    int n = int(keys.size());
    if (n == 0) return 0.0;
    if (n == 1 || T <= keys.front().time) return keys.front().value;
    if (T >= keys.back().time)            return keys.back().value;

    // Find segment
    int seg = 0;
    for (int i = 0; i < n-1; ++i) {
        if (T < keys[i+1].time) { seg = i; break; }
    }
    const BezierKey& k0 = keys[seg];
    const BezierKey& k1 = keys[seg+1];

    // Control points (time axis)
    double t0 = k0.time, t1 = k0.time + k0.outTangentDt;
    double t2 = k1.time + k1.inTangentDt, t3 = k1.time;

    // Newton-Raphson: find u s.t. BezierX(u) == T
    double u = (T - t0) / (t3 - t0);
    for (int iter = 0; iter < 8; ++iter) {
        double bt = CubicBez1D(t0, t1, t2, t3, u);
        double dt0 = 3.0*(1-u)*(1-u)*(t1-t0);
        double dt1 = 6.0*(1-u)*u*(t2-t1);
        double dt2 = 3.0*u*u*(t3-t2);
        double deriv = dt0 + dt1 + dt2;
        if (std::fabs(deriv) < 1e-10) break;
        u -= (bt - T) / deriv;
        u  = std::max(0.0, std::min(1.0, u));
    }

    // Evaluate value axis
    double v0 = k0.value, v1 = k0.value + k0.outTangentDv;
    double v2 = k1.value + k1.inTangentDv, v3 = k1.value;
    return CubicBez1D(v0, v1, v2, v3, u);
}

// ── Bake ─────────────────────────────────────────────────────────────────────

static void SimplifyCurve(std::vector<double>& times, std::vector<double>& vals,
                           double tol)
{
    // Ramer-Douglas-Peucker on (time[i], vals[i])
    int n = int(times.size());
    if (n <= 2) return;

    std::vector<bool> keep(n, false);
    keep[0] = keep[n-1] = true;

    std::function<void(int,int)> rdp = [&](int lo, int hi) {
        if (hi - lo < 2) return;
        double dx = times[hi] - times[lo], dy = vals[hi] - vals[lo];
        double len = std::sqrt(dx*dx + dy*dy);
        double maxDist = 0.0; int maxIdx = lo+1;
        for (int i = lo+1; i < hi; ++i) {
            double px = times[i] - times[lo], py = vals[i] - vals[lo];
            double dist = (len > 1e-10) ? std::fabs(px*dy - py*dx) / len
                                        : std::sqrt(px*px + py*py);
            if (dist > maxDist) { maxDist = dist; maxIdx = i; }
        }
        if (maxDist > tol) {
            keep[maxIdx] = true;
            rdp(lo, maxIdx);
            rdp(maxIdx, hi);
        }
    };
    rdp(0, n-1);

    std::vector<double> nt, nv;
    for (int i = 0; i < n; ++i)
        if (keep[i]) { nt.push_back(times[i]); nv.push_back(vals[i]); }
    times = std::move(nt);
    vals  = std::move(nv);
}

void CurveEditorPanel::BakeChannelToUsd(CurveChannel& ch)
{
    if (ch.keys.empty()) return;

    auto oldTimes  = ch.origTimes;
    auto oldValues = ch.origValues;

    int startF = int(std::floor(ch.keys.front().time));
    int endF   = int(std::ceil (ch.keys.back().time));

    std::vector<double> newTimes, newValues;
    newTimes.reserve(endF - startF + 1);
    newValues.reserve(endF - startF + 1);
    for (int f = startF; f <= endF; ++f) {
        newTimes .push_back(double(f));
        newValues.push_back(EvaluateBezier(ch.keys, double(f)));
    }
    if (m_simplifyOnBake && newTimes.size() > 2)
        SimplifyCurve(newTimes, newValues, double(m_simplifyTol));

    UsdAttribute attr = ch.attr;
    int comp = ch.component;

    auto newKeys = ch.keys;
    std::string channelKey = ChannelMetadataKey(ch);
    UsdPrim prim = m_stage ? m_stage->GetPrimAtPath(SdfPath(m_selectedPrimPath))
                            : UsdPrim{};
    VtValue oldCEMeta;
    if (prim) oldCEMeta = prim.GetCustomDataByKey(TfToken("curveEditor"));

    // Safe channel finder: looks up by attr path + component at closure call time,
    // so it remains valid even if m_channels is rebuilt between bake and undo.
    SdfPath chanAttrPath = attr.GetPath();
    auto findCh = [this, chanAttrPath, comp]() -> CurveChannel* {
        for (auto& c : m_channels)
            if (c.attr.GetPath() == chanAttrPath && c.component == comp) return &c;
        return nullptr;
    };

    if (!m_commandHistory) {
        m_suppressNotice = true;
        ClearAllTimeSamples(attr);
        for (int i = 0; i < int(newTimes.size()); ++i)
            WriteChannelValue(attr, comp, newTimes[i], newValues[i]);
        m_suppressNotice = false;
        ch.origTimes  = newTimes;
        ch.origValues = newValues;
        ch.dirty = false;
        if (prim) SaveBezierToMetadata(prim, channelKey, newKeys);
        return;
    }

    m_commandHistory->Push(std::make_unique<AttributeSetCommand>(
        "Bake F-Curve: " + ch.displayName,
        [this, attr, comp, newTimes, newValues, newKeys, channelKey, prim, findCh]() {
            m_suppressNotice = true;
            ClearAllTimeSamples(attr);
            for (int i = 0; i < int(newTimes.size()); ++i)
                WriteChannelValue(attr, comp, newTimes[i], newValues[i]);
            m_suppressNotice = false;
            if (prim) SaveBezierToMetadata(prim, channelKey, newKeys);
            if (auto* c = findCh()) {
                c->origTimes  = newTimes;
                c->origValues = newValues;
                c->keys       = newKeys;
                c->dirty      = false;
            } else {
                m_needsRefresh = true;
            }
        },
        [this, attr, comp, oldTimes, oldValues, oldCEMeta, prim, findCh]() {
            m_suppressNotice = true;
            ClearAllTimeSamples(attr);
            for (int i = 0; i < int(oldTimes.size()); ++i)
                WriteChannelValue(attr, comp, oldTimes[i], oldValues[i]);
            m_suppressNotice = false;
            if (prim) {
                if (oldCEMeta.IsEmpty())
                    prim.ClearCustomDataByKey(TfToken("curveEditor"));
                else
                    prim.SetCustomDataByKey(TfToken("curveEditor"), oldCEMeta);
            }
            if (auto* c = findCh()) {
                c->origTimes  = oldTimes;
                c->origValues = oldValues;
                FitCurveFromSamples(*c);
            } else {
                m_needsRefresh = true;
            }
        }
    ));
}

// ── Keyframe operations ───────────────────────────────────────────────────────

void CurveEditorPanel::AddKeyAt(int chanIdx, double time, double value)
{
    if (chanIdx < 0 || chanIdx >= int(m_channels.size())) return;
    CurveChannel& ch = m_channels[chanIdx];

    // Insert sorted by time
    BezierKey k;
    k.time  = time;
    k.value = value;
    k.inTangentDt  = -1.0; k.inTangentDv  = 0.0;
    k.outTangentDt =  1.0; k.outTangentDv = 0.0;

    auto it = std::lower_bound(ch.keys.begin(), ch.keys.end(), k,
        [](const BezierKey& a, const BezierKey& b){ return a.time < b.time; });
    int idx = int(it - ch.keys.begin());
    ch.keys.insert(it, k);
    RefitNeighborTangents(ch, idx);
    ch.dirty = true;

    // Select the new key
    ClearSelection();
    m_selection.push_back({chanIdx, idx, HandlePart::Key});
}

void CurveEditorPanel::DeleteSelectedKeys()
{
    // Collect per-channel indices (in reverse so erasure doesn't shift indices)
    struct ToDelete { int chanIdx; std::vector<int> keyIndices; };
    std::vector<ToDelete> byChannel;

    for (const SelectedHandle& s : m_selection) {
        if (s.part != HandlePart::Key) continue;
        bool found = false;
        for (auto& td : byChannel)
            if (td.chanIdx == s.chanIdx) { td.keyIndices.push_back(s.keyIdx); found = true; break; }
        if (!found) byChannel.push_back({s.chanIdx, {s.keyIdx}});
    }

    ClearSelection();
    for (auto& td : byChannel) {
        CurveChannel& ch = m_channels[td.chanIdx];
        std::sort(td.keyIndices.rbegin(), td.keyIndices.rend());
        for (int ki : td.keyIndices) {
            if (ki >= 0 && ki < int(ch.keys.size()))
                ch.keys.erase(ch.keys.begin() + ki);
        }
        ch.dirty = true;
    }
}

void CurveEditorPanel::ApplyAllDirty()
{
    for (auto& ch : m_channels)
        if (ch.dirty) BakeChannelToUsd(ch);
}

void CurveEditorPanel::RevertAll()
{
    for (auto& ch : m_channels)
        FitCurveFromSamples(ch);
    ClearSelection();
}

// ── Selection ────────────────────────────────────────────────────────────────

bool CurveEditorPanel::IsSelected(int ci, int ki, HandlePart p) const
{
    for (const auto& s : m_selection)
        if (s.chanIdx == ci && s.keyIdx == ki && s.part == p) return true;
    return false;
}

void CurveEditorPanel::Select(int ci, int ki, HandlePart p, bool additive)
{
    if (!additive) m_selection.clear();
    if (!IsSelected(ci, ki, p))
        m_selection.push_back({ci, ki, p});
}

void CurveEditorPanel::ClearSelection() { m_selection.clear(); }

// ── Coordinate helpers ────────────────────────────────────────────────────────

float CurveEditorPanel::TimeToX(double t, ImVec2 cMin, ImVec2 cMax) const
{
    double span = m_viewMaxTime - m_viewMinTime;
    if (span < 1e-10) return cMin.x;
    return cMin.x + float((t - m_viewMinTime) / span) * (cMax.x - cMin.x);
}
float CurveEditorPanel::ValToY(double v, ImVec2 cMin, ImVec2 cMax) const
{
    double span = m_viewMaxVal - m_viewMinVal;
    if (span < 1e-10) return cMax.y;
    return cMax.y - float((v - m_viewMinVal) / span) * (cMax.y - cMin.y);
}
double CurveEditorPanel::XToTime(float x, ImVec2 cMin, ImVec2 cMax) const
{
    float w = cMax.x - cMin.x;
    if (w < 1.f) return m_viewMinTime;
    return m_viewMinTime + double(x - cMin.x) / double(w) * (m_viewMaxTime - m_viewMinTime);
}
double CurveEditorPanel::YToVal(float y, ImVec2 cMin, ImVec2 cMax) const
{
    float h = cMax.y - cMin.y;
    if (h < 1.f) return m_viewMinVal;
    return m_viewMinVal + double(cMax.y - y) / double(h) * (m_viewMaxVal - m_viewMinVal);
}
ImVec2 CurveEditorPanel::KeyToScreen(const BezierKey& k, ImVec2 cMin, ImVec2 cMax) const
{
    return {TimeToX(k.time, cMin, cMax), ValToY(k.value, cMin, cMax)};
}

void CurveEditorPanel::FrameAll()
{
    if (m_channels.empty()) {
        m_viewMinTime = 0.0; m_viewMaxTime = 100.0;
        m_viewMinVal  =-1.0; m_viewMaxVal  =  1.0;
        return;
    }
    double tMin =  1e30, tMax = -1e30;
    double vMin =  1e30, vMax = -1e30;
    for (const auto& ch : m_channels) {
        if (!ch.visible) continue;
        for (const auto& k : ch.keys) {
            tMin = std::min(tMin, k.time);  tMax = std::max(tMax, k.time);
            vMin = std::min(vMin, k.value); vMax = std::max(vMax, k.value);
        }
    }
    if (tMin > tMax) { tMin = 0; tMax = 100; }
    if (vMin > vMax) { vMin = -1; vMax = 1;  }
    double tp = (tMax - tMin) * 0.1 + 1.0;
    double vp = (vMax - vMin) * 0.15 + 0.1;
    m_viewMinTime = tMin - tp; m_viewMaxTime = tMax + tp;
    m_viewMinVal  = vMin - vp; m_viewMaxVal  = vMax + vp;
}

// ── Grid ─────────────────────────────────────────────────────────────────────

static double NiceStep(double span, int targetLines)
{
    if (span <= 0 || targetLines <= 0) return 1.0;
    double rough = span / targetLines;
    double mag   = std::pow(10.0, std::floor(std::log10(rough)));
    double norm  = rough / mag;
    double nice  = norm < 1.5 ? 1.0 : norm < 3.5 ? 2.0 : norm < 7.5 ? 5.0 : 10.0;
    return nice * mag;
}

void CurveEditorPanel::RenderGrid(ImDrawList* dl, ImVec2 cMin, ImVec2 cMax)
{
    ImU32 gridColor = IM_COL32(60, 60, 60, 255);
    ImU32 axisColor = IM_COL32(100, 100, 100, 255);
    ImU32 textColor = IM_COL32(140, 140, 140, 255);
    float fontSize  = ImGui::GetFontSize();

    // Vertical (time)
    double tStep = NiceStep(m_viewMaxTime - m_viewMinTime, 8);
    double tStart = std::floor(m_viewMinTime / tStep) * tStep;
    for (double t = tStart; t <= m_viewMaxTime + tStep; t += tStep) {
        float x = TimeToX(t, cMin, cMax);
        if (x < cMin.x - 1.f || x > cMax.x + 1.f) continue;
        ImU32 col = std::fabs(t) < tStep * 0.01 ? axisColor : gridColor;
        dl->AddLine({x, cMin.y}, {x, cMax.y}, col);
        char buf[32]; snprintf(buf, sizeof(buf), "%.4g", t);
        dl->AddText({x + 2.f, cMax.y - fontSize - 2.f}, textColor, buf);
    }

    // Horizontal (value)
    double vStep = NiceStep(m_viewMaxVal - m_viewMinVal, 6);
    double vStart = std::floor(m_viewMinVal / vStep) * vStep;
    for (double v = vStart; v <= m_viewMaxVal + vStep; v += vStep) {
        float y = ValToY(v, cMin, cMax);
        if (y < cMin.y - 1.f || y > cMax.y + 1.f) continue;
        ImU32 col = std::fabs(v) < vStep * 0.01 ? axisColor : gridColor;
        dl->AddLine({cMin.x, y}, {cMax.x, y}, col);
        char buf[32]; snprintf(buf, sizeof(buf), "%.4g", v);
        dl->AddText({cMin.x + 2.f, y - fontSize - 1.f}, textColor, buf);
    }
}

// ── Curve drawing ────────────────────────────────────────────────────────────

void CurveEditorPanel::RenderCurve(ImDrawList* dl, CurveChannel& ch,
                                   ImVec2 cMin, ImVec2 cMax)
{
    if (!ch.visible || ch.keys.empty()) return;

    ImU32 lineCol = ImGui::ColorConvertFloat4ToU32(ch.color);
    ImU32 keyCol  = ImGui::ColorConvertFloat4ToU32(ch.color);
    ImU32 selCol  = IM_COL32(255, 220, 60, 255);
    ImU32 handleLineCol = IM_COL32(180, 180, 180, 180);

    int chanIdx = int(&ch - m_channels.data());

    // Draw curve segments
    for (int i = 0; i < int(ch.keys.size()) - 1; ++i) {
        const BezierKey& k0 = ch.keys[i];
        const BezierKey& k1 = ch.keys[i+1];
        // Subsample Bezier for smooth display
        static const int kSubs = 12;
        ImVec2 prev = KeyToScreen(k0, cMin, cMax);
        for (int s = 1; s <= kSubs; ++s) {
            double u = double(s) / double(kSubs);
            double T = k0.time + u * (k1.time - k0.time);
            double V = EvaluateBezier(ch.keys, T);
            ImVec2 cur = {TimeToX(T, cMin, cMax), ValToY(V, cMin, cMax)};
            dl->AddLine(prev, cur, lineCol, 1.5f);
            prev = cur;
        }
    }

    // Draw keys and tangent handles
    for (int ki = 0; ki < int(ch.keys.size()); ++ki) {
        const BezierKey& k = ch.keys[ki];
        ImVec2 kp = KeyToScreen(k, cMin, cMax);
        bool keySel = IsSelected(chanIdx, ki, HandlePart::Key);

        // Tangent handles — always visible
        if (ki > 0) {
            ImVec2 ih = {TimeToX(k.time + k.inTangentDt,  cMin, cMax),
                         ValToY (k.value + k.inTangentDv, cMin, cMax)};
            dl->AddLine(kp, ih, handleLineCol, 1.0f);
            bool hSel = IsSelected(chanIdx, ki, HandlePart::InTangent);
            dl->AddCircleFilled(ih, 4.f, hSel ? selCol : IM_COL32(200,200,200,220));
        }
        if (ki < int(ch.keys.size())-1) {
            ImVec2 oh = {TimeToX(k.time + k.outTangentDt,  cMin, cMax),
                         ValToY (k.value + k.outTangentDv, cMin, cMax)};
            dl->AddLine(kp, oh, handleLineCol, 1.0f);
            bool hSel = IsSelected(chanIdx, ki, HandlePart::OutTangent);
            dl->AddCircleFilled(oh, 4.f, hSel ? selCol : IM_COL32(200,200,200,220));
        }

        // Diamond keyframe marker
        float r = keySel ? 6.f : 4.5f;
        ImVec2 d[4] = {
            {kp.x,   kp.y-r},
            {kp.x+r, kp.y  },
            {kp.x,   kp.y+r},
            {kp.x-r, kp.y  },
        };
        if (keySel)
            dl->AddConvexPolyFilled(d, 4, selCol);
        else
            dl->AddConvexPolyFilled(d, 4, IM_COL32(40,40,40,255));
        dl->AddPolyline(d, 4, keyCol, ImDrawFlags_Closed, 1.5f);
    }
}

// ── Time cursor ───────────────────────────────────────────────────────────────

void CurveEditorPanel::RenderTimeCursor(ImDrawList* dl, ImVec2 cMin, ImVec2 cMax)
{
    if (m_displayTime.IsDefault()) return;
    float cx = TimeToX(m_displayTime.GetValue(), cMin, cMax);
    dl->AddLine({cx, cMin.y}, {cx, cMax.y}, IM_COL32(255, 200, 60, 200), 1.5f);
    // Small triangle at top
    dl->AddTriangleFilled({cx-5.f, cMin.y}, {cx+5.f, cMin.y}, {cx, cMin.y+9.f},
                          IM_COL32(255, 200, 60, 220));
}

// ── Canvas interaction ────────────────────────────────────────────────────────

void CurveEditorPanel::HandleCanvasInput(ImVec2 cMin, ImVec2 cMax)
{
    ImGuiIO& io = ImGui::GetIO();
    ImVec2 mouse = io.MousePos;
    bool hovered = (mouse.x >= cMin.x && mouse.x <= cMax.x &&
                    mouse.y >= cMin.y && mouse.y <= cMax.y);

    // ── Hit test helpers — defined early so the RMB popup can use them ──────────
    const float kKeyRadius    = 8.f;
    const float kHandleRadius = 7.f;

    auto hitTestKey = [&](ImVec2 p) -> std::pair<int,int> {
        for (int ci = 0; ci < int(m_channels.size()); ++ci) {
            auto& ch = m_channels[ci];
            if (!ch.visible) continue;
            for (int ki = 0; ki < int(ch.keys.size()); ++ki) {
                ImVec2 kp = KeyToScreen(ch.keys[ki], cMin, cMax);
                float dx = p.x-kp.x, dy = p.y-kp.y;
                if (dx*dx+dy*dy < kKeyRadius*kKeyRadius) return {ci, ki};
            }
        }
        return {-1,-1};
    };

    auto hitTestHandle = [&](ImVec2 p) -> SelectedHandle {
        for (int ci = 0; ci < int(m_channels.size()); ++ci) {
            const CurveChannel& ch = m_channels[ci];
            if (!ch.visible) continue;
            for (int ki = 0; ki < int(ch.keys.size()); ++ki) {
                const BezierKey& k = ch.keys[ki];
                if (ki > 0) {
                    ImVec2 ih = {TimeToX(k.time+k.inTangentDt, cMin, cMax),
                                 ValToY (k.value+k.inTangentDv, cMin, cMax)};
                    float dx=p.x-ih.x, dy=p.y-ih.y;
                    if (dx*dx+dy*dy < kHandleRadius*kHandleRadius)
                        return {ci, ki, HandlePart::InTangent};
                }
                if (ki < int(ch.keys.size())-1) {
                    ImVec2 oh = {TimeToX(k.time+k.outTangentDt, cMin, cMax),
                                 ValToY (k.value+k.outTangentDv, cMin, cMax)};
                    float dx=p.x-oh.x, dy=p.y-oh.y;
                    if (dx*dx+dy*dy < kHandleRadius*kHandleRadius)
                        return {ci, ki, HandlePart::OutTangent};
                }
            }
        }
        return {-1,-1,HandlePart::Key};
    };

    // ── RMB context menu — must be processed before the early-return so the
    // popup stays alive after the mouse leaves the canvas into the menu. ────────
    if (ImGui::IsMouseClicked(1) && hovered) {
        auto [ci, ki] = hitTestKey(mouse);
        if (ci >= 0) {
            if (!IsSelected(ci, ki, HandlePart::Key))
                Select(ci, ki, HandlePart::Key, false);
            ImGui::OpenPopup("##keyCtx");
        }
    }
    if (ImGui::BeginPopup("##keyCtx")) {
        if (ImGui::MenuItem("Delete Key"))
            DeleteSelectedKeys();
        ImGui::Separator();
        if (ImGui::MenuItem("Flatten Tangents")) {
            for (const auto& s : m_selection) {
                if (s.part != HandlePart::Key) continue;
                if (s.chanIdx < 0 || s.chanIdx >= int(m_channels.size())) continue;
                CurveChannel& ch = m_channels[s.chanIdx];
                if (s.keyIdx < 0 || s.keyIdx >= int(ch.keys.size())) continue;
                BezierKey& bk = ch.keys[s.keyIdx];
                bk.inTangentDv  = 0.0;
                bk.outTangentDv = 0.0;
                ch.dirty = true;
            }
        }
        if (ImGui::MenuItem("Break Tangents")) {
            for (const auto& s : m_selection) {
                if (s.part != HandlePart::Key) continue;
                if (s.chanIdx < 0 || s.chanIdx >= int(m_channels.size())) continue;
                CurveChannel& ch = m_channels[s.chanIdx];
                if (s.keyIdx < 0 || s.keyIdx >= int(ch.keys.size())) continue;
                ch.keys[s.keyIdx].tangentsBroken = true;
                ch.dirty = true;
            }
        }
        if (ImGui::MenuItem("Unify Tangents")) {
            for (const auto& s : m_selection) {
                if (s.part != HandlePart::Key) continue;
                if (s.chanIdx < 0 || s.chanIdx >= int(m_channels.size())) continue;
                CurveChannel& ch = m_channels[s.chanIdx];
                if (s.keyIdx < 0 || s.keyIdx >= int(ch.keys.size())) continue;
                BezierKey& bk = ch.keys[s.keyIdx];
                bk.tangentsBroken = false;
                bk.inTangentDv = -bk.outTangentDv *
                    ((-bk.inTangentDt > 1e-10) ? -bk.inTangentDt/bk.outTangentDt : 1.0);
                ch.dirty = true;
            }
        }
        if (ImGui::MenuItem("Auto Tangents")) {
            for (const auto& s : m_selection) {
                if (s.part != HandlePart::Key) continue;
                if (s.chanIdx < 0 || s.chanIdx >= int(m_channels.size())) continue;
                RefitNeighborTangents(m_channels[s.chanIdx], s.keyIdx);
                m_channels[s.chanIdx].dirty = true;
            }
        }
        ImGui::EndPopup();
    }

    if (!hovered && !m_isPanning && !m_isDragging && !m_isDraggingCursor
                 && !m_isBoxSelecting)
        return;

    // ── Time cursor drag ─────────────────────────────────────────────────────
    if (!m_displayTime.IsDefault()) {
        float cx = TimeToX(m_displayTime.GetValue(), cMin, cMax);
        bool nearCursor = (std::fabs(mouse.x - cx) < 8.f && mouse.y >= cMin.y &&
                           mouse.y <= cMin.y + 18.f);
        if (nearCursor && ImGui::IsMouseClicked(0))
            m_isDraggingCursor = true;
        if (m_isDraggingCursor && ImGui::IsMouseDown(0)) {
            double t = XToTime(mouse.x, cMin, cMax);
            if (m_snapToFrame) t = std::round(t);
            if (OnScrub) OnScrub(t);
        }
        if (ImGui::IsMouseReleased(0)) m_isDraggingCursor = false;
    }
    if (m_isDraggingCursor) return;

    // ── Pan: MMB or Alt+LMB ──────────────────────────────────────────────────
    bool wantPan = ImGui::IsMouseDown(2) ||
                   (ImGui::IsMouseDown(0) && io.KeyAlt);

    if (!m_isDragging && !m_isBoxSelecting) {
        if (hovered && wantPan && !m_isPanning) {
            m_isPanning = true;
            m_panLastMouse = mouse;
        }
        if (m_isPanning && wantPan) {
            float dx = mouse.x - m_panLastMouse.x;
            float dy = mouse.y - m_panLastMouse.y;
            float cw = cMax.x - cMin.x, ch2 = cMax.y - cMin.y;
            if (cw > 1.f) {
                double dtSpan = (m_viewMaxTime - m_viewMinTime) * double(dx) / double(cw);
                m_viewMinTime -= dtSpan; m_viewMaxTime -= dtSpan;
            }
            if (ch2 > 1.f) {
                double dvSpan = (m_viewMaxVal - m_viewMinVal) * double(dy) / double(ch2);
                m_viewMinVal += dvSpan; m_viewMaxVal += dvSpan;
            }
            m_panLastMouse = mouse;
        }
        if (!wantPan) m_isPanning = false;
    }

    // ── Zoom: scroll wheel ───────────────────────────────────────────────────
    if (hovered && std::fabs(io.MouseWheel) > 0.f && !m_isPanning) {
        float wheel = io.MouseWheel;
        double factor = std::pow(0.9, double(wheel));
        if (io.KeyShift) {
            // Zoom Y
            double pivot = YToVal(mouse.y, cMin, cMax);
            m_viewMinVal = pivot + (m_viewMinVal - pivot) * factor;
            m_viewMaxVal = pivot + (m_viewMaxVal - pivot) * factor;
        } else {
            // Zoom X
            double pivot = XToTime(mouse.x, cMin, cMax);
            m_viewMinTime = pivot + (m_viewMinTime - pivot) * factor;
            m_viewMaxTime = pivot + (m_viewMaxTime - pivot) * factor;
        }
    }

    // ── Frame all: F key ─────────────────────────────────────────────────────
    if (hovered && ImGui::IsKeyPressed(ImGuiKey_F, false))
        FrameAll();

    // ── Delete selected keys ──────────────────────────────────────────────────
    if (hovered && ImGui::IsKeyPressed(ImGuiKey_Delete, false))
        DeleteSelectedKeys();

    // ── LMB interactions ─────────────────────────────────────────────────────
    if (ImGui::IsMouseClicked(0) && hovered && !io.KeyAlt) {
        // Hit test handle first, then key, then start box-select
        SelectedHandle hitH = hitTestHandle(mouse);
        if (hitH.chanIdx >= 0) {
            bool additive = io.KeyShift || io.KeyCtrl;
            // Ensure the parent key is also selected (needed for context menu / delete)
            if (!IsSelected(hitH.chanIdx, hitH.keyIdx, HandlePart::Key))
                Select(hitH.chanIdx, hitH.keyIdx, HandlePart::Key, additive);
            Select(hitH.chanIdx, hitH.keyIdx, hitH.part, true);
            // Start drag from handle
            m_isDragging = true;
            m_dragStartMouse = mouse;
            const BezierKey& bk = m_channels[hitH.chanIdx].keys[hitH.keyIdx];
            if (hitH.part == HandlePart::InTangent) {
                m_dragStartTime  = bk.inTangentDt;
                m_dragStartValue = bk.inTangentDv;
            } else {
                m_dragStartTime  = bk.outTangentDt;
                m_dragStartValue = bk.outTangentDv;
            }
        } else {
            auto [ci, ki] = hitTestKey(mouse);
            if (ci >= 0) {
                if (!IsSelected(ci, ki, HandlePart::Key))
                    Select(ci, ki, HandlePart::Key, io.KeyShift || io.KeyCtrl);
                // Take snapshot for drag
                m_isDragging = true;
                m_dragStartMouse = mouse;
                m_dragSnapshot.clear();
                for (const auto& s : m_selection) {
                    if (s.part != HandlePart::Key) continue;
                    if (s.chanIdx < 0 || s.chanIdx >= int(m_channels.size())) continue;
                    const CurveChannel& ch = m_channels[s.chanIdx];
                    if (s.keyIdx < 0 || s.keyIdx >= int(ch.keys.size())) continue;
                    const BezierKey& bk = ch.keys[s.keyIdx];
                    m_dragSnapshot.push_back({s.chanIdx, s.keyIdx,
                        bk.time, bk.value,
                        bk.inTangentDt, bk.inTangentDv,
                        bk.outTangentDt, bk.outTangentDv});
                }
            } else {
                // Start box select
                if (!io.KeyShift && !io.KeyCtrl) ClearSelection();
                m_isBoxSelecting = true;
                m_boxSelectStart = mouse;
                m_isDragging = false;
            }
        }
    }

    // ── Dragging keys / handles ───────────────────────────────────────────────
    if (m_isDragging && ImGui::IsMouseDown(0)) {
        float dx = mouse.x - m_dragStartMouse.x;
        float dy = mouse.y - m_dragStartMouse.y;
        float cw = cMax.x - cMin.x, ch2 = cMax.y - cMin.y;
        double dt = (cw > 1.f) ? double(dx) / double(cw) * (m_viewMaxTime-m_viewMinTime) : 0.0;
        double dv = (ch2 > 1.f)? -double(dy)/ double(ch2)* (m_viewMaxVal -m_viewMinVal)  : 0.0;

        // Check if we're dragging a tangent handle
        bool draggingHandle = false;
        for (const auto& s : m_selection) {
            if (s.part == HandlePart::InTangent || s.part == HandlePart::OutTangent) {
                draggingHandle = true;
                if (s.chanIdx < 0 || s.chanIdx >= int(m_channels.size())) continue;
                CurveChannel& ch = m_channels[s.chanIdx];
                if (s.keyIdx < 0 || s.keyIdx >= int(ch.keys.size())) continue;
                BezierKey& bk = ch.keys[s.keyIdx];

                if (s.part == HandlePart::InTangent) {
                    bk.inTangentDt = std::min(-0.01, m_dragStartTime  + dt);
                    bk.inTangentDv = m_dragStartValue + dv;
                    if (!bk.tangentsBroken) {
                        // Mirror to out
                        double ratio = (bk.outTangentDt > 1e-10)
                            ? bk.outTangentDt / std::max(1e-10, -bk.inTangentDt) : 1.0;
                        bk.outTangentDv = -bk.inTangentDv * ratio;
                    }
                } else {
                    bk.outTangentDt = std::max(0.01, m_dragStartTime  + dt);
                    bk.outTangentDv = m_dragStartValue + dv;
                    if (!bk.tangentsBroken) {
                        double ratio = (-bk.inTangentDt > 1e-10)
                            ? -bk.inTangentDt / std::max(1e-10, bk.outTangentDt) : 1.0;
                        bk.inTangentDv = -bk.outTangentDv * ratio;
                    }
                }
                ch.dirty = true;
            }
        }

        if (!draggingHandle) {
            // Move key positions
            for (const auto& snap : m_dragSnapshot) {
                if (snap.chanIdx < 0 || snap.chanIdx >= int(m_channels.size())) continue;
                CurveChannel& ch = m_channels[snap.chanIdx];
                if (snap.keyIdx < 0 || snap.keyIdx >= int(ch.keys.size())) continue;
                BezierKey& bk = ch.keys[snap.keyIdx];
                double newT = snap.time + dt;
                if (m_snapToFrame) newT = std::round(newT);
                bk.time  = newT;
                bk.value = snap.value + dv;
                bk.inTangentDt  = snap.inDt;
                bk.inTangentDv  = snap.inDv;
                bk.outTangentDt = snap.outDt;
                bk.outTangentDv = snap.outDv;
                ch.dirty = true;
            }
        }
    }

    if (ImGui::IsMouseReleased(0) && m_isDragging) {
        // Re-sort each affected channel by time
        for (int ci = 0; ci < int(m_channels.size()); ++ci) {
            bool affected = false;
            for (const auto& s : m_selection)
                if (s.chanIdx == ci) { affected = true; break; }
            if (!affected) continue;
            auto& ch = m_channels[ci];
            std::stable_sort(ch.keys.begin(), ch.keys.end(),
                [](const BezierKey& a, const BezierKey& b){ return a.time < b.time; });
        }
        ClearSelection();
        m_isDragging = false;
        m_dragSnapshot.clear();
    }

    // ── Box selection ─────────────────────────────────────────────────────────
    if (m_isBoxSelecting) {
        if (ImGui::IsMouseDown(0)) {
            // Draw the selection rect
            ImDrawList* dl = ImGui::GetWindowDrawList();
            dl->AddRectFilled(m_boxSelectStart, mouse, IM_COL32(100,160,255,40));
            dl->AddRect(m_boxSelectStart, mouse, IM_COL32(100,160,255,180));
        }
        if (ImGui::IsMouseReleased(0)) {
            float xMin = std::min(m_boxSelectStart.x, mouse.x);
            float xMax = std::max(m_boxSelectStart.x, mouse.x);
            float yMin = std::min(m_boxSelectStart.y, mouse.y);
            float yMax = std::max(m_boxSelectStart.y, mouse.y);
            for (int ci = 0; ci < int(m_channels.size()); ++ci) {
                auto& ch = m_channels[ci];
                if (!ch.visible) continue;
                for (int ki = 0; ki < int(ch.keys.size()); ++ki) {
                    ImVec2 p = KeyToScreen(ch.keys[ki], cMin, cMax);
                    if (p.x >= xMin && p.x <= xMax && p.y >= yMin && p.y <= yMax)
                        if (!IsSelected(ci, ki, HandlePart::Key))
                            m_selection.push_back({ci, ki, HandlePart::Key});
                }
            }
            m_isBoxSelecting = false;
        }
    }

    // ── Double-click: add key ─────────────────────────────────────────────────
    if (hovered && ImGui::IsMouseDoubleClicked(0) && !io.KeyAlt) {
        auto [ci, ki] = hitTestKey(mouse);
        if (ci < 0 && m_activeChannel >= 0 && m_activeChannel < int(m_channels.size())) {
            double t = XToTime(mouse.x, cMin, cMax);
            double v = YToVal (mouse.y, cMin, cMax);
            if (m_snapToFrame) t = std::round(t);
            AddKeyAt(m_activeChannel, t, v);
        }
    }

}

// ── Toolbar ───────────────────────────────────────────────────────────────────

void CurveEditorPanel::RenderToolbar()
{
    bool anyDirty = false;
    for (const auto& ch : m_channels) if (ch.dirty) { anyDirty = true; break; }

    if (!anyDirty) ImGui::BeginDisabled();
    if (ImGui::Button("Apply")) ApplyAllDirty();
    if (!anyDirty) ImGui::EndDisabled();

    ImGui::SameLine();

    if (!anyDirty) ImGui::BeginDisabled();
    if (ImGui::Button("Revert")) RevertAll();
    if (!anyDirty) ImGui::EndDisabled();

    ImGui::SameLine(); ImGui::SeparatorEx(ImGuiSeparatorFlags_Vertical); ImGui::SameLine();
    ImGui::Checkbox("Simplify", &m_simplifyOnBake);
    if (m_simplifyOnBake) {
        ImGui::SameLine();
        ImGui::SetNextItemWidth(60.f);
        ImGui::DragFloat("Tol", &m_simplifyTol, 0.001f, 0.0001f, 1.f, "%.4f");
    }

    ImGui::SameLine(); ImGui::SeparatorEx(ImGuiSeparatorFlags_Vertical); ImGui::SameLine();
    ImGui::Checkbox("Snap", &m_snapToFrame);

    ImGui::SameLine(); ImGui::SeparatorEx(ImGuiSeparatorFlags_Vertical); ImGui::SameLine();
    if (!m_displayTime.IsDefault()) {
        ImGui::TextDisabled("Frame: %.1f", m_displayTime.GetValue());
    }

    ImGui::SameLine(); ImGui::SeparatorEx(ImGuiSeparatorFlags_Vertical); ImGui::SameLine();
    if (ImGui::Button("Frame All")) FrameAll();

    if (anyDirty) {
        ImGui::SameLine();
        ImGui::TextColored({1.f, 0.7f, 0.2f, 1.f}, "(unsaved edits)");
    }
}

// ── Attr list ─────────────────────────────────────────────────────────────────

void CurveEditorPanel::RenderAttrList(float listWidth)
{
    ImGui::BeginChild("##attrList", {listWidth, 0.f}, true);
    ImGui::TextDisabled("Channels");
    ImGui::Separator();

    for (int ci = 0; ci < int(m_channels.size()); ++ci) {
        CurveChannel& ch = m_channels[ci];

        ImGui::PushID(ci);

        // Color swatch
        ImVec2 swatchPos = ImGui::GetCursorScreenPos();
        ImGui::GetWindowDrawList()->AddRectFilled(
            swatchPos, {swatchPos.x + 10.f, swatchPos.y + ImGui::GetTextLineHeight()},
            ImGui::ColorConvertFloat4ToU32(ch.color));
        ImGui::Dummy({12.f, ImGui::GetTextLineHeight()});
        ImGui::SameLine();

        // Visibility toggle
        bool vis = ch.visible;
        if (ImGui::Checkbox("##v", &vis)) ch.visible = vis;
        ImGui::SameLine();

        // Name (selectable for active channel)
        bool isActive = (m_activeChannel == ci);
        if (ch.dirty) {
            ImGui::PushStyleColor(ImGuiCol_Text, ImVec4{1.f, 0.7f, 0.2f, 1.f});
        }
        if (ImGui::Selectable(ch.displayName.c_str(), isActive,
                              ImGuiSelectableFlags_None, {0.f, 0.f})) {
            m_activeChannel = ci;
        }
        if (ch.dirty) ImGui::PopStyleColor();

        if (ImGui::IsItemHovered())
            ImGui::SetTooltip("%s", ch.attr.GetName().GetText());

        ImGui::PopID();
    }

    if (m_channels.empty() && !m_selectedPrimPath.empty())
        ImGui::TextDisabled("No animated attrs");
    else if (m_selectedPrimPath.empty())
        ImGui::TextDisabled("Select a prim");

    ImGui::EndChild();
}

// ── Top-level render ──────────────────────────────────────────────────────────

void CurveEditorPanel::Render()
{
    if (m_needsRefresh) RefreshFromStage();

    RenderToolbar();
    ImGui::Separator();

    float listWidth = 180.f;
    RenderAttrList(listWidth);
    ImGui::SameLine();

    // Canvas occupies the remaining space
    ImVec2 canvasMin = ImGui::GetCursorScreenPos();
    ImVec2 avail     = ImGui::GetContentRegionAvail();
    ImVec2 canvasMax = {canvasMin.x + avail.x, canvasMin.y + avail.y};

    // Reserve the canvas area so ImGui tracks hover/focus for this region
    ImGui::InvisibleButton("##canvas", avail);

    ImDrawList* dl = ImGui::GetWindowDrawList();
    dl->AddRectFilled(canvasMin, canvasMax, IM_COL32(30, 30, 30, 255));
    dl->PushClipRect(canvasMin, canvasMax, true);

    RenderGrid(dl, canvasMin, canvasMax);

    for (auto& ch : m_channels)
        RenderCurve(dl, ch, canvasMin, canvasMax);

    RenderTimeCursor(dl, canvasMin, canvasMax);

    dl->PopClipRect();

    HandleCanvasInput(canvasMin, canvasMax);
}

} // namespace UsdLayerManager
