Editor/AutoRig/AutoRigWindow.cs

Editor dock window UI for the Auto Rigger. Provides a three-step wizard (Add Model → Review/Rig → Export) that lets the user add model files, run analysis and rigging (local or cloud/offload), preview rigs, and export skinned FBX + compiled vmdl into the project.

File AccessNetworking
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using System.IO;
using AutoRig.Analyze;
using AutoRig.Dl;
using AutoRig.Formats;
using AutoRig.Rig;
using AutoRig.Solve;
using Editor;
using Sandbox;

namespace AutoRig.Editor;

/// <summary>
/// The Auto Rig dock window: a three-step wizard (Add Model → Review Rig → Export)
/// that turns any static model into a rigged, compiled vmdl. Spacing follows
/// <see cref="UiSpacing"/> throughout (spec §3); chip/status conventions match
/// humanoid-retargeter (green = good, amber = review, red = problem, blue = working).
/// </summary>
[Dock( "Editor", "Auto Rigger (DL)", "precision_manufacturing" )]
public sealed class AutoRigWindow : Widget
{
    static AutoRigWindow _instance;

    /// <summary>The open window instance, if any (dock windows are singletons here).</summary>
    public static AutoRigWindow Instance => _instance.IsValid() ? _instance : null;

    readonly List<ModelEntry> _entries = new();

    Layout _stepHeader;
    Layout _body;
    Layout _listLayout;
    Label _statusLabel;
    Label _explanationLabel;
    ComboBox _modeCombo;
    ComboBox _modelCombo;
    Button _rigAllButton;
    Checkbox _cleanupCheck;
    Button _donorButton;

    /// <summary>Whether the post-rig cleanup phase runs (on by default).</summary>
    public bool CleanupEnabled => _cleanupCheck?.Value ?? true;

    /// <summary>Runs the cleanup phase on a fresh rig when enabled. The preview
    /// stays visually identical - it only repairs the skeleton graph, merges
    /// zero-length bones, and sanitizes weights.</summary>
    RigResult MaybeClean( ModelEntry entry, RigResult rig )
        => CleanupEnabled && entry.Mesh is not null
            ? AutoRig.Rig.SkeletonCleanup.Clean( entry.Mesh, rig )
            : rig;
    Label _donorLabel;
    DonorRig _donor;

    const string AutoModelItem = "Auto - best enabled model";

    /// <summary>Dropdown = Auto + every enabled (installed, hardware-capable) model.</summary>
    void RebuildModelCombo()
    {
        if ( _modelCombo is null )
            return;
        var previous = _modelCombo.CurrentText;
        _modelCombo.Clear();
        _modelCombo.AddItem( AutoModelItem, "auto_awesome" );
        foreach ( var choice in DlModelRegistry.EnabledChoices() )
            _modelCombo.AddItem( choice.Title, "model_training" );
        if ( !string.IsNullOrEmpty( previous ) )
            _modelCombo.TrySelectNamed( previous );
    }
    RigPreviewWidget _preview;
    Button _wiggleButton;
    JointAdjustPanel _adjustPanel;
    ModelEntry _previewEntry;
    LineEdit _outputFolderEdit;

    /// <summary>The rig mode currently selected in the toolbar.</summary>
    public RigMode SelectedMode
        => Enum.TryParse<RigMode>( _modeCombo?.CurrentText, out var mode ) ? mode : RigMode.Auto;

    /// <summary>Dock constructor (called by the editor's dock manager).</summary>
    public AutoRigWindow( Widget parent ) : base( parent )
    {
        _instance ??= this;

        Name = "AutoRig";
        WindowTitle = "Auto Rig";
        SetWindowIcon( "precision_manufacturing" );
        MinimumSize = new Vector2( 780, 480 );

        Layout = Layout.Column();
        Layout.Margin = UiSpacing.Section;
        Layout.Spacing = UiSpacing.Group;

        BuildUi();
    }

    /// <summary>Opens (or raises) the window and returns it.</summary>
    public static AutoRigWindow Open()
    {
        var window = Instance ?? EditorWindow.DockManager.Create<AutoRigWindow>();
        EditorWindow.DockManager.RaiseDock( window );
        return window;
    }

