UsdLayerManager/src/ui/SceneHierarchyPanel.h

#pragma once

#include "../core/PropertyManager.h"
#include "../core/CommandHistory.h"
#include "../core/LayerManager.h"
#include "IconManager.h"
#include <pxr/usd/usd/stage.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/usd/sdf/layer.h>
#include <pxr/usd/sdf/reference.h>
#include <imgui.h>
#include <string>
#include <vector>
#include <unordered_set>
#include <functional>

PXR_NAMESPACE_USING_DIRECTIVE

namespace UsdLayerManager {

class SceneHierarchyPanel {
public:
    SceneHierarchyPanel();
    ~SceneHierarchyPanel();
    
    void SetPropertyManager(PropertyManager* manager);
    void SetCommandHistory(CommandHistory* history) { m_commandHistory = history; }
    void SetLayerManager(LayerManager* lm) { m_layerManager = lm; }
    void SetStage(UsdStageRefPtr stage);
    void SetIconManager(IconManager* iconManager) { m_iconManager = iconManager; }
    void Render();
    
    std::string GetSelectedPrimPath() const { return m_primarySelectedPath; }
    UsdPrim GetSelectedPrim() const;

    /// Set single selection from hierarchy click (fires callback).
    /// Also clears any rect multi-selection.
    void SetSelectedPathFromClick(const std::string& path);

    /// Set single selection from viewport single-click (no callback, scroll-to).
    void SetSelectedPath(const std::string& path) {
        m_selectedPaths.clear();
        m_primarySelectedPath  = path;
        m_primarySdfPath       = path.empty() ? SdfPath() : SdfPath(path);
        if (!path.empty()) m_selectedPaths.insert(path);
        m_scrollToSelected     = !path.empty();
    }

    /// Set multi-selection from viewport rect pick (no callback, scroll-to first).
    void SetSelectedPaths(const std::vector<std::string>& paths) {
        m_selectedPaths.clear();
        m_primarySelectedPath.clear();
        m_primarySdfPath = SdfPath();
        for (const auto& p : paths) m_selectedPaths.insert(p);
        if (!paths.empty()) {
            m_primarySelectedPath = paths.front();
            m_primarySdfPath      = SdfPath(paths.front());
            m_scrollToSelected    = true;
        }
    }

    using PrimSelectCallback = std::function<void(const std::string& path)>;
    void SetOnPrimSelected(PrimSelectCallback callback) { m_onPrimSelected = callback; }

    using MultiPrimSelectCallback = std::function<void(const std::vector<std::string>&)>;
    void SetOnPrimsSelected(MultiPrimSelectCallback cb) { m_onPrimsSelected = cb; }

    /// Called when stage-level metadata (e.g. up axis) is changed via the hierarchy panel.
    using StageMetadataChangedCallback = std::function<void()>;
    void SetOnStageMetadataChanged(StageMetadataChangedCallback callback) { m_onStageMetadataChanged = callback; }
    
    enum class ArcModalMode { None, Inherit, Specialize, VariantSet, Variant };

private:
    void RenderPrimNode(const UsdPrim& prim);
    const char* GetPrimTypeIcon(const UsdPrim& prim) const;
    Icon        GetPrimTypeIconEnum(const UsdPrim& prim) const;
    void RenderContextMenu(const UsdPrim& prim);
    void RenderRemovePrimModal();
    void RenderArcModal();
    void ProcessPendingReplaceRef();
    
    PropertyManager*             m_propertyManager;
    CommandHistory*              m_commandHistory = nullptr;
    LayerManager*                m_layerManager   = nullptr;
    IconManager*                 m_iconManager = nullptr;
    UsdStageRefPtr               m_stage;

    /// Primary path: the scroll/frame target; also used for F-to-frame.
    std::string                  m_primarySelectedPath;
    SdfPath                      m_primarySdfPath;
    /// Full set of selected paths (supports multi-select from rect pick).
    std::unordered_set<std::string> m_selectedPaths;

    bool                         m_scrollToSelected  = false;

    /// Stage-local layer identifiers rebuilt once per Render() for override detection.
    /// Contains only the stage's own layers (root + sublayers + session),
    /// NOT layers that came in through references or payloads.
    std::unordered_set<std::string> m_localLayers;

    /// Visible prim order rebuilt each frame — used for Shift+RMB range select.
    std::vector<pxr::SdfPath>        m_visiblePrimOrder;
    pxr::SdfPath                     m_rangeAnchorPath;

    PrimSelectCallback           m_onPrimSelected;
    MultiPrimSelectCallback      m_onPrimsSelected;
    StageMetadataChangedCallback m_onStageMetadataChanged;

    /// Remove-prim confirmation state.
    bool                         m_showRemovePrimConfirm = false;
    SdfPath                      m_pendingRemovePrimPath;

    /// Replace-reference deferred state (file dialog must run outside popup stack).
    bool                         m_doReplaceRefPick  = false;
    SdfPath                      m_pendingReplaceRefPrim;
    pxr::SdfReference            m_pendingReplaceRef;

    /// Arc modal state (Inherit / Specialize / VariantSet / Variant).
    ArcModalMode                 m_arcModalMode       = ArcModalMode::None;
    SdfPath                      m_arcModalTargetPrim;
    std::string                  m_arcModalVarSetName;
    char                         m_arcModalBuf[512]   = {};
    bool                         m_openArcModal       = false;

    /// Inline rename state.
    SdfPath                      m_renamingPath;          ///< non-empty while renaming
    char                         m_renameBuf[256]     = {};
    bool                         m_renameJustStarted  = false;

    void         HandleKeyboardShortcuts();
    void         ToggleSelectionFromClick(const std::string& path);
    void         RangeSelectToPath(const pxr::SdfPath& path);
    SdfPath      FindUniqueChildPath(const SdfPath& parent, const std::string& baseName);
};

} // namespace UsdLayerManager