Editor/AutoRig/RigPreviewDialog.cs

Editor modal dialog for previewing and editing a rig. Shows a rendered model with selectable/draggable joints and a joint list, supports rename/delete, undo for rig edits and live pose, and reports selection info.

using System;
using System.Collections.Generic;
using AutoRig.Rig;
using Editor;
using Sandbox;

using VecN = System.Numerics.Vector3;

namespace AutoRig.Editor;

/// <summary>
/// Modal rig preview + realtime joint editor, styled after humanoid-retargeter's
/// PreviewDialog. The left pane is <see cref="RigPreviewWidget"/> (translucent
/// model, skeleton shown through it, click to select, drag a joint to see the
/// mesh follow it, scroll to zoom). The right pane is <see cref="JointListPanel"/>
/// at the same height. Right-clicking a joint (in either pane) opens a menu to
/// rename or delete it; edits write back to the entry so Export uses them.
/// </summary>
public sealed class RigPreviewDialog : Dialog
{
    readonly ModelEntry _entry;
    readonly Action _onChanged;
    readonly RigPreviewWidget _preview;
    readonly JointListPanel _jointList;
    readonly Label _summary;
    readonly string _baseline;

    // Undo: each entry restores the rig AND the live joint pose to a prior state.
    readonly Stack<(RigResult Rig, VecN[] Pose, string Selected)> _undo = new();
    Button _undoButton;

    public RigPreviewDialog( Widget parent, ModelEntry entry, Action onChanged )
        : base( parent )
    {
        _entry = entry;
        _onChanged = onChanged;
        var rig = entry.Rig;

        Window.WindowTitle = $"Preview - {entry.FileName}";
        Window.SetWindowIcon( "preview" );
        Window.SetModal( true, true );
        Window.MinimumWidth = 820;
        Window.MinimumHeight = 560;

        Layout = Layout.Column();
        Layout.Margin = 12;
        Layout.Spacing = 8;

        // ---- preview + joint list (one stretching row → both are equal height) ----
        var split = Layout.AddRow( 1 );
        split.Spacing = 8;

        _preview = new RigPreviewWidget( this );
        _preview.SetRig( entry.Mesh, rig );
        split.Add( _preview, 1 );

        _jointList = new JointListPanel( this, rig );
        split.Add( _jointList, 0 );

        // ---- info row (transport-bar idiom) ----
        var info = Layout.AddRow();
        info.Spacing = 8;
        _baseline = $"{rig.Skeleton.Joints.Count} joints · {rig.SolverName} · "
            + "click to select · drag a joint to test · right-click to edit · scroll to zoom";
        _summary = new Label( _baseline, this );
        _summary.SetStyles( $"color: {Theme.TextLight.Hex};" );
        info.Add( _summary );
        info.AddStretchCell();

        // Two-way selection sync + shared context menu.
        _preview.JointPicked = jointName =>
        {
            _jointList.Select( jointName ?? "" );
            ShowSelection( jointName );
        };
        _preview.JointContextRequested = ShowJointMenu;
        // A joint drag snapshots BEFORE it moves, so one Ctrl+Z reverts the move.
        _preview.BeforeJointMove = PushUndo;
        _preview.UndoRequested = Undo;
        _jointList.JointSelected = jointName =>
        {
            _preview.SetHighlight( jointName );
            ShowSelection( jointName );
        };
        _jointList.JointContext = ShowJointMenu;

        // ---- confirm row ----
        var confirm = Layout.AddRow();
        confirm.Spacing = 8;
        _undoButton = new Button( "Undo", "undo", this ) { Enabled = false };
        _undoButton.ToolTip = "Undo the last move / delete / rename (Ctrl+Z).";
        _undoButton.Clicked = Undo;
        confirm.Add( _undoButton );
        confirm.AddStretchCell();
        var ok = confirm.Add( new Button.Primary( "Close" ) { Icon = "check" } );
        ok.Tint = Theme.Green;
        ok.Clicked = Close;

        Window.Size = new Vector2( 980, 720 );
    }

    void PushUndo()
    {
        _undo.Push( (_entry.Rig, _preview.GetPose(), _preview.SelectedJoint ) );
        if ( _undoButton is not null )
            _undoButton.Enabled = true;
    }

    void Undo()
    {
        if ( _undo.Count == 0 )
            return;
        var (rig, pose, selected) = _undo.Pop();
        var rigChanged = !ReferenceEquals( rig, _entry.Rig );
        _entry.Rig = rig;
        if ( rigChanged )
        {
            _preview.Reload( rig );
            _jointList.Rebuild( rig );
            _onChanged?.Invoke();
        }
        _preview.SetPose( pose );   // restores the joint positions
        _preview.SetHighlight( selected ?? "" );
        _jointList.Select( selected ?? "" );
        ShowSelection( string.IsNullOrEmpty( selected ) ? null : selected );
        if ( _undoButton is not null )
            _undoButton.Enabled = _undo.Count > 0;
    }