    void BuildUi()
    {
        // ---- title header (donor idiom: large icon + name) ----
        var titleRow = Layout.AddRow();
        titleRow.Spacing = UiSpacing.Related;
        var titleIcon = new Label( "precision_manufacturing", this );
        titleIcon.SetStyles(
            $"font-family: Material Icons; font-size: 30px; color: {Theme.Blue.Hex};" );
        titleRow.Add( titleIcon );
        var titleText = new Label( "Auto Rigger (DL)", this );
        titleText.SetStyles( "font-size: 18px; font-weight: 600;" );
        titleRow.Add( titleText );
        var titleHint = new Label( "any model → animation-ready vmdl", this );
        titleHint.SetStyles( $"color: {Theme.TextLight.Hex}; margin-left: {UiSpacing.Group}px;" );
        titleRow.Add( titleHint );
        titleRow.AddStretchCell();

        // IconButton centers its glyph; a text Button reserves label space even
        // when empty, which left extra padding to the right of the cog.
        var settingsButton = new IconButton( "settings" )
        {
            ToolTip = "Settings - cloud rig server port, docker image, timeouts.",
            FixedSize = new Vector2( 28, 28 ),
            OnClick = () =>
            {
                var dialog = new SettingsDialog( this );
                dialog.Show();
            },
        };
        titleRow.Add( settingsButton );

        SettingsDialog.LoadAndApply();   // persisted settings take effect at open

        // ---- wizard step strip (progress indicator - all steps live on this page) ----
        _stepHeader = Layout.AddRow();
        _stepHeader.Spacing = UiSpacing.Group;
        _stepChips[0] = AddStepChip( 1, "Add Model" );
        _stepChips[1] = AddStepChip( 2, "Export" );
        _stepHeader.AddStretchCell();

        // ---- body: step 1 (Add Model) ----
        _body = Layout.AddColumn( 1 );
        _body.Margin = new Sandbox.UI.Margin( 0, UiSpacing.Section, 0, 0 );
        _body.Spacing = UiSpacing.Group;

        var addRow = _body.AddRow();
        addRow.Spacing = UiSpacing.Related;
        var addButton = new Button( "Add Files…", "add", this );
        addButton.Clicked = AddFiles;
        addRow.Add( addButton );

        _rigAllButton = new Button( "Rig All", "engineering", this );
        _rigAllButton.ToolTip = "Rig every analyzed model in the list, one after "
            + "another, with the selected model.";
        _rigAllButton.Clicked = () => _ = RigAll();
        addRow.Add( _rigAllButton );

        // DL-only pivot (2026-07-02): every rig runs the neural pipeline. Model
        // dropdown auto-picks the best ENABLED model per mesh; Manage opens the
        // enable/disable + download list. Geometric solvers stay in Code/ unused.
        var modelLabel = new Label( "Model", this );
        modelLabel.SetStyles( $"color: {Theme.TextLight.Hex}; margin-left: {UiSpacing.Group}px;" );
        addRow.Add( modelLabel );
        _modelCombo = new ComboBox( this );
        _modelCombo.MinimumWidth = 220;
        RebuildModelCombo();
        addRow.Add( _modelCombo );

        var manageButton = new Button( "Manage Models…", "model_training", this );
        manageButton.ToolTip = "Download deep-learning models and choose which are "
            + "enabled. Models this machine cannot run are locked off automatically.";
        manageButton.Clicked = () =>
        {
            var dialog = new CatalogDialog( this,
                _entries.FirstOrDefault( e => e.Mesh is not null )?.Mesh.Positions.Length ?? 0 );
            dialog.Window.Show();
        };
        addRow.Add( manageButton );

        var cloudButton = new Button( "Cloud… (experimental)", "cloud", this );
        cloudButton.ToolTip = "EXPERIMENTAL. Rig on a rented vast.ai GPU when this "
            + "machine can't run a model. Renting always asks first and shows the "
            + "cost; the instance is destroyed automatically afterwards (only ours). "
            + "Remote pipelines are best-effort and some models may not work yet.";
        cloudButton.Clicked = OpenCloudDialog;
        addRow.Add( cloudButton );

        _cleanupCheck = new Checkbox( "Cleanup phase", this ) { Value = true };
        _cleanupCheck.ToolTip = "After rigging, run a foolproof cleanup: repair the "
            + "skeleton (single root, valid hierarchy, merge zero-length bones), "
            + "sanitize skin weights, and orient bones along their limbs so the "
            + "exported model imports clean in any DCC tool. The preview looks "
            + "practically identical - it only tidies the rig, it never moves joints.";
        addRow.Add( _cleanupCheck );
        addRow.AddStretchCell();

        DlModelRegistry.Changed = () =>
        {
            RebuildModelCombo();
            RebuildList();
        };

        var split = _body.AddRow( 1 );
        split.Spacing = UiSpacing.Section;

        var left = split.AddColumn( 1 );
        left.Spacing = UiSpacing.RowGap;
        _listLayout = left.AddColumn();
        _listLayout.Spacing = UiSpacing.RowGap;
        left.AddStretchCell();

        // Preview viewport, wiggle test and adjust panel removed from the UI on
        // user request (2026-07-02) - the window is Add → Export. The preview
        // machinery stays in the codebase, dormant.

        // ---- step 2: export ----
        var exportHeader = new Label( "2. Export", this );
        exportHeader.SetStyles(
            $"font-weight: 600; color: {Theme.Blue.Hex}; "
            + $"margin-top: {UiSpacing.HeaderTop}px; margin-bottom: {UiSpacing.HeaderBottom}px;" );
        _body.Add( exportHeader );
        var exportRow = _body.AddRow();
        exportRow.Spacing = UiSpacing.Related;
        var folderLabel = new Label( "Output folder", this );
        folderLabel.SetStyles( $"color: {Theme.TextLight.Hex};" );
        exportRow.Add( folderLabel );
        _outputFolderEdit = new LineEdit( this ) { Text = "models/autorig" };
        exportRow.Add( _outputFolderEdit, 1 );
        var exportButton = new Button( "Export All", "save", this );
        exportButton.Clicked = () => _ = ExportAll();
        exportRow.Add( exportButton );

        RebuildList();

        // ---- status bar ----
        var statusRow = Layout.AddRow();
        statusRow.Margin = new Sandbox.UI.Margin( 0, UiSpacing.Group, 0, 0 );
        _statusLabel = new Label( "Ready.", this );
        _statusLabel.SetStyles( $"color: {Theme.TextLight.Hex};" );
        statusRow.Add( _statusLabel );
        statusRow.AddStretchCell();
    }

