#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 <pxr/usd/usd/prim.h>
#include <pxr/usd/usd/primRange.h>
#include <pxr/usd/usd/references.h>
#include <pxr/usd/usd/primCompositionQuery.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/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);
    // No scroll — user clicked the item directly, it's already visible.
    m_scrollToSelected = false;
    if (m_onPrimSelected) m_onPrimSelected(path);
}

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;
}

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

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

        // 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;

        if (ImGui::BeginTable("##primtree", 4, tblFlags)) {
            // Col 0 stretches; cols 1-3 are small fixed-width icon columns.
            ImGui::TableSetupColumn("##prim", ImGuiTableColumnFlags_WidthStretch);
            ImGui::TableSetupColumn("##type", ImGuiTableColumnFlags_WidthFixed, kIconW);
            ImGui::TableSetupColumn("##vis",  ImGuiTableColumnFlags_WidthFixed, kIconW);
            ImGui::TableSetupColumn("##ref",  ImGuiTableColumnFlags_WidthFixed, kIconW);

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

            ImGui::EndTable();
        }

        // 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()) {
                                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();

        // ---- Add Reference ----
        if (ImGui::MenuItem("Add Reference...")) {
            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()) {
                // Derive a valid USD prim name from the file's stem.
                std::string stem = std::filesystem::path(filePath).stem().string();
                std::string xformName = SanitizeUsdName(stem);
                if (xformName.empty()) xformName = "Reference";

                // Avoid name collision: append _N if the path already exists.
                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::EndPopup();
    }

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

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

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();

    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();
    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;
    }

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

    bool open = ImGui::TreeNodeEx(displayName.c_str(), flags);

    if (pushedColor)
        ImGui::PopStyleColor();

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

    // Selection on click (not on toggle arrow).
    if (ImGui::IsItemClicked() && !ImGui::IsItemToggledOpen())
        SetSelectedPathFromClick(primStr);

    // 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());
    }

    // Context menu (must follow the last widget = the tree node).
    RenderContextMenu(prim);

    // ── 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;
        ImVec4 visTint = isInvisible ? ImVec4(0.45f, 0.45f, 0.45f, 0.6f)
                                     : ImVec4(0.9f,  0.9f,  0.9f,  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);

        if (ImGui::ImageButton("##vis", 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");
    }

    // ── 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) {
        for (const auto& child : prim.GetChildren())
            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();

        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();

        // ---- Reference operations ----
        if (ImGui::MenuItem("Add Reference...")) {
            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());
                }
            }
        }

        bool hasRefs = prim.HasAuthoredReferences();

        // ---- Replace Reference ----
        if (ImGui::BeginMenu("Replace Reference", 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())) {
                        // NOTE: file dialog is blocking — close popup first via deferred path.
                        m_pendingReplaceRef     = oldRef;
                        m_pendingReplaceRefPrim = prim.GetPath();
                        m_doReplaceRefPick      = true;
                    }
                    anyRepl = true;
                }
            }
            if (!anyRepl) {
                ImGui::TextDisabled("(no direct references)");
            }
            ImGui::EndMenu();
        }

        // ---- Remove Reference ----
        if (ImGui::BeginMenu("Remove Reference", hasRefs)) {
            // Collect direct reference arcs via composition query.
            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 References")) {
                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();
        }

        // ---- 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::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