#include "SceneHierarchyPanel.h"
#include "../utils/Logger.h"
#include "../utils/FileDialog.h"
#include "../core/commands/CreatePrimCommand.h"
#include "../core/commands/DeletePrimCommand.h"
#include "../core/commands/AddReferenceCommand.h"
#include "../core/commands/ReplaceReferenceCommand.h"
#include "../core/commands/RenamePrimCommand.h"
#include "../core/commands/ReparentPrimCommand.h"
#include "../core/commands/GroupPrimsCommand.h"
#include "../core/commands/LayerCommands.h"
#include <pxr/usd/usd/prim.h>
#include <pxr/usd/usd/primRange.h>
#include <pxr/usd/usd/references.h>
#include <pxr/usd/usd/payloads.h>
#include <pxr/usd/usd/inherits.h>
#include <pxr/usd/usd/specializes.h>
#include <pxr/usd/usd/variantSets.h>
#include <pxr/usd/usd/primCompositionQuery.h>
#include <pxr/usd/sdf/payload.h>
#include <pxr/usd/usdGeom/imageable.h>
#include <pxr/usd/usdGeom/metrics.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdGeom/xform.h>
#include <pxr/usd/usdGeom/xformCommonAPI.h>
#include <pxr/usd/sdf/layer.h>
#include <pxr/usd/sdf/reference.h>
#include <pxr/usd/sdf/primSpec.h>
#include <pxr/base/tf/token.h>
#include <filesystem>
#include <algorithm>
#include <cctype>
#include <unordered_set>
#include <memory>

PXR_NAMESPACE_USING_DIRECTIVE