    readonly Label[] _stepChips = new Label[2];

    Label AddStepChip( int number, string title )
    {
        var label = new Label( $"{number}. {title}", this );
        label.SetStyles(
            $"color: {Theme.TextLight.Hex}; padding: {UiSpacing.Related}px {UiSpacing.Group}px;" );
        _stepHeader.Add( label );
        return label;
    }

    /// <summary>Highlights how far the user has gotten (everything lives on one page).</summary>
    void UpdateSteps()
    {
        if ( _stepChips[0] is null )
            return;
        var reached = _entries.Any( e => e.Status == EntryStatus.Exported ) ? 2 : 1;
        for ( var i = 0; i < _stepChips.Length; i++ )
        {
            var active = i + 1 <= reached;
            var current = i + 1 == reached;
            _stepChips[i].SetStyles(
                $"color: {(active ? Theme.Blue : Theme.TextLight).Hex}; "
                + $"font-weight: {(current ? 600 : 400)}; "
                + (current ? $"background-color: {Theme.ControlBackground.Hex}; " : "")
                + $"border-radius: 10px; padding: {UiSpacing.Related}px {UiSpacing.Group}px;" );
        }
    }

    void SetStatus( string text, Color color )
    {
        if ( _statusLabel is null )
            return;
        _statusLabel.Text = text;
        _statusLabel.SetStyles( $"color: {color.Hex};" );
    }

    // ================================================================== step 1

    /// <summary>Loads a rigged donor model for Transfer mode.</summary>
    void PickDonor()
    {
        var path = EditorUtility.OpenFileDialog(
            "Select rigged donor model", "Rigged Models (*.glb *.gltf *.fbx)", null );
        if ( string.IsNullOrEmpty( path ) )
            return;
        try
        {
            _donor = DonorLoader.Load( File.ReadAllBytes( path ), Path.GetFileName( path ) );
            _donorLabel.Text =
                $"{Path.GetFileName( path )} ({_donor.Skeleton.Joints.Count} joints)";
            SetStatus( $"Donor ready: {Path.GetFileName( path )} - pick Transfer mode and Rig.",
                Theme.Green );
        }
        catch ( Exception e )
        {
            _donor = null;
            _donorLabel.Text = "";
            SetStatus( FirstLine( e.Message ), Theme.Red );
        }
    }

