UsdLayerManager/src/ui/MaterialEditorPanel.h

#pragma once

#include "../core/CommandHistory.h"
#include "../core/MaterialManager.h"
#include "IconManager.h"
#include "MaterialPreviewRenderer.h"
#include "NodeThumbnailCache.h"
#include <pxr/usd/usd/stage.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/usd/sdf/valueTypeName.h>
#include <pxr/base/vt/value.h>
#include <imgui_node_editor.h>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

namespace UsdLayerManager {

/// Node-graph editor for authoring UsdShade material networks, with a
/// shader-ball preview rendered through Hydra.
class MaterialEditorPanel {
public:
    MaterialEditorPanel();
    ~MaterialEditorPanel();

    void SetStage(pxr::UsdStageRefPtr stage);
    void SetMaterialManager(MaterialManager* mgr) { m_materialManager = mgr; }
    void SetCommandHistory(CommandHistory* history) { m_commandHistory = history; }
    void SetIconManager(IconManager* icons) { m_iconManager = icons; }

    /// Forwarded to the shader-ball preview so it matches the main viewport's
    /// global color-correction settings.
    void SetColorCorrectionFromPrefs(int ccMode, const std::string& ocioDisplay,
                                     const std::string& ocioView,
                                     const std::string& ocioColorSpace,
                                     const std::string& ocioLook);

    /// Opens the material at pathStr in the graph work area, creating it (as a
    /// "Material" prim) if it doesn't exist yet. Also selects it in the browser.
    void OpenOrCreateMaterial(const std::string& pathStr);

    /// Tracks the current selection (mirrors PropertyPanel/ViewportPanel's
    /// SetSelectedPrimPath) so the toolbar can offer "create + bind" / "bind"
    /// affordances, and so selecting a prim with an already-bound material
    /// auto-loads it into the canvas.
    void SetTargetPrimPath(const std::string& path);

    /// Splitter-adjustable column widths, persisted in preferences.ini by
    /// Application. Setter clamps to the same range as the splitters.
    void GetColumnWidths(float& browserWidth, float& previewWidth) const {
        browserWidth = m_browserWidth;
        previewWidth = m_previewWidth;
    }
    void SetColumnWidths(float browserWidth, float previewWidth);

    /// Whether toggling a node's pin display mode / revealing a hidden pin
    /// authors USD customData (persisted) or stays session-only. Pushed in by
    /// Application from AppPreferences, live-updated when the user flips the
    /// checkbox in Preferences -> Material.
    void SetKeepGraphNodeViewSettingsInUsd(bool value) { m_keepGraphNodeViewSettingsInUsd = value; }

    void Render();

private:
    /// Resolved endpoint of a rendered pin, keyed by its ax::NodeEditor PinId.
    struct PinInfo {
        pxr::SdfPath nodePath;
        std::string  name;
        bool         isOutput = false;
        pxr::SdfValueTypeName typeName;
        /// True for the collapsed "more pins" connector nub rendered on a
        /// node's title row when it's hiding unconnected pins (Hypershade
        /// style) — name/typeName are unset. Connecting to/from it opens
        /// RenderPinPickMenu() to choose which actual pin to wire, instead of
        /// creating the connection immediately. See HandleCreateAndDelete().
        bool         isMeta = false;
    };
    /// Resolved endpoint of a rendered link, keyed by its ax::NodeEditor LinkId.
    struct LinkInfo {
        pxr::SdfPath destNode;
        std::string  destInput;
    };

    /// Hypershade-style per-node pin display state. Always the render-time
    /// source of truth (never resync-clobbered); optionally mirrored into USD
    /// customData when m_keepGraphNodeViewSettingsInUsd is set. See
    /// GetPinDisplayOverride()/IsPinVisible().
    struct PinDisplayOverride {
        PinDisplayMode mode = PinDisplayMode::All;
    };
    /// A connection dragged onto/from a collapsed meta pin, awaiting the user
    /// picking which actual pin on nodePath/isOutput to wire up. otherPin is
    /// the already-resolved endpoint on the far side of the link (itself meta
    /// if the user dragged directly between two collapsed nubs, in which case
    /// picking chains into a second RenderPinPickMenu() for that side).
    struct PendingPinPick {
        pxr::SdfPath nodePath;
        bool         isOutput = false;
        PinInfo      otherPin;
    };