    protected override void OnKeyPress( KeyEvent e )
    {
        if ( e.Key == KeyCode.Z && e.HasCtrl )
        {
            Undo();
            return;
        }
        base.OnKeyPress( e );
    }

    void ShowSelection( string jointName )
    {
        if ( jointName is null )
        {
            _summary.Text = _baseline;
            return;
        }
        var rig = _entry.Rig;
        var index = rig.Skeleton.IndexOf( jointName );
        if ( index < 0 )
        {
            _summary.Text = _baseline;
            return;
        }
        var joint = rig.Skeleton.Joints[index];
        var parent = joint.Parent >= 0 ? rig.Skeleton.Joints[joint.Parent].Name : "none";
        _summary.Text = $"{jointName} · parent: {parent} · "
            + $"({joint.Position.X:0.##}, {joint.Position.Y:0.##}, {joint.Position.Z:0.##})";
    }

    void ShowJointMenu( string jointName )
    {
        var index = _entry.Rig.Skeleton.IndexOf( jointName );
        if ( index < 0 )
            return;
        var isRoot = _entry.Rig.Skeleton.Joints[index].Parent < 0;

        var menu = new Menu();
        menu.AddOption( "Rename…", "edit", () => RenameJoint( index ) );
        var delete = menu.AddOption( "Delete joint", "delete", () => DeleteJoint( index ) );
        delete.Enabled = !isRoot;   // the root anchors the skeleton
        if ( isRoot )
            delete.StatusTip = "The root joint cannot be deleted.";
        menu.OpenAtCursor();
    }

    void RenameJoint( int index )
    {
        var current = _entry.Rig.Skeleton.Joints[index].Name;
        var prompt = new RenamePrompt( this, current, name =>
        {
            var edited = JointEdits.RenameJoint( _entry.Rig, index, name );
            var newName = edited.Skeleton.Joints[index].Name;
            if ( newName == current )
                return;   // blank or collision - nothing changed
            PushUndo();
            ApplyEdit( edited, newName );
        } );
        prompt.Window.Show();
    }

    void DeleteJoint( int index )
    {
        RigResult edited;
        try
        {
            edited = JointEdits.DeleteJoint( _entry.Rig, index );
            edited.Weights.Validate( _entry.Mesh, edited.Skeleton );
        }
        catch ( Exception )
        {
            return;   // malformed edit - leave the rig untouched
        }
        if ( ReferenceEquals( edited, _entry.Rig ) )
            return;   // nothing removed (e.g. root)
        PushUndo();
        ApplyEdit( edited, null );
    }

    /// <summary>Swaps the edited rig onto the entry, rebuilds both panes, keeps
    /// the given joint selected, and refreshes the window row (so Export uses it).</summary>
    void ApplyEdit( RigResult rig, string keepSelected )
    {
        _entry.Rig = rig;
        _preview.Reload( rig );
        _jointList.Rebuild( rig );
        if ( keepSelected is not null && rig.Skeleton.IndexOf( keepSelected ) >= 0 )
        {
            _preview.SetHighlight( keepSelected );
            _jointList.Select( keepSelected );
            ShowSelection( keepSelected );
        }
        else
        {
            _preview.SetHighlight( "" );
            _jointList.Select( "" );
            _summary.Text = $"{rig.Skeleton.Joints.Count} joints · {rig.SolverName} · "
                + "click to select · drag a joint to test · right-click to edit · scroll to zoom";
        }
        _onChanged?.Invoke();
    }

    /// <summary>Tiny modal name prompt: a text field + OK/Cancel.</summary>
    sealed class RenamePrompt : Dialog
    {
        public RenamePrompt( Widget parent, string current, Action<string> onAccept )
            : base( parent )
        {
            Window.WindowTitle = "Rename joint";
            Window.SetWindowIcon( "edit" );
            Window.SetModal( true, true );
            Window.MinimumWidth = 320;

            Layout = Layout.Column();
            Layout.Margin = 12;
            Layout.Spacing = 8;

            var edit = new LineEdit( this ) { Text = current };
            edit.SelectAll();
            Layout.Add( edit );

            var buttons = Layout.AddRow();
            buttons.Spacing = 8;
            buttons.AddStretchCell();
            var cancel = buttons.Add( new Button( "Cancel" ) );
            cancel.Clicked = Close;
            var ok = buttons.Add( new Button.Primary( "Rename" ) { Icon = "check" } );
            ok.Clicked = () =>
            {
                onAccept?.Invoke( edit.Text );
                Close();
            };

            Window.Size = new Vector2( 340, 110 );
        }
    }
}