    /// <summary>Textures referenced (not embedded) by the source file: try the
    /// recorded relative path, then just its filename, beside the model.</summary>
    static void ResolveSiblingTextures( AutoRig.Mesh.RigMesh mesh, string modelPath )
    {
        var directory = Path.GetDirectoryName( modelPath );
        if ( string.IsNullOrEmpty( directory ) )
            return;
        foreach ( var material in mesh.Materials )
        {
            if ( material.BaseColorImage is not null
                || string.IsNullOrEmpty( material.SourceImagePath ) )
                continue;
            var relative = material.SourceImagePath.Replace( '\\', Path.DirectorySeparatorChar );
            foreach ( var candidate in new[]
            {
                Path.Combine( directory, relative ),
                Path.Combine( directory, Path.GetFileName( relative ) ),
            } )
            {
                try
                {
                    if ( File.Exists( candidate ) )
                    {
                        material.BaseColorImage = File.ReadAllBytes( candidate );
                        break;
                    }
                }
                catch ( Exception )
                {
                    // unreadable candidate - try the next, else stay untextured
                }
            }
        }
    }

    void AddFiles()
    {
        var path = EditorUtility.OpenFileDialog(
            "Select model", "Model Files (*.fbx *.glb *.gltf *.obj)", null );
        if ( string.IsNullOrEmpty( path ) )
            return;
        AddFile( path );
    }

    /// <summary>Adds one file and analyzes it in the background.</summary>
    public async void AddFile( string path )
    {
        var entry = new ModelEntry
        {
            Path = path,
            FileName = System.IO.Path.GetFileName( path ),
        };
        _entries.Add( entry );
        RebuildList();
        SetStatus( $"Analyzing {entry.FileName}…", Theme.Blue );

        try
        {
            var bytes = await Task.Run( () => System.IO.File.ReadAllBytes( path ) );
            var (mesh, analysis) = await Task.Run( () =>
            {
                var loaded = MeshLoader.Load( bytes, entry.FileName );
                return (loaded, MeshAnalyzer.Analyze( loaded ));
            } );
            ResolveSiblingTextures( mesh, path );
            entry.Mesh = mesh;
            entry.Analysis = analysis;
            entry.Status = EntryStatus.Ready;
            entry.Detail = analysis.Classification.Explanation;
            ShowInPreview( entry );
            // Adding a model only analyzes it now - the user drives rigging with
            // the row's "Rig" button (per model) or "Rig All" (the whole list).
            SetStatus( $"{entry.FileName}: {analysis.Summary} - press Rig (or Rig All).",
                Theme.Green );
        }
        catch ( Exception e )
        {
            entry.Status = EntryStatus.Failed;
            entry.Detail = FirstLine( e.Message );
            SetStatus( $"{entry.FileName}: {entry.Detail}", Theme.Red );
        }
        RebuildList();
    }