    void RenderToolbar();
    /// Property editor for the node selected in the graph, shown below the
    /// shader-ball preview. Values apply live while a widget is being
    /// dragged; one undoable command is pushed when the edit ends.
    void RenderSelectedNodeProperties();
    void RenderInputValueWidget(const ShaderGraphNode& node, const ShaderPinInfo& input,
                                const pxr::UsdPrim& prim);
    void CommitInputEdit(const ShaderGraphNode& node, const ShaderPinInfo& input);
    /// Hypershade-style browser: lists every material on the stage; the
    /// selected one is loaded into the graph work area via "Show Graph"
    /// (or double-click).
    void RenderMaterialBrowser();
    /// Creates a uniquely-named empty material under /Materials and opens it.
    void CreateNewMaterial();
    /// Preset graphs for the open material (one undoable command each):
    /// UsdPreviewSurface + UsdUVTextures fed by an st primvar reader, or a
    /// MaterialX standard_surface + image nodes fed by a texcoord node.
    void CreateUsdPresetGraph();
    void CreateMaterialXPresetGraph();
    void CreateArnoldPresetGraph();
    void RenderNodeGraphCanvas();
    void RenderPreviewPanel();
    void HandleCreateAndDelete();
    /// Nuke-style TAB popup: type-to-filter shader node list; Enter or click
    /// creates the highlighted node at canvasPos. When linkSource is non-null
    /// (a connection was dragged onto empty canvas), the list is filtered to
    /// nodes with a type-connectable pin and the chosen node is auto-wired to
    /// linkSource.
    void RenderNodeSearchMenu(const ImVec2& canvasPos, const PinInfo* linkSource = nullptr);
    /// Creates the node and returns its prim path (empty on failure).
    pxr::SdfPath CreateShaderNode(const std::string& shaderId, const ImVec2& canvasPos);
    /// Nudges a creation point right until the new node's estimated rectangle
    /// clears every existing node's rendered rectangle. Overlapping nodes
    /// fight over the editor's hit test: the buried one becomes unmovable and
    /// its pins unhoverable. Must be called with the editor context current.
    ImVec2 FindFreeCanvasSpot(ImVec2 desired, const std::string& shaderId) const;
    /// Binds an existing material to targetPath, undoably, restoring whatever
    /// direct binding (if any) targetPath had before.
    void BindMaterialToTarget(const pxr::SdfPath& materialPath, const pxr::SdfPath& targetPath);
    /// Creates a new material named after m_targetPrimPath under /Materials,
    /// binds it there, and opens it in the canvas.
    void CreateAndBindMaterialForTarget();
    void CreateConnection(const PinInfo& destInput, const PinInfo& sourceOutput);
    void DeleteNode(const pxr::SdfPath& nodePath);
    void DisconnectAttr(const pxr::SdfPath& destNode, const std::string& destInput);
    /// Resolved on-disk path of a texture node's image (its Asset-typed `file`
    /// input), or empty if unauthored/unresolvable. Feeds the thumbnail cache.
    std::string ResolveTextureFilePath(const ShaderGraphNode& node) const;
    void SyncFromUsd();
    void PersistNodePosition(ax::NodeEditor::NodeId nodeId);
    bool IsPinLinked(const pxr::SdfPath& nodePath, const std::string& pinName, bool isOutput) const;

