UsdLayerManager/src/ui/MaterialEditorPanel.h

#pragma once

#include "../core/CommandHistory.h"
#include "../core/MaterialManager.h"
#include "IconManager.h"
#include "MaterialPreviewRenderer.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>

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

    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;
        pxr::SdfValueTypeName typeName;
    };
    /// Resolved endpoint of a rendered link, keyed by its ax::NodeEditor LinkId.
    struct LinkInfo {
        pxr::SdfPath destNode;
        std::string  destInput;
    };

    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();
    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.
    void RenderNodeSearchMenu(const ImVec2& canvasPos);
    void CreateShaderNode(const std::string& shaderId, const ImVec2& canvasPos);
    /// 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);
    void SyncFromUsd();
    void PersistNodePosition(ax::NodeEditor::NodeId nodeId);
    bool IsPinLinked(const pxr::SdfPath& nodePath, const std::string& pinName, bool isOutput) const;

    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;

    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;

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

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