    void RebuildList()
    {
        if ( _listLayout is null )
            return;
        UpdateSteps();
        _listLayout.Clear( true );

        if ( _rigAllButton is not null )
            _rigAllButton.Enabled = _entries.Any( e => e.Analysis is not null
                && e.Status is not (EntryStatus.Analyzing or EntryStatus.Rigging) );

        if ( _entries.Count == 0 )
        {
            var empty = new Label(
                "Add a model to get started - any static mesh: characters, animals, vehicles, weapons, props.",
                this );
            empty.SetStyles( $"color: {Theme.TextLight.Hex}; margin: {UiSpacing.Group}px;" );
            _listLayout.Add( empty );
            return;
        }

        foreach ( var entry in _entries )
        {
            var captured = entry;
            // Card per model (donor idiom): recessed rounded background.
            var card = new Widget( this );
            card.SetStyles(
                $"background-color: {Theme.ControlBackground.Hex}; border-radius: 6px;" );
            card.Layout = Layout.Row();
            var row = card.Layout;
            row.Spacing = UiSpacing.Related;
            row.Margin = new Sandbox.UI.Margin(
                UiSpacing.Group, UiSpacing.Related, UiSpacing.Group, UiSpacing.Related );
            _listLayout.Add( card );

            var (icon, tone) = entry.Status switch
            {
                EntryStatus.Analyzing => ("sync", Theme.Blue),
                EntryStatus.Ready => ("check_circle", Theme.Green),
                EntryStatus.Rigging => ("sync", Theme.Blue),
                EntryStatus.Rigged => ("engineering", Theme.Green),
                EntryStatus.Exported => ("task_alt", Theme.Green),
                _ => ("error", Theme.Red),
            };
            var iconLabel = new Label( icon, this );
            iconLabel.SetStyles( $"font-family: Material Icons; font-size: 16px; color: {tone.Hex};" );
            row.Add( iconLabel );

            var name = new Label( entry.FileName, this );
            name.MouseClick = () => ShowInPreview( captured );
            row.Add( name );

            var chip = new Label( entry.ChipText, this );
            chip.ToolTip = entry.Detail;
            chip.SetStyles( $"color: {tone.Hex}; padding: 0px {UiSpacing.Related}px;" );
            row.Add( chip );
            row.AddStretchCell();

            if ( entry.Status is EntryStatus.Ready or EntryStatus.Rigged )
            {
                var rigButton = new Button( entry.Status == EntryStatus.Rigged ? "Re-rig" : "Rig", "engineering", this );
                rigButton.SetStyles( $"padding: 2px {UiSpacing.Group}px; margin-right: {UiSpacing.Related}px;" );
                rigButton.Clicked = () => _ = RigEntry( captured, SelectedMode );
                row.Add( rigButton );

                // Offload this model's rig onto a GPU the user already has running.
                var offloadButton = new Button( "Offload…", "dns", this );
                offloadButton.ToolTip = "Run this rig on one of your already-rented "
                    + "vast.ai GPUs (experimental). Opens a picker of running instances.";
                offloadButton.SetStyles( $"padding: 2px {UiSpacing.Group}px; margin-right: {UiSpacing.Related}px;" );
                offloadButton.Clicked = () => OpenOffloadPicker( captured );
                row.Add( offloadButton );
            }

            // Preview: disabled until this entry's rigging has FINISHED (user
            // requirement) - enabled only with a rig result in hand.
            var previewButton = new IconButton( "visibility" )
            {
                ToolTip = entry.Rig is not null
                    ? "Preview the rig - translucent model, skeleton with joints, scroll to zoom"
                    : "Rig first - preview unlocks when rigging has finished",
                FixedSize = new Vector2( 24, 24 ),
                Enabled = entry.Rig is not null
                    && entry.Status is EntryStatus.Rigged or EntryStatus.Exported,
                OnClick = () =>
                {
                    if ( captured.Rig is null )
                        return;
                    // The dialog edits captured.Rig in place (delete/rename/move);
                    // RebuildList refreshes the row so Export uses the edited rig.
                    var dialog = new RigPreviewDialog( this, captured, RebuildList );
                    dialog.Window.Show();
                },
            };
            row.Add( previewButton );

            // IconButton centers its glyph (a text Button reserves label space,
            // which pushed the X left no matter the padding).
            var remove = new IconButton( "close" )
            {
                ToolTip = "Remove from the list",
                FixedSize = new Vector2( 24, 24 ),
                OnClick = () =>
                {
                    _entries.Remove( captured );
                    RebuildList();
                },
            };
            row.Add( remove );
        }
    }

    static string FirstLine( string text )
    {
        var i = text.IndexOfAny( new[] { '\r', '\n' } );
        return i < 0 ? text : text[..i];
    }

    // ================================================================== step 2

    /// <summary>Rigs every analyzed entry in turn (sequential - neural inference
    /// is heavy, so we never run two at once).</summary>
    public async Task RigAll()
    {
        foreach ( var entry in _entries.ToList() )
        {
            if ( entry.Analysis is null )
                continue;   // still analyzing, or failed to load
            if ( entry.Status is EntryStatus.Analyzing or EntryStatus.Rigging )
                continue;
            await RigEntry( entry, SelectedMode );
        }
    }