    /// Seeds m_pinDisplayOverrides for any node path not already present
    /// (i.e. not touched yet this session) from the just-synced m_graph's
    /// pinDisplayMode (customData-backed). Never overwrites an existing
    /// entry, so in-session toggles survive resyncs.
    void SeedPinDisplayOverrides();
    /// Creates the entry (defaulted to All) on first access.
    PinDisplayOverride& GetPinDisplayOverride(const pxr::SdfPath& nodePath);
    /// True if pin should be drawn this frame: All mode shows everything;
    /// Connected mode shows only linked pins (linked pins are never hidden —
    /// Link() unconditionally references pin IDs that must have been drawn —
    /// so a hidden pin becomes visible the moment it's connected, via
    /// RenderPinPickMenu()); Collapsed mode hides every real pin, linked or
    /// not, in favor of the single connector nub per side (see
    /// EffectivePinId(), which redirects existing links to that nub).
    bool IsPinVisible(const ShaderGraphNode& node, const std::string& pinName, bool isOutput) const;
    /// Cycles a node's display mode: All -> Connected -> Collapsed -> All.
    void CyclePinDisplayMode(const pxr::SdfPath& nodePath);
    /// Sets every node currently in m_graph to mode (toolbar "set all" button).
    void ApplyPinDisplayModeToAllNodes(PinDisplayMode mode);
    /// Authors the node's current PinDisplayOverride into USD customData via
    /// an undoable command, mirroring PersistNodePosition. No-op unless
    /// m_keepGraphNodeViewSettingsInUsd is set.
    void PersistPinDisplayState(const pxr::SdfPath& nodePath);
    /// PinId to use when drawing a link touching (nodePath, pinName,
    /// isOutput): the pin's own id normally, or that node's collapsed meta
    /// nub id when the node is in Collapsed mode (so existing links visually
    /// converge on the nub instead of referencing an undrawn real pin).
    uintptr_t EffectivePinId(const pxr::SdfPath& nodePath, const std::string& pinName, bool isOutput) const;
    /// Renders the deferred "choose pin" popup for m_pendingPinPick (Maya
    /// Hypershade style): a connection was dragged onto/from a collapsed meta
    /// pin, and the user now picks which actual attribute it wires to. Must
    /// be called from within the NE::Suspend()/Resume() block, like the
    /// other popups.
    void RenderPinPickMenu();
    /// Finishes a pick: creates the real connection to/from chosen, unless
    /// the far side of the pending pick was itself a meta pin, in which case
    /// this chains into a second pick for that side instead.
    void CompletePinPick(const PinInfo& chosen);

    static bool SaveNodeSettingsCallback(ax::NodeEditor::NodeId nodeId,
                                         const char* data, size_t size,
                                         ax::NodeEditor::SaveReasonFlags reason,
                                         void* userPointer);

    ax::NodeEditor::EditorContext* m_editorContext = nullptr;
    pxr::UsdStageRefPtr m_stage;
    MaterialManager*    m_materialManager = nullptr;
    CommandHistory*     m_commandHistory  = nullptr;
    IconManager*        m_iconManager     = nullptr;

    /// Material currently shown in the graph work area (empty = none).
    pxr::SdfPath m_materialPath;
    /// Material highlighted in the browser list; becomes m_materialPath when
    /// the user clicks "Show Graph" (or double-clicks the entry).
    pxr::SdfPath m_browserSelection;
    /// Currently-selected prim, tracked for the "create + bind" / "bind"
    /// toolbar affordances; empty when nothing (or a non-prim path) is selected.
    pxr::SdfPath m_targetPrimPath;

    ShaderGraphSnapshot m_graph;
    /// NodeId (uintptr_t) -> prim path, for nodes already known to the editor
    /// this session (seeded position, so re-syncing won't fight live drags).
    std::unordered_map<uintptr_t, pxr::SdfPath> m_nodeIdToPath;
    /// Rebuilt every frame from the current m_graph while rendering.
    std::unordered_map<uintptr_t, PinInfo>  m_pinIdToInfo;
    std::unordered_map<uintptr_t, LinkInfo> m_linkIdToInfo;

    MaterialPreviewRenderer m_preview;
    /// Per-node in-canvas thumbnails (texture images + material shader balls).
    NodeThumbnailCache m_thumbnails;

    ImVec2 m_pendingCreateNodePos{0.0f, 0.0f};