namespace UsdLayerManager {

/// Convert a file base name (e.g. "my asset.v01") into a valid USD prim name.
/// USD identifiers: [A-Za-z_][A-Za-z0-9_]*
static std::string SanitizeUsdName(const std::string& raw) {
    std::string result;
    result.reserve(raw.size());
    for (char c : raw) {
        if (std::isalnum(static_cast<unsigned char>(c)) || c == '_') {
            result += c;
        } else {
            result += '_';
        }
    }
    if (result.empty() || std::isdigit(static_cast<unsigned char>(result[0]))) {
        result = "_" + result;
    }
    return result;
}

SceneHierarchyPanel::SceneHierarchyPanel()
    : m_propertyManager(nullptr)
    , m_stage(nullptr) {
}

SceneHierarchyPanel::~SceneHierarchyPanel() {
}

void SceneHierarchyPanel::SetPropertyManager(PropertyManager* manager) {
    m_propertyManager = manager;
}

void SceneHierarchyPanel::SetStage(UsdStageRefPtr stage) {
    m_stage = stage;
    m_selectedPaths.clear();
    m_primarySelectedPath.clear();
    m_primarySdfPath  = SdfPath();
    m_scrollToSelected = false;
}

UsdPrim SceneHierarchyPanel::GetSelectedPrim() const {
    if (m_stage && !m_primarySelectedPath.empty()) {
        return m_stage->GetPrimAtPath(SdfPath(m_primarySelectedPath));
    }
    return UsdPrim();
}

void SceneHierarchyPanel::SetSelectedPathFromClick(const std::string& path) {
    m_selectedPaths.clear();
    m_primarySelectedPath = path;
    m_primarySdfPath      = path.empty() ? SdfPath() : SdfPath(path);
    if (!path.empty()) m_selectedPaths.insert(path);
    if (!path.empty()) m_rangeAnchorPath = SdfPath(path);
    m_scrollToSelected = false;
    if (m_onPrimSelected)   m_onPrimSelected(path);
    if (m_onPrimsSelected) {
        std::vector<std::string> v;
        if (!path.empty()) v.push_back(path);
        m_onPrimsSelected(v);
    }
}

void SceneHierarchyPanel::ToggleSelectionFromClick(const std::string& path) {
    if (m_selectedPaths.count(path)) {
        m_selectedPaths.erase(path);
        if (m_primarySelectedPath == path) {
            m_primarySelectedPath = m_selectedPaths.empty() ? "" : *m_selectedPaths.begin();
            m_primarySdfPath = m_primarySelectedPath.empty()
                ? SdfPath() : SdfPath(m_primarySelectedPath);
        }
    } else {
        m_selectedPaths.insert(path);
        m_primarySelectedPath = path;
        m_primarySdfPath      = SdfPath(path);
        m_rangeAnchorPath     = SdfPath(path);
    }
    m_scrollToSelected = false;
    if (m_onPrimSelected) m_onPrimSelected(m_primarySelectedPath);
    if (m_onPrimsSelected) {
        std::vector<std::string> v(m_selectedPaths.begin(), m_selectedPaths.end());
        m_onPrimsSelected(v);
    }
}

void SceneHierarchyPanel::RangeSelectToPath(const SdfPath& path) {
    auto it1 = std::find(m_visiblePrimOrder.begin(), m_visiblePrimOrder.end(), m_rangeAnchorPath);
    auto it2 = std::find(m_visiblePrimOrder.begin(), m_visiblePrimOrder.end(), path);
    if (it1 != m_visiblePrimOrder.end() && it2 != m_visiblePrimOrder.end()) {
        int idx1 = static_cast<int>(it1 - m_visiblePrimOrder.begin());
        int idx2 = static_cast<int>(it2 - m_visiblePrimOrder.begin());
        if (idx1 > idx2) std::swap(idx1, idx2);
        for (int i = idx1; i <= idx2; ++i)
            m_selectedPaths.insert(m_visiblePrimOrder[i].GetString());
    } else {
        m_selectedPaths.insert(path.GetString());
    }
    m_primarySelectedPath = path.GetString();
    m_primarySdfPath      = path;
    m_scrollToSelected    = false;
    if (m_onPrimSelected) m_onPrimSelected(m_primarySelectedPath);
    if (m_onPrimsSelected) {
        std::vector<std::string> v(m_selectedPaths.begin(), m_selectedPaths.end());
        m_onPrimsSelected(v);
    }
}

const char* SceneHierarchyPanel::GetPrimTypeIcon(const UsdPrim& prim) const {
    // Kept for legacy callers; returns a short ASCII label.
    if (prim.IsPseudoRoot()) return "W";
    std::string t = prim.GetTypeName().GetString();
    if (t.find("Mesh")     != std::string::npos) return "G";
    if (t.find("Camera")   != std::string::npos) return "C";
    if (t.find("Light")    != std::string::npos) return "L";
    if (t.find("Material") != std::string::npos) return "S";
    if (t.find("Shader")   != std::string::npos) return "S";
    if (t.find("Xform")    != std::string::npos) return "X";
    if (t.find("Scope")    != std::string::npos) return "O";
    if (prim.IsModel()) return "M";
    return "P";
}

Icon SceneHierarchyPanel::GetPrimTypeIconEnum(const UsdPrim& prim) const {
    if (prim.IsPseudoRoot()) return Icon::Globe;
    std::string t = prim.GetTypeName().GetString();
    if (t.find("Mesh")     != std::string::npos ||
        t.find("Subdiv")   != std::string::npos) return Icon::Cube;
    if (t.find("Camera")   != std::string::npos) return Icon::Camera;
    if (t.find("Light")    != std::string::npos) return Icon::Lightbulb;
    if (t.find("Material") != std::string::npos) return Icon::Swatchbook;
    if (t.find("Shader")   != std::string::npos) return Icon::Code;
    if (t.find("Xform")    != std::string::npos) return Icon::ObjectGroup;
    if (t.find("Scope")    != std::string::npos) return Icon::FolderOpen;
    if (prim.IsModel()) return Icon::LayerGroup;
    return Icon::CircleDot;
}

SdfPath SceneHierarchyPanel::FindUniqueChildPath(const SdfPath& parent,
                                                  const std::string& baseName) {
    std::string name = baseName;
    int n = 1;
    while (m_stage->GetPrimAtPath(parent.AppendChild(TfToken(name))).IsValid())
        name = baseName + "_" + std::to_string(n++);
    return parent.AppendChild(TfToken(name));
}

void SceneHierarchyPanel::HandleKeyboardShortcuts() {
    if (!ImGui::IsWindowFocused(ImGuiFocusedFlags_RootAndChildWindows)) return;
    // Don't steal keys while any InputText (rename, arc modal, etc.) is active.
    if (ImGui::GetIO().WantTextInput) return;

    // F2 — start inline rename on the primary selection.
    if (ImGui::IsKeyPressed(ImGuiKey_F2, false) && !m_primarySdfPath.IsEmpty()) {
        UsdPrim prim = m_stage->GetPrimAtPath(m_primarySdfPath);
        if (prim.IsValid()) {
            m_renamingPath = m_primarySdfPath;
            std::string name = prim.GetName().GetString();
            strncpy(m_renameBuf, name.c_str(), sizeof(m_renameBuf) - 1);
            m_renameBuf[sizeof(m_renameBuf) - 1] = '\0';
            m_renameJustStarted = true;
            m_scrollToSelected  = true; // ensure the prim is visible
        }
    }

    // Ctrl+G — group selected prims under a new Xform.
    if (ImGui::GetIO().KeyCtrl &&
        ImGui::IsKeyPressed(ImGuiKey_G, false) &&
        !m_selectedPaths.empty() && m_commandHistory && m_stage) {

        std::vector<SdfPath> srcs;
        for (const auto& s : m_selectedPaths)
            srcs.emplace_back(s);

        // Filter out paths that are descendants of other selected paths
        // (they'll move with their ancestor, so grouping them separately is wrong).
        std::vector<SdfPath> filtered;
        for (const SdfPath& p : srcs) {
            bool hasAncestorInSet = false;
            for (const SdfPath& other : srcs) {
                if (other != p && p.HasPrefix(other)) {
                    hasAncestorInSet = true;
                    break;
                }
            }
            if (!hasAncestorInSet)
                filtered.push_back(p);
        }
        if (filtered.empty()) return;

        // Compute common parent of the filtered set.
        SdfPath commonParent = filtered[0].GetParentPath();
        for (const SdfPath& p : filtered) {
            while (!p.HasPrefix(commonParent))
                commonParent = commonParent.GetParentPath();
        }

        SdfPath groupPath = FindUniqueChildPath(commonParent, "group1");
        m_commandHistory->Push(std::make_unique<GroupPrimsCommand>(m_stage, filtered, groupPath));
        SetSelectedPathFromClick(groupPath.GetString());
    }

    // Ctrl+H — hide selected prims (visibility = invisible).
    if (ImGui::GetIO().KeyCtrl && ImGui::IsKeyPressed(ImGuiKey_H, false) &&
        !m_selectedPaths.empty() && m_stage) {
        for (const auto& pathStr : m_selectedPaths) {
            UsdPrim prim = m_stage->GetPrimAtPath(SdfPath(pathStr));
            if (prim.IsValid() && prim.IsA<UsdGeomImageable>())
                UsdGeomImageable(prim).GetVisibilityAttr().Set(UsdGeomTokens->invisible);
        }
    }

    // Shift+H — show selected prims (visibility = inherited).
    if (ImGui::GetIO().KeyShift && ImGui::IsKeyPressed(ImGuiKey_H, false) &&
        !m_selectedPaths.empty() && m_stage) {
        for (const auto& pathStr : m_selectedPaths) {
            UsdPrim prim = m_stage->GetPrimAtPath(SdfPath(pathStr));
            if (prim.IsValid() && prim.IsA<UsdGeomImageable>())
                UsdGeomImageable(prim).GetVisibilityAttr().Set(UsdGeomTokens->inherited);
        }
    }

    // Delete — open the remove-prim confirmation modal.
    if (ImGui::IsKeyPressed(ImGuiKey_Delete, false) && !m_primarySdfPath.IsEmpty()) {
        auto rootLayer = m_stage->GetRootLayer();
        if (rootLayer && !!rootLayer->GetPrimAtPath(m_primarySdfPath)) {
            m_pendingRemovePrimPath = m_primarySdfPath;
            m_showRemovePrimConfirm = true;
        }
    }
}

void SceneHierarchyPanel::Render() {
    if (!m_stage) {
        ImGui::TextDisabled("No stage loaded");
        return;
    }

    HandleKeyboardShortcuts();

    // ── Edit-target layer switcher ────────────────────────────────────
    {
        struct LayerEntry { SdfLayerHandle layer; std::string label; };
        std::vector<LayerEntry> layers;

        auto sessionLayer = m_stage->GetSessionLayer();
        if (sessionLayer)
            layers.push_back({ sessionLayer, "Session" });

        auto rootLayer = m_stage->GetRootLayer();
        if (rootLayer) {
            std::string name = rootLayer->IsAnonymous()
                ? "anonymous"
                : std::filesystem::path(rootLayer->GetIdentifier()).filename().string();
            layers.push_back({ rootLayer, "Root: " + name });

            for (const auto& subPath : rootLayer->GetSubLayerPaths()) {
                auto sub = SdfLayer::FindOrOpenRelativeToLayer(rootLayer, subPath);
                if (!sub) continue;
                std::string subName = sub->IsAnonymous()
                    ? sub->GetIdentifier()
                    : std::filesystem::path(sub->GetIdentifier()).filename().string();
                layers.push_back({ SdfLayerHandle(sub), subName });
            }
        }

        std::string currentLabel = "\xe2\x80\x94"; // em dash
        auto etLayer = m_stage->GetEditTarget().GetLayer();
        if (etLayer) {
            for (const auto& e : layers) {
                if (e.layer && e.layer->GetIdentifier() == etLayer->GetIdentifier()) {
                    currentLabel = e.label;
                    break;
                }
            }
        }

        bool renderLayerActive = m_renderLayerManager && !m_renderLayerManager->IsDefaultActive();
        ImGui::BeginDisabled(renderLayerActive);
        ImGui::SetNextItemWidth(-1.0f);
        if (ImGui::BeginCombo("##layerswitch", currentLabel.c_str(),
                              ImGuiComboFlags_HeightRegular)) {
            for (const auto& e : layers) {
                bool isCurrent = (e.layer && etLayer &&
                                  e.layer->GetIdentifier() == etLayer->GetIdentifier());
                if (ImGui::Selectable(e.label.c_str(), isCurrent)) {
                    if (m_layerManager)
                        m_layerManager->SetEditTarget(e.layer->GetIdentifier());
                    else
                        m_stage->SetEditTarget(UsdEditTarget(e.layer));
                }
                if (isCurrent)
                    ImGui::SetItemDefaultFocus();
            }
            ImGui::EndCombo();
        }
        bool comboHovered = ImGui::IsItemHovered();
        ImGui::EndDisabled();
        if (renderLayerActive && comboHovered)
            ImGui::SetTooltip("Edit target is locked to the active render layer");
        ImGui::Spacing();
    }

    RenderSublayerSection();
    ImGui::Spacing();

    auto paths = m_propertyManager->GetPrimPaths();
    if (paths.empty()) {
        ImGui::TextDisabled("No prims in stage");
    } else {
        UsdPrim root = m_stage->GetPseudoRoot();

        // Rebuilt each frame as RenderPrimNode visits prims — used for range select.
        m_visiblePrimOrder.clear();

        // Rebuild local-layer set once per frame (used by RenderPrimNode to
        // detect attribute overrides).  GetLayerStack() returns only the stage's
        // own layers — root layer, sublayers, session layer — NOT reference layers.
        m_localLayers.clear();
        for (const auto& layer : m_stage->GetLayerStack())
            m_localLayers.insert(layer->GetIdentifier());

        // ── ImGui Demo "Tables/Tree view" pattern ──────────────────────────
        //  Col 0 │ Col 1  │ Col 2 │ Col 3
        //  ▶ Prim│ Type   │ Vis   │ Ref
        //
        // The tree node lives in col 0 with ImGuiTreeNodeFlags_SpanAllColumns.
        // This makes the selection highlight, IsItemClicked, and SetScrollHereY
        // all operate on the FULL ROW rect — the correct ImGui tree-in-table model.
        const float kIconW = ImGui::GetTextLineHeight() + 4.0f;  // small fixed col width

        const ImGuiTableFlags tblFlags =
            ImGuiTableFlags_NoBordersInBody |
            ImGuiTableFlags_NoPadOuterX     |
            ImGuiTableFlags_RowBg           |
            ImGuiTableFlags_SizingFixedFit;

        const float kSpecW = ImGui::CalcTextSize("over").x + 2.f;
        if (ImGui::BeginTable("##primtree", 5, tblFlags)) {
            // Col 0 stretches; cols 1-4 are small fixed-width columns.
            ImGui::TableSetupColumn("##prim", ImGuiTableColumnFlags_WidthStretch);
            ImGui::TableSetupColumn("##type", ImGuiTableColumnFlags_WidthFixed, kIconW);
            ImGui::TableSetupColumn("##vis",  ImGuiTableColumnFlags_WidthFixed, kIconW);
            ImGui::TableSetupColumn("##ref",  ImGuiTableColumnFlags_WidthFixed, kIconW);
            ImGui::TableSetupColumn("##spec", ImGuiTableColumnFlags_WidthFixed, kSpecW);

            for (const auto& child : root.GetChildren())
                RenderPrimNode(child);

            ImGui::EndTable();
        }

        // ── Stage-root drop zone ───────────────────────────────────────────
        // Invisible strip below the tree: drop here to reparent a prim to root.
        ImGui::InvisibleButton("##rootdrop", ImVec2(-1.0f, 8.0f));
        if (ImGui::BeginDragDropTarget()) {
            if (const ImGuiPayload* payload = ImGui::AcceptDragDropPayload("USD_PRIM_PATH")) {
                SdfPath srcPath(static_cast<const char*>(payload->Data));
                // Only reparent if not already a root prim.
                if (m_commandHistory && srcPath.GetPathElementCount() > 1) {
                    std::string name = FindUniqueChildPath(SdfPath("/"), srcPath.GetName()).GetName();
                    m_commandHistory->Push(std::make_unique<ReparentPrimCommand>(
                        m_stage, srcPath, SdfPath("/"), name));
                    SetSelectedPathFromClick(SdfPath("/").AppendChild(TfToken(name)).GetString());
                }
            }
            ImGui::EndDragDropTarget();
        }
        if (ImGui::IsItemHovered())
            ImGui::SetTooltip("Drop here to reparent to stage root");

        // Deselect when left-clicking on blank space (no prim item hovered).
        if (ImGui::IsWindowHovered(ImGuiHoveredFlags_RootAndChildWindows) &&
            ImGui::IsMouseClicked(ImGuiMouseButton_Left) &&
            !ImGui::IsAnyItemHovered())
        {
            SetSelectedPathFromClick("");
        }
    }

    // Window-level right-click context menu (blank area) for stage-level operations.
    if (ImGui::BeginPopupContextWindow("StageContextMenu",
            ImGuiPopupFlags_MouseButtonRight | ImGuiPopupFlags_NoOpenOverItems)) {
        ImGui::TextDisabled("Stage");
        ImGui::Separator();

        // ---- Create Prim ----
        static const char* kPrimTypes[] = {
            "Xform", "Scope",
            "Mesh", "Sphere", "Cube", "Cylinder", "Cone", "Capsule",
            "Camera",
            "SphereLight", "DomeLight", "RectLight", "DiskLight",
            "CylinderLight", "DistantLight"
        };
        if (ImGui::BeginMenu("Create Prim")) {
            for (const char* typeName : kPrimTypes) {
                if (ImGui::MenuItem(typeName)) {
                    std::string baseName = typeName;
                    std::string finalName = baseName;
                    int suffix = 1;
                    while (m_stage->GetPrimAtPath(SdfPath("/" + finalName)).IsValid()) {
                        finalName = baseName + "_" + std::to_string(suffix++);
                    }
                    SdfPath primPath("/" + finalName);
                    if (m_commandHistory) {
                        auto cmd = std::make_unique<CreatePrimCommand>(
                            m_stage, primPath, TfToken(typeName));
                        m_commandHistory->Push(std::move(cmd));
                        UsdPrim newPrim = m_stage->GetPrimAtPath(primPath);
                        if (newPrim.IsValid()) SetSelectedPathFromClick(primPath.GetString());
                    } else {
                        try {
                            UsdPrim newPrim = m_stage->DefinePrim(primPath, TfToken(typeName));
                            if (newPrim.IsValid()) {
                                // Set correct camera orientation for the current up-axis
                                if (std::string(typeName) == "Camera") {
                                    bool isZUp = (UsdGeomGetStageUpAxis(m_stage) == UsdGeomTokens->z);
                                    if (isZUp) {
                                        pxr::UsdGeomXformCommonAPI xformAPI(newPrim);
                                        xformAPI.SetRotate(
                                            pxr::GfVec3f(90.f, 0.f, 0.f),
                                            pxr::UsdGeomXformCommonAPI::RotationOrderXYZ,
                                            pxr::UsdTimeCode::Default());
                                    }
                                }
                                LOG_INFO("Created prim '" + primPath.GetString() + "' of type " + baseName);
                                SetSelectedPathFromClick(primPath.GetString());
                            } else {
                                LOG_ERROR("Failed to create prim of type: " + baseName);
                            }
                        } catch (const std::exception& e) {
                            LOG_ERROR(std::string("Create prim error: ") + e.what());
                        }
                    }
                }
            }
            ImGui::EndMenu();
        }

        ImGui::Separator();

        // ---- Set Up Axis ----
        {
            TfToken currentUpAxis = UsdGeomGetStageUpAxis(m_stage);
            bool isYUp = (currentUpAxis == UsdGeomTokens->y);
            bool isZUp = (currentUpAxis == UsdGeomTokens->z);
            if (ImGui::BeginMenu("Set Up Axis")) {
                if (ImGui::MenuItem("Y Up", nullptr, isYUp, !isYUp)) {
                    if (UsdGeomSetStageUpAxis(m_stage, UsdGeomTokens->y)) {
                        LOG_INFO("Stage up axis set to Y");
                        if (m_onStageMetadataChanged) m_onStageMetadataChanged();
                    } else {
                        LOG_ERROR("Failed to set stage up axis to Y");
                    }
                }
                if (ImGui::MenuItem("Z Up", nullptr, isZUp, !isZUp)) {
                    if (UsdGeomSetStageUpAxis(m_stage, UsdGeomTokens->z)) {
                        LOG_INFO("Stage up axis set to Z");
                        if (m_onStageMetadataChanged) m_onStageMetadataChanged();
                    } else {
                        LOG_ERROR("Failed to set stage up axis to Z");
                    }
                }
                ImGui::EndMenu();
            }
        }

        ImGui::Separator();

        // ---- Reference submenu ----
        if (ImGui::BeginMenu("Reference")) {
            if (ImGui::MenuItem("Add...")) {
                std::string filePath = FileDialog::OpenFile(
                    "USD Files (*.usd;*.usda;*.usdc;*.usdz)\0*.usd;*.usda;*.usdc;*.usdz\0All Files (*.*)\0*.*\0",
                    "Add Reference File");
                if (!filePath.empty()) {
                    std::string stem = std::filesystem::path(filePath).stem().string();
                    std::string xformName = SanitizeUsdName(stem);
                    if (xformName.empty()) xformName = "Reference";

                    std::string finalName = xformName;
                    int suffix = 1;
                    while (m_stage->GetPrimAtPath(SdfPath("/" + finalName)).IsValid())
                        finalName = xformName + "_" + std::to_string(suffix++);

                    SdfPath xformPath("/" + finalName);
                    if (m_commandHistory) {
                        m_commandHistory->Push(std::make_unique<AddReferenceCommand>(
                            m_stage, xformPath, filePath));
                    } else {
                        try {
                            UsdPrim xformPrim = m_stage->DefinePrim(xformPath, TfToken("Xform"));
                            if (xformPrim.IsValid()) {
                                bool ok = xformPrim.GetReferences().AddReference(filePath);
                                if (ok)
                                    LOG_INFO("Added reference '" + filePath + "' under prim: " + xformPath.GetString());
                                else
                                    LOG_ERROR("Failed to add reference '" + filePath + "' to: " + xformPath.GetString());
                            } else {
                                LOG_ERROR("Failed to define Xform prim: " + xformPath.GetString());
                            }
                        } catch (const std::exception& e) {
                            LOG_ERROR(std::string("Add reference error: ") + e.what());
                        }
                    }
                }
            }
            ImGui::EndMenu();
        }

        // ---- Payload submenu ----
        if (ImGui::BeginMenu("Payload")) {
            if (ImGui::MenuItem("Add...")) {
                std::string filePath = FileDialog::OpenFile(
                    "USD Files (*.usd;*.usda;*.usdc;*.usdz)\0*.usd;*.usda;*.usdc;*.usdz\0All Files (*.*)\0*.*\0",
                    "Add Payload File");
                if (!filePath.empty()) {
                    std::string stem = std::filesystem::path(filePath).stem().string();
                    std::string xformName = SanitizeUsdName(stem);
                    if (xformName.empty()) xformName = "Payload";

                    std::string finalName = xformName;
                    int suffix = 1;
                    while (m_stage->GetPrimAtPath(SdfPath("/" + finalName)).IsValid())
                        finalName = xformName + "_" + std::to_string(suffix++);

                    SdfPath xformPath("/" + finalName);
                    try {
                        UsdPrim xformPrim = m_stage->DefinePrim(xformPath, TfToken("Xform"));
                        if (xformPrim.IsValid()) {
                            bool ok = xformPrim.GetPayloads().AddPayload(filePath);
                            if (ok)
                                LOG_INFO("Added payload '" + filePath + "' under prim: " + xformPath.GetString());
                            else
                                LOG_ERROR("Failed to add payload '" + filePath + "' to: " + xformPath.GetString());
                        } else {
                            LOG_ERROR("Failed to define Xform prim: " + xformPath.GetString());
                        }
                    } catch (const std::exception& e) {
                        LOG_ERROR(std::string("Add payload error: ") + e.what());
                    }
                }
            }
            ImGui::EndMenu();
        }

        ImGui::EndPopup();
    }

    // Deferred confirm modal for prim removal (must be opened outside any popup stack).
    RenderRemovePrimModal();

    // Deferred modal for arc operations (Inherit / Specialize / VariantSet / Variant).
    RenderArcModal();

    // Deferred file-dialog for reference replacement (must run outside any popup stack).
    ProcessPendingReplaceRef();
}

// ---------------------------------------------------------------------------
void SceneHierarchyPanel::RenderSublayerSection() {
    if (!ImGui::CollapsingHeader("Sublayers", ImGuiTreeNodeFlags_DefaultOpen))
        return;

    if (!m_layerManager) {
        ImGui::TextDisabled("  No layer manager");
        return;
    }

    auto sublayers = m_layerManager->GetSublayers();
    if (sublayers.empty()) {
        ImGui::TextDisabled("  No sublayers");
        return;
    }

    const float rowHeight   = ImGui::GetTextLineHeightWithSpacing();
    const int   visibleRows = std::min<int>(static_cast<int>(sublayers.size()), 6);
    float tableHeight = rowHeight * (visibleRows + 1); // +1 for header row

    if (!ImGui::BeginTable("HierarchySublayerTable", 6,
            ImGuiTableFlags_RowBg | ImGuiTableFlags_BordersInnerV |
            ImGuiTableFlags_SizingFixedFit | ImGuiTableFlags_ScrollY,
            ImVec2(0, tableHeight)))
        return;

    ImGui::TableSetupScrollFreeze(0, 1);
    ImGui::TableSetupColumn("##drag",  ImGuiTableColumnFlags_WidthFixed, 14.0f);
    ImGui::TableSetupColumn("##dirty", ImGuiTableColumnFlags_WidthFixed, 14.0f);
    ImGui::TableSetupColumn("Name",    ImGuiTableColumnFlags_WidthStretch);
    ImGui::TableSetupColumn("##et",    ImGuiTableColumnFlags_WidthFixed, 28.0f);
    ImGui::TableSetupColumn("Muted",   ImGuiTableColumnFlags_WidthFixed, 55.0f);
    ImGui::TableSetupColumn("##order", ImGuiTableColumnFlags_WidthFixed, 44.0f);

    // Custom header row: pen icon for the ET column, text for the rest.
    ImGui::TableNextRow(ImGuiTableRowFlags_Headers);
    ImGui::TableSetColumnIndex(2); ImGui::TableHeader("Name");
    ImGui::TableSetColumnIndex(3);
    {
        float cellW = ImGui::GetContentRegionAvail().x;
        float iconW = 14.0f;
        float off   = (cellW - iconW) * 0.5f;
        if (off > 0.0f) ImGui::SetCursorPosX(ImGui::GetCursorPosX() + off);
        if (m_iconManager)
            ImGui::Image(ImTextureRef(m_iconManager->Get(Icon::Pen)), ImVec2(iconW, iconW));
        else
            ImGui::TextUnformatted("ET");
        if (ImGui::IsItemHovered()) ImGui::SetTooltip("Edit Target");
    }
    ImGui::TableSetColumnIndex(4); ImGui::TableHeader("Muted");
    ImGui::TableSetColumnIndex(5); ImGui::TableHeader("Order");

    for (int i = 0; i < static_cast<int>(sublayers.size()); i++) {
        const auto& li = sublayers[i];
        ImGui::TableNextRow();
        ImGui::PushID(i);

        // Col 0: invisible span-all Selectable for row selection / drag / context menu.
        ImGui::TableSetColumnIndex(0);
        ImVec2 rowStart = ImGui::GetCursorScreenPos(); // save before Selectable moves cursor

        bool isSelected = (m_sublayerSelectedIdx == i);
        if (ImGui::Selectable("##row", isSelected,
                ImGuiSelectableFlags_SpanAllColumns |
                ImGuiSelectableFlags_AllowOverlap,
                ImVec2(0, 0)))
            m_sublayerSelectedIdx = i;

        // Context menu: must come immediately after the span-all Selectable so
        // BeginPopupContextItem sees it as the "last item" and responds to
        // right-click anywhere in the row.
        RenderSublayerContextMenu(i, sublayers);

        // Drag source: also attached to the Selectable (left-button drag).
        if (ImGui::BeginDragDropSource(ImGuiDragDropFlags_SourceAllowNullID)) {
            ImGui::SetDragDropPayload("HIER_SUBLAYER_IDX", &i, sizeof(int));
            ImGui::Text("%s", li.displayName.c_str());
            ImGui::EndDragDropSource();
        }

        // Drop target on the Selectable.
        if (ImGui::BeginDragDropTarget()) {
            if (const ImGuiPayload* payload =
                    ImGui::AcceptDragDropPayload("HIER_SUBLAYER_IDX")) {
                int src = *static_cast<const int*>(payload->Data);
                int dst = i;
                if (src != dst) {
                    std::vector<std::string> before, after;
                    for (const auto& sl : sublayers) before.push_back(sl.identifier);
                    after = before;
                    std::string moved = after[src];
                    after.erase(after.begin() + src);
                    after.insert(after.begin() + dst, moved);
                    if (m_commandHistory) {
                        m_commandHistory->Push(std::make_unique<LayerReorderCommand>(
                            m_layerManager, before, after));
                    } else {
                        LayerReorderCommand cmd(m_layerManager, before, after);
                        static_cast<ICommand&>(cmd).Execute();
                    }
                }
            }
            ImGui::EndDragDropTarget();
        }

        // "=" drag handle: overlay at col 0 using the saved screen position so it
        // sits inside the row, not below it (TableSetColumnIndex only resets X).
        ImGui::SetCursorScreenPos(rowStart);
        ImGui::TextDisabled("=");

        // Col 1: dirty indicator
        ImGui::TableSetColumnIndex(1);
        if (li.isDirty)
            ImGui::TextColored(ImVec4(1.0f, 0.6f, 0.2f, 1.0f), "*");
        else
            ImGui::TextDisabled(" ");

        bool isRenderLayer = RenderLayerManager::IsRenderLayerSublayer(li.layer);
        bool renderLayerActive = m_renderLayerManager && !m_renderLayerManager->IsDefaultActive();

        // Col 2: layer name
        ImGui::TableSetColumnIndex(2);
        ImVec4 nameColor = li.isMuted
            ? ImVec4(0.5f, 0.5f, 0.5f, 1.0f)
            : ImVec4(1.0f, 1.0f, 1.0f, 1.0f);
        ImGui::TextColored(nameColor, "%s", li.displayName.c_str());
        if (ImGui::IsItemHovered() && !li.realPath.empty())
            ImGui::SetTooltip("%s\n%s", li.identifier.c_str(), li.realPath.c_str());
        if (isRenderLayer) {
            ImGui::SameLine();
            ImGui::TextColored(ImVec4(0.9f, 0.7f, 0.2f, 1.0f), "[RL]");
            if (ImGui::IsItemHovered())
                ImGui::SetTooltip("Managed from the Render Layer panel");
        }

        // Col 3: edit target checkbox — checking sets this layer as edit target;
        // unchecking does nothing (there is always an active edit target).
        ImGui::TableSetColumnIndex(3);
        {
            bool isET = li.isEditTarget;
            ImGui::PushStyleVar(ImGuiStyleVar_FramePadding, ImVec2(2.0f, 1.0f));
            ImGui::BeginDisabled(renderLayerActive);
            if (ImGui::Checkbox("##et", &isET) && isET)
                m_layerManager->SetEditTarget(li.identifier);
            ImGui::EndDisabled();
            ImGui::PopStyleVar();
            if (ImGui::IsItemHovered())
                ImGui::SetTooltip(renderLayerActive ? "Edit target is locked to the active render layer"
                                  : li.isEditTarget  ? "Edit target (active)"
                                                      : "Set as edit target");
        }

        // Col 4: mute checkbox — disabled for render-layer sublayers, which
        // RenderLayerManager mutes/unmutes exclusively as layers are switched.
        ImGui::TableSetColumnIndex(4);
        {
            bool muted = li.isMuted;
            ImGui::PushStyleVar(ImGuiStyleVar_FramePadding, ImVec2(2.0f, 1.0f));
            ImGui::BeginDisabled(isRenderLayer);
            if (ImGui::Checkbox("##muted", &muted)) {
                if (muted) m_layerManager->MuteLayer(li.identifier);
                else       m_layerManager->UnmuteLayer(li.identifier);
            }
            ImGui::EndDisabled();
            ImGui::PopStyleVar();
            if (isRenderLayer && ImGui::IsItemHovered())
                ImGui::SetTooltip("Managed from the Render Layer panel");
        }

        // Col 5: up/down reorder icon buttons.
        ImGui::TableSetColumnIndex(5);
        {
            bool canUp   = (i > 0);
            bool canDown = (i < static_cast<int>(sublayers.size()) - 1);

            auto buildBeforeAfter = [&](int a, int b,
                    std::vector<std::string>& before,
                    std::vector<std::string>& after) {
                for (const auto& sl : sublayers) before.push_back(sl.identifier);
                after = before;
                std::swap(after[a], after[b]);
            };

            ImVec2 iconSize(14.0f, 14.0f);

            ImGui::BeginDisabled(!canUp);
            ImTextureID upTex = m_iconManager ? m_iconManager->Get(Icon::ChevronUp) : ImTextureID_Invalid;
            if (m_iconManager ? ImGui::ImageButton("##up", ImTextureRef(upTex), iconSize)
                              : ImGui::ArrowButton("##up", ImGuiDir_Up)) {
                std::vector<std::string> before, after;
                buildBeforeAfter(i, i - 1, before, after);
                if (m_commandHistory)
                    m_commandHistory->Push(
                        std::make_unique<LayerReorderCommand>(m_layerManager, before, after));
                else {
                    LayerReorderCommand cmd(m_layerManager, before, after);
                    static_cast<ICommand&>(cmd).Execute();
                }
            }
            ImGui::EndDisabled();
            if (ImGui::IsItemHovered()) ImGui::SetTooltip("Move Up");

            ImGui::SameLine(0.0f, 2.0f);

            ImGui::BeginDisabled(!canDown);
            ImTextureID downTex = m_iconManager ? m_iconManager->Get(Icon::ChevronDown) : ImTextureID_Invalid;
            if (m_iconManager ? ImGui::ImageButton("##down", ImTextureRef(downTex), iconSize)
                              : ImGui::ArrowButton("##down", ImGuiDir_Down)) {
                std::vector<std::string> before, after;
                buildBeforeAfter(i, i + 1, before, after);
                if (m_commandHistory)
                    m_commandHistory->Push(
                        std::make_unique<LayerReorderCommand>(m_layerManager, before, after));
                else {
                    LayerReorderCommand cmd(m_layerManager, before, after);
                    static_cast<ICommand&>(cmd).Execute();
                }
            }
            ImGui::EndDisabled();
            if (ImGui::IsItemHovered()) ImGui::SetTooltip("Move Down");
        }

        ImGui::PopID();
    }

    ImGui::EndTable();
}

// ---------------------------------------------------------------------------
void SceneHierarchyPanel::RenderSublayerContextMenu(int i,
        const std::vector<LayerInfo>& sublayers) {
    if (!ImGui::BeginPopupContextItem(("ctx_sublayer_" + std::to_string(i)).c_str()))
        return;

    const auto& li = sublayers[i];
    bool isRenderLayer = RenderLayerManager::IsRenderLayerSublayer(li.layer);
    bool renderLayerActive = m_renderLayerManager && !m_renderLayerManager->IsDefaultActive();

    if (ImGui::MenuItem(li.isMuted ? "Unmute" : "Mute", nullptr, false, !isRenderLayer)) {
        if (li.isMuted) m_layerManager->UnmuteLayer(li.identifier);
        else            m_layerManager->MuteLayer(li.identifier);
    }
    if (isRenderLayer && ImGui::IsItemHovered())
        ImGui::SetTooltip("Managed from the Render Layer panel");

    if (!li.isEditTarget &&
        ImGui::MenuItem("Set as Edit Target", nullptr, false, !renderLayerActive))
        m_layerManager->SetEditTarget(li.identifier);
    if (renderLayerActive && !li.isEditTarget && ImGui::IsItemHovered())
        ImGui::SetTooltip("Edit target is locked to the active render layer");

    ImGui::Separator();

    if (ImGui::MenuItem("Remove", nullptr, false, !isRenderLayer)) {
        if (m_commandHistory)
            m_commandHistory->Push(
                std::make_unique<LayerRemoveCommand>(m_layerManager, i));
        else
            m_layerManager->RemoveSublayer(i);
    }
    if (isRenderLayer && ImGui::IsItemHovered())
        ImGui::SetTooltip("Managed from the Render Layer panel");

    ImGui::EndPopup();
}

void SceneHierarchyPanel::RenderPrimNode(const UsdPrim& prim) {
    if (!prim.IsValid()) return;

    std::string displayName = prim.GetName().GetString();
    if (displayName.empty()) displayName = prim.GetPath().GetString();

    std::string typeName = prim.GetTypeName().GetString();
    SdfPath     primPath = prim.GetPath();
    std::string primStr  = primPath.GetString();

    bool isSelected  = (m_selectedPaths.count(primStr) > 0);
    bool isActive    = prim.IsActive();
    bool isImageable = prim.IsA<UsdGeomImageable>();
    bool isInvisible = false;
    bool hasRefs     = prim.HasAuthoredReferences();
    bool hasChildren = !prim.GetChildren().empty();
    SdfSpecifier spec = prim.GetSpecifier();

    if (isImageable) {
        UsdGeomImageable img(prim);
        isInvisible = (img.ComputeVisibility() == UsdGeomTokens->invisible);
    }

    // ── Colour coding ───────────────────────────────────────────────────────
    // Orange: prim (or any of its attributes) has an opinion in the stage's
    //         own layers → indicates a local override on top of references.
    // Blue:   prim has references but NO local attribute override.
    // Both colours are dimmed when the prim is inactive.
    // ────────────────────────────────────────────────────────────────────────
    bool hasOverride = false;
    if (!m_localLayers.empty()) {
        for (const auto& attr : prim.GetAuthoredAttributes()) {
            for (const auto& spec : attr.GetPropertyStack()) {
                if (m_localLayers.count(spec->GetLayer()->GetIdentifier())) {
                    hasOverride = true;
                    break;
                }
            }
            if (hasOverride) break;
        }
    }

    // Force-open ancestor nodes when scrolling to the primary selection.
    bool isAncestorOfPrimary = m_scrollToSelected &&
                               !m_primarySdfPath.IsEmpty() &&
                               !m_primarySdfPath.IsRootPrimPath() &&
                               m_primarySdfPath.HasPrefix(primPath) &&
                               (m_primarySdfPath != primPath);
    if (isAncestorOfPrimary)
        ImGui::SetNextItemOpen(true, ImGuiCond_Always);

    // ── ImGui Demo tree-in-table pattern ────────────────────────────────────
    // Tree node goes in Col 0 with SpanAllColumns.  This makes the full row
    // rect the "item" for selection highlight, IsItemClicked, and scroll.
    // Subsequent columns are filled AFTER the tree node open/close decision.
    // ────────────────────────────────────────────────────────────────────────

    ImGui::TableNextRow();
    m_visiblePrimOrder.push_back(primPath);  // for Shift+LMB range select
    ImGui::TableNextColumn();   // Col 0 — prim name + tree arrow

    ImGui::PushID(primStr.c_str());

    // Determine final text colour for the prim name.
    // Priority: override (orange) > reference (blue) > inactive (dim) > default.
    // Alpha is reduced when the prim is inactive.
    const float alpha = isActive ? 1.0f : 0.45f;
    bool pushedColor = false;
    if (hasOverride) {
        // Orange  — local attribute override present
        ImGui::PushStyleColor(ImGuiCol_Text, ImVec4(1.0f, 0.60f, 0.10f, alpha));
        pushedColor = true;
    } else if (hasRefs) {
        // Blue  — has references, no local overrides
        ImGui::PushStyleColor(ImGuiCol_Text, ImVec4(0.40f, 0.70f, 1.0f, alpha));
        pushedColor = true;
    } else if (!isActive) {
        // Dim grey for inactive prims with no other colour
        ImGui::PushStyleColor(ImGuiCol_Text, ImVec4(0.45f, 0.45f, 0.45f, 1.0f));
        pushedColor = true;
    }

    bool isRenaming = (m_renamingPath == primPath);

    ImGuiTreeNodeFlags flags = ImGuiTreeNodeFlags_OpenOnArrow |
                               ImGuiTreeNodeFlags_SpanAllColumns;    // full-row item rect
    if (!isRenaming)   flags |= ImGuiTreeNodeFlags_OpenOnDoubleClick;
    if (isSelected)    flags |= ImGuiTreeNodeFlags_Selected;
    if (!hasChildren)  flags |= ImGuiTreeNodeFlags_Leaf |
                                ImGuiTreeNodeFlags_NoTreePushOnOpen;

    bool open    = false;
    bool renamed = false;   // true if rename committed this frame (prim now expired)
    if (isRenaming) {
        // Render just the arrow (hidden label) so we keep indentation and open/close,
        // then overlay an InputText inline for the prim name.
        open = ImGui::TreeNodeEx("##renaming", flags);
        if (pushedColor)
            ImGui::PopStyleColor();
        ImGui::SameLine();
        ImGui::SetNextItemWidth(-1.0f);
        // firstFrame: skip deactivation-cancel on the exact frame SetKeyboardFocusHere
        // activates the widget — IsItemDeactivated() can fire spuriously on that frame.
        bool firstFrame = m_renameJustStarted;
        if (m_renameJustStarted) {
            ImGui::SetKeyboardFocusHere();
            m_renameJustStarted = false;
        }
        bool commit = ImGui::InputText("##rename_input", m_renameBuf, sizeof(m_renameBuf),
                                       ImGuiInputTextFlags_EnterReturnsTrue |
                                       ImGuiInputTextFlags_AutoSelectAll);
        bool canceled = !firstFrame && ImGui::IsItemDeactivated() && !commit;
        if (commit && m_renameBuf[0] != '\0') {
            std::string newName = SanitizeUsdName(m_renameBuf);
            if (!newName.empty() && newName != displayName && m_commandHistory) {
                m_commandHistory->Push(std::make_unique<RenamePrimCommand>(
                    m_stage, primPath, newName));
                SdfPath newPath = primPath.GetParentPath().AppendChild(TfToken(newName));
                SetSelectedPathFromClick(newPath.GetString());
                renamed = true;  // prim at primPath is now expired — stop here
            }
            m_renamingPath = SdfPath();
        }
        if (canceled)
            m_renamingPath = SdfPath();
    } else {
        open = ImGui::TreeNodeEx(displayName.c_str(), flags);
        if (pushedColor)
            ImGui::PopStyleColor();
    }

    // ── Early-out after rename: prim is expired, close table row and return ──
    if (renamed) {
        ImGui::TableNextColumn();   // Col 1
        ImGui::TableNextColumn();   // Col 2
        ImGui::TableNextColumn();   // Col 3
        ImGui::TableNextColumn();   // Col 4
        if (open && hasChildren)
            ImGui::TreePop();
        ImGui::PopID();
        return;
    }

    // ── Scroll-to-selection (SpanAllColumns gives correct full-row rect) ────
    if (m_scrollToSelected && primStr == m_primarySelectedPath) {
        ImGui::SetScrollHereY(0.5f);
        m_scrollToSelected = false;
    }

    if (!isRenaming) {
        // Selection on LMB click (not on toggle arrow).
        // Shift+LMB: range-select from anchor. Ctrl+LMB: toggle. Plain: replace.
        if (ImGui::IsItemClicked() && !ImGui::IsItemToggledOpen()) {
            if (ImGui::GetIO().KeyShift)
                RangeSelectToPath(primPath);
            else if (ImGui::GetIO().KeyCtrl)
                ToggleSelectionFromClick(primStr);
            else
                SetSelectedPathFromClick(primStr);
        }

        // Double-click to start inline rename.
        if (ImGui::IsItemHovered() && ImGui::IsMouseDoubleClicked(ImGuiMouseButton_Left)) {
            m_renamingPath = primPath;
            strncpy(m_renameBuf, displayName.c_str(), sizeof(m_renameBuf) - 1);
            m_renameBuf[sizeof(m_renameBuf) - 1] = '\0';
            m_renameJustStarted = true;
        }

        // Tooltip.
        if (ImGui::IsItemHovered()) {
            std::string tip = "Type: " + (typeName.empty() ? "(unknown)" : typeName) +
                              "\nPath: " + primStr +
                              "\nActive: " + (isActive ? "Yes" : "No");
            if (isImageable)
                tip += std::string("\nVisibility: ") + (isInvisible ? "Invisible" : "Visible");
            if (hasRefs)
                tip += "\nHas references";
            if (m_selectedPaths.size() > 1)
                tip += "\n\n" + std::to_string(m_selectedPaths.size()) + " prims selected";
            ImGui::SetTooltip("%s", tip.c_str());
        }

        // ── Drag-and-drop source ─────────────────────────────────────────
        if (ImGui::BeginDragDropSource()) {
            ImGui::SetDragDropPayload("USD_PRIM_PATH",
                                      primStr.c_str(), primStr.size() + 1,
                                      ImGuiCond_Once);
            ImGui::Text("Move: %s", displayName.c_str());
            ImGui::EndDragDropSource();
        }

        // ── Drag-and-drop target ─────────────────────────────────────────
        if (ImGui::BeginDragDropTarget()) {
            if (const ImGuiPayload* payload = ImGui::AcceptDragDropPayload("USD_PRIM_PATH")) {
                SdfPath src(static_cast<const char*>(payload->Data));
                // Guard: don't reparent onto self or onto a descendant.
                if (src != primPath && !primPath.HasPrefix(src) && m_commandHistory) {
                    std::string name = FindUniqueChildPath(primPath, src.GetName()).GetName();
                    m_commandHistory->Push(std::make_unique<ReparentPrimCommand>(
                        m_stage, src, primPath, name));
                    SetSelectedPathFromClick(primPath.AppendChild(TfToken(name)).GetString());
                }
            }
            ImGui::EndDragDropTarget();
        }

        // Context menu on RMB.
        RenderContextMenu(prim);

        // A context menu op (Unparent, Group, etc.) may have moved this prim,
        // expiring it. Bail out before touching the dead prim object.
        if (!prim.IsValid()) {
            ImGui::TableNextColumn();   // Col 1
            ImGui::TableNextColumn();   // Col 2
            ImGui::TableNextColumn();   // Col 3
            ImGui::TableNextColumn();   // Col 4
            if (open && hasChildren)
                ImGui::TreePop();
            ImGui::PopID();
            return;
        }
    }

    // ── Col 1: Prim-type icon ───────────────────────────────────────────────
    ImGui::TableNextColumn();
    {
        const float     iconSz  = ImGui::GetTextLineHeight();
        const ImVec2    iconVec(iconSz, iconSz);
        ImTextureID     id      = m_iconManager ? m_iconManager->Get(GetPrimTypeIconEnum(prim))
                                                : ImTextureID_Invalid;
        ImVec4 tint = isActive ? ImVec4(1,1,1,1) : ImVec4(0.45f,0.45f,0.45f,1);
        ImGui::SetCursorPosX(ImGui::GetCursorPosX() + 1.f);
        ImGui::ImageWithBg(ImTextureRef(id), iconVec,
                           ImVec2(0,0), ImVec2(1,1), ImVec4(0,0,0,0), tint);
    }

    // ── Col 2: Visibility toggle ────────────────────────────────────────────
    ImGui::TableNextColumn();
    if (isImageable) {
        const float  iconSz  = ImGui::GetTextLineHeight();
        const ImVec2 iconVec(iconSz, iconSz);
        Icon   visIcon = isInvisible ? Icon::EyeSlash : Icon::Eye;
        // Invisible: EyeSlash in orange; Visible: Eye at full alpha (dimmer tint).
        ImVec4 visTint = isInvisible ? ImVec4(1.00f, 0.60f, 0.15f, 1.0f)
                                     : ImVec4(0.60f, 0.60f, 0.60f, 1.0f);
        ImTextureID id = m_iconManager ? m_iconManager->Get(visIcon) : ImTextureID_Invalid;

        ImGui::PushStyleColor(ImGuiCol_Button,        ImVec4(0,0,0,0));
        ImGui::PushStyleColor(ImGuiCol_ButtonHovered, ImVec4(1,1,1,0.12f));
        ImGui::PushStyleColor(ImGuiCol_ButtonActive,  ImVec4(1,1,1,0.20f));
        ImGui::PushStyleVar(ImGuiStyleVar_FramePadding, ImVec2(0,0));
        ImGui::SetCursorPosX(ImGui::GetCursorPosX() + 1.f);

        std::string visId = "##vis_" + primStr;
        if (ImGui::ImageButton(visId.c_str(), ImTextureRef(id), iconVec,
                               ImVec2(0,0), ImVec2(1,1), ImVec4(0,0,0,0), visTint)) {
            try {
                UsdGeomImageable img(prim);
                UsdAttribute visAttr = img.GetVisibilityAttr();
                if (isInvisible) {
                    visAttr.Set(UsdGeomTokens->inherited);
                    LOG_INFO("Made visible: " + primStr);
                } else {
                    visAttr.Set(UsdGeomTokens->invisible);
                    LOG_INFO("Made invisible: " + primStr);
                }
            } catch (const std::exception& e) {
                LOG_ERROR(std::string("Toggle visibility: ") + e.what());
            }
        }
        if (ImGui::IsItemHovered())
            ImGui::SetTooltip(isInvisible ? "Invisible — click to show"
                                          : "Visible — click to hide");
        ImGui::PopStyleVar();
        ImGui::PopStyleColor(3);
    } else {
        ImGui::Dummy(ImVec2(ImGui::GetTextLineHeight(), ImGui::GetTextLineHeight()));
    }

    // ── Col 3: Reference indicator ──────────────────────────────────────────
    ImGui::TableNextColumn();
    if (hasRefs) {
        const float  iconSz  = ImGui::GetTextLineHeight();
        const ImVec2 iconVec(iconSz, iconSz);
        ImTextureID  id = m_iconManager ? m_iconManager->Get(Icon::Link) : ImTextureID_Invalid;
        ImVec4 tint(0.45f, 0.75f, 1.0f, 1.0f);
        ImGui::SetCursorPosX(ImGui::GetCursorPosX() + 1.f);
        ImGui::ImageWithBg(ImTextureRef(id), iconVec,
                           ImVec2(0,0), ImVec2(1,1), ImVec4(0,0,0,0), tint);
        if (ImGui::IsItemHovered())
            ImGui::SetTooltip("Has references");
    }

    // ── Col 4: Specifier (def / over / class) ──────────────────────────────
    ImGui::TableNextColumn();
    {
        const char* label = nullptr;
        ImVec4      color;
        switch (spec) {
        case SdfSpecifierOver:
            label = "over";
            color = ImVec4(1.00f, 0.60f, 0.10f, isActive ? 1.0f : 0.45f);
            break;
        case SdfSpecifierClass:
            label = "class";
            color = ImVec4(0.40f, 0.70f, 1.00f, isActive ? 1.0f : 0.45f);
            break;
        default: // SdfSpecifierDef
            label = "def";
            color = ImVec4(0.45f, 0.45f, 0.45f, isActive ? 0.55f : 0.30f);
            break;
        }
        ImGui::TextColored(color, "%s", label);
    }

    // ── Recurse into children ───────────────────────────────────────────────
    // TreePop must be called in the SAME column as TreeNodeEx (col 0).
    // Since we called TableNextColumn three more times above, we must move
    // back to col 0 before TreePop.  The correct ImGui demo pattern is to
    // recurse BEFORE filling other columns, but we need icons on the same row.
    // Solution: recurse here (after columns), but ImGui only needs TreePop to
    // be inside the same Begin/End pair — column doesn't matter for TreePop.
    if (open && hasChildren) {
        // Snapshot children: stage mutations inside RenderPrimNode (reparent, group)
        // invalidate the live GetChildren() iterator.
        std::vector<UsdPrim> children(prim.GetChildren().begin(),
                                      prim.GetChildren().end());
        for (const auto& child : children)
            RenderPrimNode(child);
        ImGui::TreePop();
    }

    ImGui::PopID();
}

void SceneHierarchyPanel::RenderContextMenu(const UsdPrim& prim) {
    if (!prim.IsValid() || prim.IsPseudoRoot()) return;

    if (ImGui::BeginPopupContextItem("PrimContextMenu")) {
        std::string primName = prim.GetName().GetString();
        ImGui::TextDisabled("%s", primName.c_str());
        ImGui::Separator();

        bool isActive = prim.IsActive();
        if (ImGui::MenuItem(isActive ? "Deactivate" : "Activate")) {
            try {
                prim.SetActive(!isActive);
                LOG_INFO(std::string(!isActive ? "Activated" : "Deactivated") + " prim: " + prim.GetPath().GetString());
            } catch (const std::exception& e) {
                LOG_ERROR(std::string("Failed to toggle active: ") + e.what());
            }
        }

        bool isImageable = prim.IsA<UsdGeomImageable>();
        if (isImageable) {
            UsdGeomImageable img(prim);
            TfToken vis = img.ComputeVisibility();
            bool isInvisible = (vis == UsdGeomTokens->invisible);
            if (ImGui::MenuItem(isInvisible ? "Make Visible" : "Make Invisible")) {
                try {
                    UsdAttribute visAttr = img.GetVisibilityAttr();
                    if (isInvisible) {
                        visAttr.Set(UsdGeomTokens->inherited);
                    } else {
                        visAttr.Set(UsdGeomTokens->invisible);
                    }
                    LOG_INFO(std::string(isInvisible ? "Made visible" : "Made invisible") + ": " + prim.GetPath().GetString());
                } catch (const std::exception& e) {
                    LOG_ERROR(std::string("Failed to toggle visibility: ") + e.what());
                }
            }
        }

        ImGui::Separator();

        // ── Add Child Prim ────────────────────────────────────────────────────
        static const char* kChildPrimTypes[] = {
            "Xform", "Scope",
            "Mesh", "Sphere", "Cube", "Cylinder", "Cone", "Capsule",
            "Camera",
            "SphereLight", "DomeLight", "RectLight", "DiskLight",
            "CylinderLight", "DistantLight"
        };
        if (ImGui::BeginMenu("Add Child")) {
            SdfPath parentPath = prim.GetPath();
            for (const char* typeName : kChildPrimTypes) {
                if (ImGui::MenuItem(typeName)) {
                    std::string baseName = typeName;
                    SdfPath childPath = FindUniqueChildPath(parentPath, baseName);
                    if (m_commandHistory) {
                        m_commandHistory->Push(std::make_unique<CreatePrimCommand>(
                            m_stage, childPath, TfToken(typeName)));
                        UsdPrim child = m_stage->GetPrimAtPath(childPath);
                        if (child.IsValid()) {
                            SetSelectedPathFromClick(childPath.GetString());
                            m_scrollToSelected = true;
                        }
                    } else {
                        UsdPrim child = m_stage->DefinePrim(childPath, TfToken(typeName));
                        if (child.IsValid()) {
                            SetSelectedPathFromClick(childPath.GetString());
                            m_scrollToSelected = true;
                        } else {
                            LOG_ERROR("Failed to create child prim: " + childPath.GetString());
                        }
                    }
                }
            }
            ImGui::EndMenu();
        }

        ImGui::Separator();

        // ── Rename / Group ────────────────────────────────────────────────────
        if (ImGui::MenuItem("Rename", "F2")) {
            m_renamingPath = prim.GetPath();
            std::string name = prim.GetName().GetString();
            strncpy(m_renameBuf, name.c_str(), sizeof(m_renameBuf) - 1);
            m_renameBuf[sizeof(m_renameBuf) - 1] = '\0';
            m_renameJustStarted = true;
            m_scrollToSelected  = true;
        }
        bool canUnparent = (prim.GetPath().GetPathElementCount() > 1);
        if (ImGui::MenuItem("Unparent to Root", nullptr, false, canUnparent)) {
            if (m_commandHistory && m_stage) {
                std::string name = FindUniqueChildPath(SdfPath("/"), prim.GetName().GetString()).GetName();
                m_commandHistory->Push(std::make_unique<ReparentPrimCommand>(
                    m_stage, prim.GetPath(), SdfPath("/"), name));
                SetSelectedPathFromClick(SdfPath("/").AppendChild(TfToken(name)).GetString());
            }
        }
        if (ImGui::MenuItem("Group", "Ctrl+G")) {
            std::vector<SdfPath> srcs;
            if (m_selectedPaths.count(prim.GetPath().GetString()))
                for (const auto& s : m_selectedPaths) srcs.emplace_back(s);
            else
                srcs.push_back(prim.GetPath());
            if (!srcs.empty() && m_commandHistory && m_stage) {
                // Filter descendants
                std::vector<SdfPath> filtered;
                for (const SdfPath& p : srcs) {
                    bool hasAncestor = false;
                    for (const SdfPath& other : srcs)
                        if (other != p && p.HasPrefix(other)) { hasAncestor = true; break; }
                    if (!hasAncestor) filtered.push_back(p);
                }
                if (!filtered.empty()) {
                    SdfPath commonParent = filtered[0].GetParentPath();
                    for (const SdfPath& p : filtered)
                        while (!p.HasPrefix(commonParent)) commonParent = commonParent.GetParentPath();
                    SdfPath groupPath = FindUniqueChildPath(commonParent, "group1");
                    m_commandHistory->Push(std::make_unique<GroupPrimsCommand>(m_stage, filtered, groupPath));
                    SetSelectedPathFromClick(groupPath.GetString());
                }
            }
        }

        // Unparent or Group may have moved this prim, expiring it.
        // Close the popup now before touching any more prim state.
        if (!prim.IsValid()) {
            ImGui::EndPopup();
            return;
        }

        ImGui::Separator();

        bool hasChildren = !prim.GetChildren().empty();
        if (ImGui::MenuItem("Expand Children", nullptr, false, hasChildren)) {
            ImGui::GetStateStorage()->SetInt(ImGui::GetID(prim.GetPath().GetText()), 1);
        }
        if (ImGui::MenuItem("Collapse Children", nullptr, false, hasChildren)) {
            ImGui::GetStateStorage()->SetInt(ImGui::GetID(prim.GetPath().GetText()), 0);
        }

        ImGui::Separator();

        bool hasRefs     = prim.HasAuthoredReferences();
        bool hasPayloads = prim.HasAuthoredPayloads();

        // ---- Reference submenu ----
        if (ImGui::BeginMenu("Reference")) {
            if (ImGui::MenuItem("Add...")) {
                std::string filePath = FileDialog::OpenFile(
                    "USD Files (*.usd;*.usda;*.usdc;*.usdz)\0*.usd;*.usda;*.usdc;*.usdz\0All Files (*.*)\0*.*\0",
                    "Add Reference File");
                if (!filePath.empty()) {
                    try {
                        bool ok = prim.GetReferences().AddReference(filePath);
                        if (ok)
                            LOG_INFO("Added reference '" + filePath + "' to prim: " + prim.GetPath().GetString());
                        else
                            LOG_ERROR("Failed to add reference '" + filePath + "' to: " + prim.GetPath().GetString());
                    } catch (const std::exception& e) {
                        LOG_ERROR(std::string("Add reference error: ") + e.what());
                    }
                }
            }
            ImGui::Separator();

            if (ImGui::BeginMenu("Replace", hasRefs)) {
                UsdPrimCompositionQuery::Filter replFilter;
                replFilter.arcTypeFilter = UsdPrimCompositionQuery::ArcTypeFilter::Reference;
                replFilter.dependencyTypeFilter = UsdPrimCompositionQuery::DependencyTypeFilter::Direct;
                UsdPrimCompositionQuery replQuery(prim, replFilter);

                bool anyRepl = false;
                for (auto& arc : replQuery.GetCompositionArcs()) {
                    SdfReferenceEditorProxy editor;
                    SdfReference oldRef;
                    if (arc.GetIntroducingListEditor(&editor, &oldRef)) {
                        std::string label = oldRef.GetAssetPath().empty()
                            ? "(internal reference)"
                            : oldRef.GetAssetPath();
                        if (ImGui::MenuItem(label.c_str())) {
                            m_pendingReplaceRef     = oldRef;
                            m_pendingReplaceRefPrim = prim.GetPath();
                            m_doReplaceRefPick      = true;
                        }
                        anyRepl = true;
                    }
                }
                if (!anyRepl)
                    ImGui::TextDisabled("(no direct references)");
                ImGui::EndMenu();
            }

            if (ImGui::BeginMenu("Remove", hasRefs)) {
                UsdPrimCompositionQuery::Filter filter;
                filter.arcTypeFilter = UsdPrimCompositionQuery::ArcTypeFilter::Reference;
                filter.dependencyTypeFilter = UsdPrimCompositionQuery::DependencyTypeFilter::Direct;
                UsdPrimCompositionQuery query(prim, filter);

                bool anyListed = false;
                for (auto& arc : query.GetCompositionArcs()) {
                    SdfReferenceEditorProxy editor;
                    SdfReference ref;
                    if (arc.GetIntroducingListEditor(&editor, &ref)) {
                        std::string label = ref.GetAssetPath().empty()
                            ? "(internal reference)"
                            : ref.GetAssetPath();
                        if (ImGui::MenuItem(label.c_str())) {
                            try {
                                prim.GetReferences().RemoveReference(ref);
                                LOG_INFO("Removed reference '" + label + "' from: " + prim.GetPath().GetString());
                            } catch (const std::exception& e) {
                                LOG_ERROR(std::string("Remove reference error: ") + e.what());
                            }
                        }
                        anyListed = true;
                    }
                }
                if (anyListed) ImGui::Separator();
                if (ImGui::MenuItem("Clear All")) {
                    try {
                        prim.GetReferences().ClearReferences();
                        LOG_INFO("Cleared all references on: " + prim.GetPath().GetString());
                    } catch (const std::exception& e) {
                        LOG_ERROR(std::string("Clear references error: ") + e.what());
                    }
                }
                ImGui::EndMenu();
            }

            ImGui::EndMenu();
        }

        // ---- Payload submenu ----
        if (ImGui::BeginMenu("Payload")) {
            if (ImGui::MenuItem("Add...")) {
                std::string filePath = FileDialog::OpenFile(
                    "USD Files (*.usd;*.usda;*.usdc;*.usdz)\0*.usd;*.usda;*.usdc;*.usdz\0All Files (*.*)\0*.*\0",
                    "Add Payload File");
                if (!filePath.empty()) {
                    try {
                        bool ok = prim.GetPayloads().AddPayload(filePath);
                        if (ok)
                            LOG_INFO("Added payload '" + filePath + "' to prim: " + prim.GetPath().GetString());
                        else
                            LOG_ERROR("Failed to add payload '" + filePath + "' to: " + prim.GetPath().GetString());
                    } catch (const std::exception& e) {
                        LOG_ERROR(std::string("Add payload error: ") + e.what());
                    }
                }
            }
            ImGui::Separator();

            if (ImGui::BeginMenu("Remove", hasPayloads)) {
                UsdPrimCompositionQuery::Filter filter;
                filter.arcTypeFilter = UsdPrimCompositionQuery::ArcTypeFilter::Payload;
                filter.dependencyTypeFilter = UsdPrimCompositionQuery::DependencyTypeFilter::Direct;
                UsdPrimCompositionQuery query(prim, filter);

                bool anyListed = false;
                for (auto& arc : query.GetCompositionArcs()) {
                    SdfPayloadEditorProxy editor;
                    SdfPayload payload;
                    if (arc.GetIntroducingListEditor(&editor, &payload)) {
                        std::string label = payload.GetAssetPath().empty()
                            ? "(internal payload)"
                            : payload.GetAssetPath();
                        if (ImGui::MenuItem(label.c_str())) {
                            try {
                                prim.GetPayloads().RemovePayload(payload);
                                LOG_INFO("Removed payload '" + label + "' from: " + prim.GetPath().GetString());
                            } catch (const std::exception& e) {
                                LOG_ERROR(std::string("Remove payload error: ") + e.what());
                            }
                        }
                        anyListed = true;
                    }
                }
                if (anyListed) ImGui::Separator();
                if (ImGui::MenuItem("Clear All")) {
                    try {
                        prim.GetPayloads().ClearPayloads();
                        LOG_INFO("Cleared all payloads on: " + prim.GetPath().GetString());
                    } catch (const std::exception& e) {
                        LOG_ERROR(std::string("Clear payloads error: ") + e.what());
                    }
                }
                ImGui::EndMenu();
            }

            ImGui::EndMenu();
        }

        // ---- Inherit submenu ----
        bool hasInherits = prim.HasAuthoredInherits();
        if (ImGui::BeginMenu("Inherit")) {
            if (ImGui::MenuItem("Add...")) {
                m_arcModalMode       = ArcModalMode::Inherit;
                m_arcModalTargetPrim = prim.GetPath();
                m_arcModalBuf[0]     = '\0';
                m_openArcModal       = true;
            }
            ImGui::Separator();
            if (ImGui::BeginMenu("Remove", hasInherits)) {
                SdfPathVector paths = prim.GetInherits().GetAllDirectInherits();
                bool any = false;
                for (const SdfPath& ipath : paths) {
                    if (ImGui::MenuItem(ipath.GetText())) {
                        try {
                            prim.GetInherits().RemoveInherit(ipath);
                            LOG_INFO("Removed inherit '" + ipath.GetString() + "' from: " + prim.GetPath().GetString());
                        } catch (const std::exception& e) {
                            LOG_ERROR(std::string("Remove inherit error: ") + e.what());
                        }
                    }
                    any = true;
                }
                if (any) ImGui::Separator();
                if (ImGui::MenuItem("Clear All")) {
                    try {
                        prim.GetInherits().ClearInherits();
                        LOG_INFO("Cleared all inherits on: " + prim.GetPath().GetString());
                    } catch (const std::exception& e) {
                        LOG_ERROR(std::string("Clear inherits error: ") + e.what());
                    }
                }
                ImGui::EndMenu();
            }
            ImGui::EndMenu();
        }

        // ---- Specialize submenu ----
        bool hasSpecializes = prim.HasAuthoredSpecializes();
        if (ImGui::BeginMenu("Specialize")) {
            if (ImGui::MenuItem("Add...")) {
                m_arcModalMode       = ArcModalMode::Specialize;
                m_arcModalTargetPrim = prim.GetPath();
                m_arcModalBuf[0]     = '\0';
                m_openArcModal       = true;
            }
            ImGui::Separator();
            if (ImGui::BeginMenu("Remove", hasSpecializes)) {
                UsdPrimCompositionQuery::Filter filter;
                filter.arcTypeFilter        = UsdPrimCompositionQuery::ArcTypeFilter::Specialize;
                filter.dependencyTypeFilter = UsdPrimCompositionQuery::DependencyTypeFilter::Direct;
                UsdPrimCompositionQuery query(prim, filter);

                bool any = false;
                for (auto& arc : query.GetCompositionArcs()) {
                    SdfPathEditorProxy editor;
                    SdfPath specPath;
                    if (arc.GetIntroducingListEditor(&editor, &specPath)) {
                        if (ImGui::MenuItem(specPath.GetText())) {
                            try {
                                prim.GetSpecializes().RemoveSpecialize(specPath);
                                LOG_INFO("Removed specialize '" + specPath.GetString() + "' from: " + prim.GetPath().GetString());
                            } catch (const std::exception& e) {
                                LOG_ERROR(std::string("Remove specialize error: ") + e.what());
                            }
                        }
                        any = true;
                    }
                }
                if (any) ImGui::Separator();
                if (ImGui::MenuItem("Clear All")) {
                    try {
                        prim.GetSpecializes().ClearSpecializes();
                        LOG_INFO("Cleared all specializes on: " + prim.GetPath().GetString());
                    } catch (const std::exception& e) {
                        LOG_ERROR(std::string("Clear specializes error: ") + e.what());
                    }
                }
                ImGui::EndMenu();
            }
            ImGui::EndMenu();
        }

        // ---- VariantSet submenu ----
        if (ImGui::BeginMenu("VariantSet")) {
            if (ImGui::MenuItem("Add...")) {
                m_arcModalMode       = ArcModalMode::VariantSet;
                m_arcModalTargetPrim = prim.GetPath();
                m_arcModalBuf[0]     = '\0';
                m_openArcModal       = true;
            }
            bool hasVarSets = prim.HasVariantSets();
            if (hasVarSets) {
                ImGui::Separator();
                std::vector<std::string> vsNames;
                prim.GetVariantSets().GetNames(&vsNames);
                for (const std::string& vsName : vsNames) {
                    if (ImGui::BeginMenu(vsName.c_str())) {
                        if (ImGui::MenuItem("Add Variant...")) {
                            m_arcModalMode       = ArcModalMode::Variant;
                            m_arcModalTargetPrim = prim.GetPath();
                            m_arcModalVarSetName = vsName;
                            m_arcModalBuf[0]     = '\0';
                            m_openArcModal       = true;
                        }
                        UsdVariantSet vs         = prim.GetVariantSet(vsName);
                        std::vector<std::string> varNames = vs.GetVariantNames();
                        if (!varNames.empty()) {
                            ImGui::Separator();
                            std::string currentSel = vs.GetVariantSelection();
                            for (const std::string& varName : varNames) {
                                bool isCurrent = (varName == currentSel);
                                if (ImGui::MenuItem(varName.c_str(), nullptr, isCurrent)) {
                                    try {
                                        vs.SetVariantSelection(varName);
                                        LOG_INFO("Set variant '" + varName + "' on '" + vsName + "' for: " + prim.GetPath().GetString());
                                    } catch (const std::exception& e) {
                                        LOG_ERROR(std::string("Set variant selection error: ") + e.what());
                                    }
                                }
                            }
                        }
                        ImGui::EndMenu();
                    }
                }
            }
            ImGui::EndMenu();
        }

        // ---- Prim removal ----
        // Only show "Remove Prim" for prims that have a local spec authored in the root
        // layer.  Prims brought in purely via composition from an external referenced
        // stage have no local spec and cannot be removed directly.
        ImGui::Separator();
        {
            auto rootLayer = m_stage->GetRootLayer();
            bool hasLocalSpec = rootLayer && !!rootLayer->GetPrimAtPath(prim.GetPath());
            if (ImGui::MenuItem("Remove Prim", nullptr, false, hasLocalSpec)) {
                // Defer to the confirm modal — can't open a modal from inside a popup.
                m_pendingRemovePrimPath = prim.GetPath();
                m_showRemovePrimConfirm = true;
            }
        }

        ImGui::EndPopup();
    }
}

void SceneHierarchyPanel::RenderRemovePrimModal() {
    if (m_showRemovePrimConfirm) {
        ImGui::OpenPopup("Remove Prim##confirm");
        m_showRemovePrimConfirm = false;
    }

    // Centre the modal over the main viewport.
    ImVec2 center = ImGui::GetMainViewport()->GetCenter();
    ImGui::SetNextWindowPos(center, ImGuiCond_Always, ImVec2(0.5f, 0.5f));
    ImGui::SetNextWindowSize(ImVec2(420, 0), ImGuiCond_Always);

    if (ImGui::BeginPopupModal("Remove Prim##confirm", nullptr,
            ImGuiWindowFlags_AlwaysAutoResize | ImGuiWindowFlags_NoMove)) {

        ImGui::TextUnformatted("Are you sure you want to remove this prim?");
        ImGui::Spacing();
        ImGui::TextColored(ImVec4(1.0f, 0.8f, 0.2f, 1.0f), "%s",
            m_pendingRemovePrimPath.GetText());
        ImGui::Spacing();
        ImGui::TextDisabled("This will remove the prim spec from the root layer.\n"
                            "Child prims authored locally will also be removed.");
        ImGui::Separator();

        float buttonWidth = 120.0f;
        float spacing = ImGui::GetStyle().ItemSpacing.x;
        float totalW = buttonWidth * 2.0f + spacing;
        ImGui::SetCursorPosX((ImGui::GetContentRegionAvail().x - totalW) * 0.5f +
                              ImGui::GetCursorPosX());

        if (ImGui::Button("Remove", ImVec2(buttonWidth, 0))) {
            if (m_stage && !m_pendingRemovePrimPath.IsEmpty()) {
                if (m_commandHistory) {
                    // Snapshot the spec BEFORE deletion, then push.
                    auto cmd = std::make_unique<DeletePrimCommand>(
                        m_stage, m_pendingRemovePrimPath);
                    std::string removedStr = m_pendingRemovePrimPath.GetString();
                    m_commandHistory->Push(std::move(cmd));
                    // Clear selection if removed prim was selected.
                    if (m_primarySelectedPath == removedStr) {
                        m_primarySelectedPath.clear();
                        m_primarySdfPath = SdfPath();
                        m_selectedPaths.clear();
                        if (m_onPrimSelected) m_onPrimSelected("");
                    } else {
                        m_selectedPaths.erase(removedStr);
                    }
                } else {
                    try {
                        bool ok = m_stage->RemovePrim(m_pendingRemovePrimPath);
                        if (ok) {
                            LOG_INFO("Removed prim: " + m_pendingRemovePrimPath.GetString());
                            std::string removedStr = m_pendingRemovePrimPath.GetString();
                            if (m_primarySelectedPath == removedStr) {
                                m_primarySelectedPath.clear();
                                m_primarySdfPath = SdfPath();
                                m_selectedPaths.clear();
                                if (m_onPrimSelected) m_onPrimSelected("");
                            } else {
                                m_selectedPaths.erase(removedStr);
                            }
                        } else {
                            LOG_ERROR("Failed to remove prim: " + m_pendingRemovePrimPath.GetString());
                        }
                    } catch (const std::exception& e) {
                        LOG_ERROR(std::string("Remove prim error: ") + e.what());
                    }
                }
                m_pendingRemovePrimPath = SdfPath();
            }
            ImGui::CloseCurrentPopup();
        }

        ImGui::SameLine();

        if (ImGui::Button("Cancel", ImVec2(buttonWidth, 0))) {
            m_pendingRemovePrimPath = SdfPath();
            ImGui::CloseCurrentPopup();
        }

        ImGui::EndPopup();
    }
}

void SceneHierarchyPanel::RenderArcModal() {
    if (m_openArcModal) {
        ImGui::OpenPopup("Arc Operation##arc");
        m_openArcModal = false;
    }

    ImVec2 center = ImGui::GetMainViewport()->GetCenter();
    ImGui::SetNextWindowPos(center, ImGuiCond_Always, ImVec2(0.5f, 0.5f));
    ImGui::SetNextWindowSize(ImVec2(420, 0), ImGuiCond_Always);

    if (ImGui::BeginPopupModal("Arc Operation##arc", nullptr,
            ImGuiWindowFlags_AlwaysAutoResize | ImGuiWindowFlags_NoMove)) {

        const char* title     = "Arc Operation";
        const char* prompt    = "Path:";
        bool        isNameInput = false;

        switch (m_arcModalMode) {
            case ArcModalMode::Inherit:
                title  = "Add Inherit";
                prompt = "Inherit path (e.g. /ClassName):";
                break;
            case ArcModalMode::Specialize:
                title  = "Add Specialize";
                prompt = "Specialize path (e.g. /BasePrim):";
                break;
            case ArcModalMode::VariantSet:
                title       = "Add VariantSet";
                prompt      = "VariantSet name:";
                isNameInput = true;
                break;
            case ArcModalMode::Variant:
                title       = "Add Variant";
                prompt      = "Variant name:";
                isNameInput = true;
                break;
            default:
                ImGui::CloseCurrentPopup();
                ImGui::EndPopup();
                return;
        }

        ImGui::TextUnformatted(title);
        ImGui::Separator();
        ImGui::Spacing();
        ImGui::TextDisabled("Prim: %s", m_arcModalTargetPrim.GetText());
        if (m_arcModalMode == ArcModalMode::Variant)
            ImGui::TextDisabled("VariantSet: %s", m_arcModalVarSetName.c_str());
        ImGui::Spacing();
        ImGui::TextUnformatted(prompt);
        ImGui::SetNextItemWidth(-1.0f);
        if (ImGui::IsWindowAppearing())
            ImGui::SetKeyboardFocusHere();
        bool confirm = ImGui::InputText("##arcbuf", m_arcModalBuf,
                                        sizeof(m_arcModalBuf),
                                        ImGuiInputTextFlags_EnterReturnsTrue);
        ImGui::Spacing();
        ImGui::Separator();

        bool inputNonEmpty = (m_arcModalBuf[0] != '\0');
        float buttonWidth  = 120.0f;
        float spacing      = ImGui::GetStyle().ItemSpacing.x;
        float totalW       = buttonWidth * 2.0f + spacing;
        ImGui::SetCursorPosX(
            (ImGui::GetContentRegionAvail().x - totalW) * 0.5f + ImGui::GetCursorPosX());

        // Capture Enter-confirm before BeginDisabled so it works regardless of button state.
        bool doConfirm = (confirm && inputNonEmpty);

        if (!inputNonEmpty) ImGui::BeginDisabled();
        if (ImGui::Button("OK", ImVec2(buttonWidth, 0)))
            doConfirm = inputNonEmpty;
        if (!inputNonEmpty) ImGui::EndDisabled();

        ImGui::SameLine();

        if (ImGui::Button("Cancel", ImVec2(buttonWidth, 0))) {
            m_arcModalMode       = ArcModalMode::None;
            m_arcModalTargetPrim = SdfPath();
            m_arcModalVarSetName.clear();
            ImGui::CloseCurrentPopup();
        }

        if (doConfirm) {
            std::string raw(m_arcModalBuf);
            std::string val = isNameInput ? SanitizeUsdName(raw) : raw;

            if (m_stage && !m_arcModalTargetPrim.IsEmpty()) {
                UsdPrim prim = m_stage->GetPrimAtPath(m_arcModalTargetPrim);
                if (prim.IsValid()) {
                    try {
                        switch (m_arcModalMode) {
                            case ArcModalMode::Inherit: {
                                SdfPath ipath(val);
                                bool ok = prim.GetInherits().AddInherit(ipath);
                                if (ok)
                                    LOG_INFO("Added inherit '" + val + "' to: " + m_arcModalTargetPrim.GetString());
                                else
                                    LOG_ERROR("Failed to add inherit '" + val + "' to: " + m_arcModalTargetPrim.GetString());
                                break;
                            }
                            case ArcModalMode::Specialize: {
                                SdfPath spath(val);
                                bool ok = prim.GetSpecializes().AddSpecialize(spath);
                                if (ok)
                                    LOG_INFO("Added specialize '" + val + "' to: " + m_arcModalTargetPrim.GetString());
                                else
                                    LOG_ERROR("Failed to add specialize '" + val + "' to: " + m_arcModalTargetPrim.GetString());
                                break;
                            }
                            case ArcModalMode::VariantSet: {
                                prim.GetVariantSets().AddVariantSet(val);
                                LOG_INFO("Added variantSet '" + val + "' to: " + m_arcModalTargetPrim.GetString());
                                break;
                            }
                            case ArcModalMode::Variant: {
                                UsdVariantSet vs = prim.GetVariantSet(m_arcModalVarSetName);
                                bool ok = vs.AddVariant(val);
                                if (ok)
                                    LOG_INFO("Added variant '" + val + "' to variantSet '" + m_arcModalVarSetName + "' on: " + m_arcModalTargetPrim.GetString());
                                else
                                    LOG_ERROR("Failed to add variant '" + val + "' to variantSet '" + m_arcModalVarSetName + "'");
                                break;
                            }
                            default: break;
                        }
                    } catch (const std::exception& e) {
                        LOG_ERROR(std::string("Arc operation error: ") + e.what());
                    }
                } else {
                    LOG_ERROR("Arc modal: prim no longer valid: " + m_arcModalTargetPrim.GetString());
                }
            }
            m_arcModalMode       = ArcModalMode::None;
            m_arcModalTargetPrim = SdfPath();
            m_arcModalVarSetName.clear();
            ImGui::CloseCurrentPopup();
        }

        ImGui::EndPopup();
    }
}

void SceneHierarchyPanel::ProcessPendingReplaceRef() {
    if (!m_doReplaceRefPick) return;
    m_doReplaceRefPick = false;

    if (!m_stage || m_pendingReplaceRefPrim.IsEmpty()) return;

    std::string newPath = FileDialog::OpenFile(
        "USD Files (*.usd;*.usda;*.usdc;*.usdz)\0*.usd;*.usda;*.usdc;*.usdz\0All Files (*.*)\0*.*\0",
        "Replace Reference File");
    if (newPath.empty()) return;

    UsdPrim prim = m_stage->GetPrimAtPath(m_pendingReplaceRefPrim);
    if (!prim.IsValid()) {
        LOG_ERROR("Replace reference: prim no longer valid: " + m_pendingReplaceRefPrim.GetString());
        return;
    }

    try {
        // Build the new SdfReference preserving prim path and layer offset.
        SdfReference newRef(newPath,
                            m_pendingReplaceRef.GetPrimPath(),
                            m_pendingReplaceRef.GetLayerOffset());

        if (m_commandHistory) {
            m_commandHistory->Push(std::make_unique<ReplaceReferenceCommand>(
                m_stage, m_pendingReplaceRefPrim, m_pendingReplaceRef, newRef));
        } else {
            UsdReferences refs = prim.GetReferences();
            bool removed = refs.RemoveReference(m_pendingReplaceRef);
            if (!removed) {
                LOG_ERROR("Replace reference: failed to remove old reference '" +
                          m_pendingReplaceRef.GetAssetPath() + "'");
            } else {
                bool added = refs.AddReference(newRef);
                if (added) {
                    LOG_INFO("Replaced reference '" + m_pendingReplaceRef.GetAssetPath() +
                             "' -> '" + newPath + "' on prim: " + m_pendingReplaceRefPrim.GetString());
                } else {
                    LOG_ERROR("Replace reference: failed to add new reference '" + newPath + "'");
                }
            }
        }
    } catch (const std::exception& e) {
        LOG_ERROR(std::string("Replace reference error: ") + e.what());
    }

    m_pendingReplaceRefPrim = SdfPath();
    m_pendingReplaceRef     = SdfReference();
}

} // namespace UsdLayerManager