    /// <summary>Generates (or regenerates) the rig for an entry in the background.</summary>
    public async Task RigEntry( ModelEntry entry, RigMode mode )
    {
        if ( entry.Analysis is null || entry.Status is EntryStatus.Analyzing or EntryStatus.Failed )
            return;

        entry.StateToken++;
        entry.Status = EntryStatus.Rigging;
        RebuildList();
        SetStatus( $"Rigging {entry.FileName} ({mode})…", Theme.Blue );

        try
        {
            // DL-only: dropdown picks the model (Auto = best enabled for THIS mesh).
            // Nothing usable = a clear failure with instructions; solver errors
            // surface as-is - never a silent geometric fallback.
            var selected = _modelCombo?.CurrentText;
            var choice = selected is null || selected == AutoModelItem
                ? DlModelRegistry.SelectBest( entry.Mesh?.Positions.Length ?? 1 )
                : DlModelRegistry.EnabledChoices().FirstOrDefault( c => c.Title == selected )
                    ?? DlModelRegistry.SelectBest( entry.Mesh?.Positions.Length ?? 1 );
            if ( choice is null )
            {
                const string msg = "No usable deep-learning model - open Manage Models… "
                    + "and download/enable one (or load your own checkpoint).";
                SetStatus( msg, Theme.Yellow );
                FlashError( entry, msg );
                return;
            }
            SetStatus( $"Rigging {entry.FileName} with {choice.Title}…", Theme.Blue );
            var rig = await Task.Run( () =>
            {
                var loaded = DlModelRegistry.LoadRunnableModel( choice );
                return loaded switch
                {
                    RigNetBundle bundle => DeepLearningSolver.Rig( entry.Analysis, bundle ),
                    AutoRig.Dl.UniRig.UniRigSkeletonModel unirig
                        => UniRigSolver.Rig( entry.Analysis, unirig ),
                    AutoRig.Dl.UniRig.SkinTokensModel skinTokens
                        => SkinTokensSolver.Rig( entry.Analysis, skinTokens ),
                    AutoRig.Dl.MagicArticulate.MagicArticulateModel magic
                        => MagicArticulateSolver.Rig( entry.Analysis, magic ),
                    AutoRig.Dl.Puppeteer.PuppeteerModel puppeteer
                        => PuppeteerSolver.Rig( entry.Analysis, puppeteer ),
                    AutoRig.Dl.RigAnything.RigAnythingModel rigAnything
                        => RigAnythingSolver.Rig( entry.Analysis, rigAnything ),
                    AutoRig.Dl.Anymate.AnymateModel anymate
                        => AnymateSolver.Rig( entry.Analysis, anymate ),
                    _ => throw new FormatException( "Unknown model type." ),
                };
            } );
            rig = MaybeClean( entry, rig );
            entry.Rig = rig;
            entry.Status = EntryStatus.Rigged;
            entry.Detail = rig.Explanation;
            entry.StateToken++;
            ShowInPreview( entry );
            SetStatus(
                $"{entry.FileName}: {rig.Skeleton.Joints.Count} joints via {rig.SolverName}"
                + (rig.Degraded ? " (degraded)" : ""),
                rig.Degraded ? Theme.Yellow : Theme.Green );
        }
        catch ( Exception e )
        {
            SetStatus( $"{entry.FileName}: {FirstLine( e.Message )}", Theme.Red );
            FlashError( entry, FirstLine( e.Message ) );   // red for 4s, then restores the row
        }
        RebuildList();
    }

    /// <summary>The cloud model for an entry: the toolbar dropdown pick, or Auto
    /// → best enabled model for that mesh. Null when nothing is enabled.</summary>
    CatalogEntry ResolveCloudModel( ModelEntry entry )
    {
        var selected = _modelCombo?.CurrentText;
        var choice = selected is null || selected == AutoModelItem
            ? DlModelRegistry.SelectBest( entry.Mesh?.Positions.Length ?? 1 )
            : DlModelRegistry.EnabledChoices().FirstOrDefault( c => c.Title == selected )
                ?? DlModelRegistry.SelectBest( entry.Mesh?.Positions.Length ?? 1 );
        return choice?.Entry;
    }

    /// <summary>Runs a cloud rig (rent or offload) for an entry and applies the
    /// result, driving the row status. Rethrows so the dialog shows it too.</summary>
    async Task RunCloud( ModelEntry entry, Func<Task<RigResult>> run )
    {
        entry.StateToken++;
        entry.Status = EntryStatus.Rigging;
        RebuildList();
        try
        {
            var rig = MaybeClean( entry, await run() );
            entry.Rig = rig;
            entry.Status = EntryStatus.Rigged;
            entry.Detail = rig.Explanation;
            entry.StateToken++;
            ShowInPreview( entry );
            SetStatus( $"{entry.FileName}: {rig.Skeleton.Joints.Count} joints via {rig.SolverName}", Theme.Green );
            RebuildList();
        }
        catch ( Exception e )
        {
            SetStatus( $"{entry.FileName}: {FirstLine( e.Message )}", Theme.Red );
            FlashError( entry, FirstLine( e.Message ) );   // red for 4s, then restores the row
            throw;
        }
    }