    /// Alt+RMB drag distance not yet converted into a discrete wheel-zoom
    /// step (the node editor only zooms in wheel increments).
    float m_zoomDragAccum = 0.0f;

    char m_nodeSearchBuf[128] = "";
    int  m_nodeSearchSelected = 0;

    /// Nuke-style drag-to-create: set when a connection was dragged from a pin
    /// onto empty canvas. m_openLinkDragMenu requests the create-node popup on
    /// the next frame (deferred so OpenPopup runs while the editor is
    /// suspended); m_linkDragSourcePin is the pin to auto-wire the new node to.
    bool    m_openLinkDragMenu = false;
    PinInfo m_linkDragSourcePin;

    /// Hypershade-style pin display state, keyed by node path string. Always
    /// the render-time source of truth; see PinDisplayOverride.
    std::unordered_map<std::string, PinDisplayOverride> m_pinDisplayOverrides;
    /// Whether node toggles author USD customData (see
    /// SetKeepGraphNodeViewSettingsInUsd).
    bool m_keepGraphNodeViewSettingsInUsd = false;
    /// Deferred "choose pin" popup request: set from HandleCreateAndDelete()
    /// when a link is dropped on/dragged from a collapsed meta pin, actually
    /// opened after NE::Suspend(), like m_openLinkDragMenu.
    bool            m_openPinPickMenu = false;
    PendingPinPick  m_pendingPinPick;

    /// Toolbar toggle: whether node thumbnails (texture images / material
    /// shader-ball renders) are shown at all. Off skips the reserved space
    /// and the decode/render work entirely, not just the visual.
    bool m_showNodeThumbnails = true;
    /// Toolbar "set all nodes" pin-display button: the mode last broadcast
    /// to every node (and the icon/tooltip shown), independent of what any
    /// individual node is actually in after the user tweaks one by hand.
    PinDisplayMode m_toolbarPinDisplayMode = PinDisplayMode::All;

    /// Column widths of the browser (left) and preview (right) sections,
    /// user-adjustable via the splitters between them and the canvas.
    float m_browserWidth = 220.0f;
    float m_previewWidth = 320.0f;

    /// Fraction of the browser column's height given to the materials list;
    /// the create-node list takes the rest. User-adjustable via the horizontal
    /// splitter between them.
    float m_browserListRatio = 0.45f;

    /// Inline-rename state for the materials list (double-click a row to start;
    /// mirrors SceneHierarchyPanel): the material being renamed (empty = none),
    /// the edit buffer, and a one-frame flag to focus the InputText on start.
    pxr::SdfPath m_renamingMaterial;
    char         m_materialRenameBuf[128] = "";
    bool         m_materialRenameJustStarted = false;

    /// Search text of the browser's persistent create-node list.
    char m_browserNodeSearchBuf[128] = "";
    /// Shader id clicked in the browser's create-node list; created at the
    /// canvas view center on the next RenderNodeGraphCanvas (ScreenToCanvas
    /// is only valid inside the editor's Begin/End).
    std::string m_pendingCreateShaderId;

    /// Debug event-trace state (see the logging block at the end of
    /// RenderNodeGraphCanvas): whether the last LMB press landed on a node
    /// rect, whether that press's drag-start has been logged yet, and the
    /// editor NodeIds hit at press time (their live positions are logged on
    /// every drag-move so a node that stops following the mouse is visible).
    bool m_debugPressOnNode = false;
    bool m_debugDragLogged  = false;
    std::vector<uintptr_t> m_debugDragNodeIds;

    /// First node currently selected in the graph editor (empty = none);
    /// drives the properties section under the preview.
    pxr::SdfPath m_selectedNodePath;
    /// Value/authored-state of the input being edited, captured when its
    /// widget activates, so the commit command can restore it on undo. Only
    /// one ImGui widget can be active at a time, so one slot suffices.
    pxr::VtValue m_preEditValue;
    bool         m_preEditWasAuthored = false;
};

} // namespace UsdLayerManager