    /// <summary>Shows an error on a row TRANSIENTLY: the row turns red with the
    /// message, then after 4s restores itself to its resting state (Rigged if it
    /// has a rig, else Ready) instead of staying red forever. A newer action on
    /// the row (bumping StateToken) cancels the restore so it never clobbers it.</summary>
    async void FlashError( ModelEntry entry, string message )
    {
        entry.Status = EntryStatus.Failed;
        entry.Detail = message;
        var token = ++entry.StateToken;
        RebuildList();

        await Task.Delay( 4000 );
        if ( entry.StateToken != token )
            return;   // something newer touched this row - leave it be

        if ( entry.Rig is not null )
        {
            entry.Status = EntryStatus.Rigged;
            entry.Detail = entry.Rig.Explanation;
        }
        else if ( entry.Analysis is not null )
        {
            entry.Status = EntryStatus.Ready;
            entry.Detail = entry.Analysis.Classification.Explanation;
        }
        else
            return;   // never reached a good state - keep the failure visible

        entry.StateToken++;
        RebuildList();
    }

    /// <summary>Opens the vast.ai dialog wired to RENT a GPU for the first
    /// analyzed entry with the currently selected model (Auto → best).</summary>
    void OpenCloudDialog()
    {
        var entry = _entries.FirstOrDefault( e => e.Analysis is not null );
        if ( entry is null )
        {
            SetStatus( "Add a model file first - the cloud rig uploads that mesh.", Theme.Yellow );
            return;
        }
        var model = ResolveCloudModel( entry );
        if ( model is null )
        {
            SetStatus( "Enable a model first (Manage Models) - the cloud rig runs "
                + "the selected model remotely.", Theme.Yellow );
            return;
        }

        var dialog = new VastDialog( this, model,
            ( offer, session ) => RunCloud( entry,
                () => session.RigAsync( entry.Analysis, model.Id, model.Title, offer ) ),
            // Provision the chosen model onto an already-running box (no rig, no
            // rent, no destroy) so rigs can be offloaded there later.
            ( instanceId, provModel, session ) =>
                session.ProvisionOnExisting( instanceId, provModel.Id ) );
        dialog.Window.Show();
    }

    /// <summary>Row "Offload…": pick one of the user's already-running GPUs and
    /// run this entry's rig on it (reuse a box; no new rental).</summary>
    void OpenOffloadPicker( ModelEntry entry )
    {
        if ( entry.Analysis is null )
        {
            SetStatus( $"{entry.FileName}: still analyzing - offload once it is ready.", Theme.Yellow );
            return;
        }
        var model = ResolveCloudModel( entry );
        if ( model is null )
        {
            SetStatus( "Enable a model first (Manage Models) before offloading.", Theme.Yellow );
            return;
        }
        var picker = new GpuPickerDialog( this, model,
            ( instanceId, destroyAfter, session ) => RunCloud( entry,
                () => session.RigOnExisting(
                    entry.Analysis, model.Id, model.Title, instanceId, destroyAfter ) ) );
        picker.Window.Show();
    }

    /// <summary>Entries with a generated rig, ready for export.</summary>
    public IReadOnlyList<ModelEntry> RiggedEntries
        => _entries.Where( e => e.Rig is not null ).ToList();

    static RigNetBundle _loadedModel;
    static string _loadedModelPath;
    static RigNetBundle _modelOverride;

    /// <summary>Filename of the user-loaded checkpoint override, null when unset.</summary>
    public static string ModelOverrideLabel { get; private set; }

    /// <summary>A user-picked checkpoint bundle takes priority over the catalog install.</summary>
    public static void SetModelOverride( RigNetBundle bundle, string label )
    {
        _modelOverride = bundle;
        ModelOverrideLabel = bundle is null ? null : $"Using your checkpoint: {label}";
    }

    /// <summary>
    /// Resolves the active neural model: the user's own checkpoint override when
    /// set, otherwise the first installed catalog model (loaded off the main
    /// thread, cached until the file changes). Null (→ geometric fallback with an
    /// explanation) when neither exists.
    /// </summary>
    static async Task<RigNetBundle> LoadInstalledModel()
    {
        if ( _modelOverride is not null )
            return _modelOverride;

        var entry = ModelCatalog.Entries.FirstOrDefault( CatalogDialog.IsInstalled );
        if ( entry is null )
            return null;
        var path = Path.Combine( CatalogDialog.InstallDir( entry ), entry.Files[0].FileName );
        if ( _loadedModel is not null && _loadedModelPath == path )
            return _loadedModel;

        var bundle = await Task.Run( () => RigNetBundle.FromCheckpointsZip( File.ReadAllBytes( path ) ) );
        _loadedModel = bundle;
        _loadedModelPath = path;
        return bundle;
    }

    // ================================================================== step 3

    /// <summary>Exports every rigged entry: writes fbx + vmdl into the project's assets
    /// and compiles; per-entry failures are isolated onto their row chips.</summary>
    public async Task ExportAll()
    {
        var rigged = RiggedEntries;
        if ( rigged.Count == 0 )
        {
            SetStatus( "Nothing to export - rig a model first.", Theme.Yellow );
            return;
        }

        var folder = (_outputFolderEdit?.Text ?? "models/autorig").Trim().Trim( '/', '\\' );
        var assetsRoot = Project.Current?.GetAssetsPath();
        if ( string.IsNullOrEmpty( assetsRoot ) )
        {
            SetStatus( "No active project - open a project before exporting.", Theme.Red );
            return;
        }
        var directory = System.IO.Path.Combine( assetsRoot, folder );
        System.IO.Directory.CreateDirectory( directory );

        var failures = 0;
        foreach ( var entry in rigged )
        {
            try
            {
                SetStatus( $"Exporting {entry.FileName}…", Theme.Blue );
                var baseName = System.IO.Path.GetFileNameWithoutExtension( entry.FileName );
                var bundle = await Task.Run( () =>
                    AutoRig.Export.RigExporter.Export( entry.Mesh, entry.Rig, baseName, folder ) );

                // Collision-safe file names (never overwrite user content).
                var fbxPath = UniquePath( directory, bundle.FbxFileName );
                var vmdlName = System.IO.Path.GetFileNameWithoutExtension( fbxPath ) + ".vmdl";
                var vmdlPath = System.IO.Path.Combine( directory, vmdlName );
                var vmdlText = bundle.Vmdl.Replace(
                    bundle.FbxFileName, $"{folder}/{System.IO.Path.GetFileName( fbxPath )}" );

                await Task.Run( () =>
                {
                    System.IO.File.WriteAllBytes( fbxPath, bundle.Fbx );
                    System.IO.File.WriteAllText( vmdlPath, vmdlText );
                    // Companion textures/materials are OURS to refresh - plain
                    // overwrite (UniquePath only protects user content).
                    foreach ( var (extraName, bytes) in bundle.ExtraFiles )
                        System.IO.File.WriteAllBytes(
                            System.IO.Path.Combine( directory, extraName ), bytes );
                } );

                foreach ( var (extraName, _) in bundle.ExtraFiles )
                    AssetSystem.RegisterFile( System.IO.Path.Combine( directory, extraName ) );
                AssetSystem.RegisterFile( fbxPath );
                var vmdlAsset = AssetSystem.RegisterFile( vmdlPath );
                vmdlAsset?.Compile( true );

                entry.Status = EntryStatus.Exported;
                entry.Detail = $"{folder}/{vmdlName}";
            }
            catch ( Exception e )
            {
                failures++;
                entry.Status = EntryStatus.Failed;
                entry.Detail = FirstLine( e.Message );
            }
            RebuildList();
        }

        SetStatus(
            failures == 0
                ? $"Exported {rigged.Count} model(s) to {folder}/ - skinned .fbx "
                    + "(Blender-importable) + compiled .vmdl."
                : $"Exported with {failures} failure(s) - check the row chips.",
            failures == 0 ? Theme.Green : Theme.Yellow );
    }

    static string UniquePath( string directory, string fileName )
    {
        var path = System.IO.Path.Combine( directory, fileName );
        if ( !System.IO.File.Exists( path ) )
            return path;
        var stem = System.IO.Path.GetFileNameWithoutExtension( fileName );
        var extension = System.IO.Path.GetExtension( fileName );
        var n = 2;
        while ( System.IO.File.Exists( System.IO.Path.Combine( directory, $"{stem}_{n}{extension}" ) ) )
            n++;
        return System.IO.Path.Combine( directory, $"{stem}_{n}{extension}" );
    }

    /// <summary>Shows an entry's mesh + rig in the preview viewport and adjust panel.</summary>
    public void ShowInPreview( ModelEntry entry )
    {
        if ( _preview is null || entry?.Mesh is null )
            return;
        var entryChanged = !ReferenceEquals( _previewEntry, entry );
        _previewEntry = entry;
        _preview.SetRig( entry.Mesh, entry.Rig );
        if ( _explanationLabel is not null )
            _explanationLabel.Text = entry.Rig?.Explanation ?? entry.Detail ?? "";
        if ( entryChanged )
            _adjustPanel?.SetEntry( entry );
        else
            _adjustPanel?.Refresh();   // same entry, new rig (auto-rig on add)
    }
}