using System;
using System.IO;
using System.Linq;
using System.Threading;
using AutoRig.Dl;
using Editor;
using Sandbox;

namespace AutoRig.Editor;

/// <summary>
/// The deep-learning model catalog (spec §7-§8): one row per model with license,
/// size, a hardware-advisor verdict chip for the CURRENT mesh, and Download with
/// progress + cancel. When no model fits this machine, rows explain why and the
/// geometric solver remains the path. Spacing per UiSpacing throughout.
/// </summary>
public sealed class CatalogDialog : Dialog
{
    readonly int _vertexCount;
    CancellationTokenSource _cancel;

    public CatalogDialog( Widget parent, int vertexCount ) : base( parent )
    {
        _vertexCount = vertexCount;
        Window.Title = "Manage Deep Learning Models";
        Window.SetWindowIcon( "model_training" );
        Window.MinimumSize = new Vector2( 640, 320 );

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

        var header = new Label(
            "Models download directly from their authors' official distribution. "
            + "Each shows whether it fits this machine and the current mesh.", this );
        header.SetStyles( $"color: {Theme.TextLight.Hex}; margin-bottom: {UiSpacing.HeaderBottom}px;" );
        Layout.Add( header );

        foreach ( var entry in ModelCatalog.Entries )
            BuildRow( entry );

        BuildOwnCheckpointRow();

        Layout.AddStretchCell();
        var footer = Layout.AddRow();
        footer.AddStretchCell();
        var close = new Button( "Close", this );
        close.Clicked = () => Window.Close();
        footer.Add( close );
    }

    /// <summary>Spec §7 mode 2: point the solver at your own RigNet checkpoints zip.</summary>
    void BuildOwnCheckpointRow()
    {
        var box = Layout.AddColumn();
        box.Spacing = UiSpacing.Related;
        box.Margin = new Sandbox.UI.Margin(
            UiSpacing.Related, UiSpacing.RowGap, UiSpacing.Related, UiSpacing.RowGap );

        var title = new Label( "Your own checkpoint", this );
        title.SetStyles( "font-weight: 600;" );
        box.Add( title );
        var hint = new Label(
            "Load a RigNet-format checkpoints zip you trained or downloaded yourself "
            + "(must contain gcn_meanshift / rootnet / bonenet / skinnet model_best.pth.tar).", this );
        hint.SetStyles( $"color: {Theme.TextLight.Hex};" );
        box.Add( hint );

        var row = box.AddRow();
        row.Spacing = UiSpacing.Related;
        var status = new Label( AutoRigWindow.ModelOverrideLabel ?? "", this );
        status.SetStyles( $"color: {Theme.TextLight.Hex};" );
        var load = new Button( "Load checkpoint zip…", "folder_open", this );
        load.Clicked = async () =>
        {
            var path = EditorUtility.OpenFileDialog(
                "Select RigNet checkpoints zip", "Checkpoint Bundles (*.zip)", null );
            if ( string.IsNullOrEmpty( path ) )
                return;
            load.Enabled = false;
            status.Text = "Loading…";
            status.SetStyles( $"color: {Theme.Blue.Hex};" );
            try
            {
                var bytes = File.ReadAllBytes( path );
                var bundle = await System.Threading.Tasks.Task.Run(
                    () => Solve.RigNetBundle.FromCheckpointsZip( bytes ) );
                DlModelRegistry.SetUserCheckpoint( bundle, Path.GetFileName( path ) );
                status.Text = $"Loaded {Path.GetFileName( path )} - it is now in the Model dropdown.";
                status.SetStyles( $"color: {Theme.Green.Hex};" );
            }
            catch ( Exception e )
            {
                status.Text = FirstLine( e.Message );
                status.SetStyles( $"color: {Theme.Red.Hex};" );
            }
            load.Enabled = true;
        };
        row.Add( load );
        row.Add( status );
        row.AddStretchCell();
    }

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

    void BuildRow( CatalogEntry entry )
    {
        // Editor assemblies are not whitelisted - probe the machine here and inject.
        var available = GC.GetGCMemoryInfo().TotalAvailableMemoryBytes;
        var (verdict, why) = HardwareAdvisor.Assess( entry, Math.Max( _vertexCount, 1 ), available );
        var tone = verdict switch
        {
            Verdict.Recommended => Theme.Green,
            Verdict.Slow => Theme.Yellow,
            _ => Theme.Red,
        };

        var box = Layout.AddColumn();
        box.Spacing = UiSpacing.Related;
        box.Margin = new Sandbox.UI.Margin( UiSpacing.Related, UiSpacing.RowGap, UiSpacing.Related, UiSpacing.RowGap );

        var titleRow = box.AddRow();
        titleRow.Spacing = UiSpacing.Related;

        // Enable toggle - LOCKED OFF when this machine cannot run the model or the
        // runtime does not support the architecture yet.
        var meetsHardware = DlModelRegistry.MeetsHardware( entry );
        var usable = meetsHardware && entry.RuntimeSupported;
        var enableCheck = new Checkbox( "", this )
        {
            Value = usable && IsInstalled( entry ) && DlModelRegistry.IsEnabled( entry.Id ),
            Enabled = usable && IsInstalled( entry ),
            ToolTip = !entry.RuntimeSupported
                ? "Not runnable yet: the pure-C# port of this architecture is planned."
                : !meetsHardware
                    ? $"Disabled: this machine does not meet the model's requirements. {why}"
                    : IsInstalled( entry )
                        ? "Enabled models appear in the Model dropdown; Auto picks the best one."
                        : "Download the model first.",
        };
        if ( !usable )
            DlModelRegistry.SetEnabled( entry.Id, false );
        enableCheck.Toggled = () => DlModelRegistry.SetEnabled( entry.Id, enableCheck.Value );
        titleRow.Add( enableCheck );

        var title = new Label( entry.Title, this );
        title.SetStyles( $"font-weight: 600;" );
        titleRow.Add( title );
        var verdictChip = new Label( verdict.ToString(), this );
        verdictChip.ToolTip = why;
        verdictChip.SetStyles( $"color: {tone.Hex}; padding: 0px {UiSpacing.Related}px;" );
        titleRow.Add( verdictChip );
        titleRow.AddStretchCell();

        var summary = new Label( entry.Summary, this );
        summary.SetStyles( $"color: {Theme.TextLight.Hex};" );
        box.Add( summary );
        var license = new Label( $"License: {entry.License}", this );
        license.ToolTip = entry.Attribution;
        license.SetStyles( $"color: {Theme.TextLight.Hex};" );
        box.Add( license );

        var actionRow = box.AddRow();
        actionRow.Spacing = UiSpacing.Related;
        var status = new Label( IsInstalled( entry ) ? "Installed." : "", this );
        status.SetStyles( $"color: {Theme.TextLight.Hex};" );

        var download = new Button( IsInstalled( entry ) ? "Re-download" : "Download", "download", this );
        var cancelButton = new Button( "Cancel", "close", this ) { Visible = false };
        download.Enabled = verdict != Verdict.Insufficient && entry.RuntimeSupported;
        if ( !entry.RuntimeSupported )
            download.ToolTip = "Runtime support for this architecture is planned - "
                + "download will unlock with it.";
        else if ( verdict == Verdict.Insufficient )
            download.ToolTip = why;

        download.Clicked = async () =>
        {
            _cancel = new CancellationTokenSource();
            download.Visible = false;
            cancelButton.Visible = true;
            try
            {
                foreach ( var file in entry.Files )
                {
                    await DownloadManager.Download( entry, file, InstallDir( entry ),
                        ( done, total ) => status.Text = total > 0
                            ? $"{file.FileName}: {done / 1048576.0:0.0} / {total / 1048576.0:0.0} MB ({done * 100.0 / total:0}%)"
                            : $"{file.FileName}: {done / 1048576.0:0.0} MB",
                        _cancel.Token );
                }
                status.Text = "Installed and enabled.";
                status.SetStyles( $"color: {Theme.Green.Hex};" );
                DlModelRegistry.Changed?.Invoke();
            }
            catch ( OperationCanceledException )
            {
                status.Text = "Cancelled (partial file kept for resume).";
                status.SetStyles( $"color: {Theme.Yellow.Hex};" );
            }
            catch ( Exception e )
            {
                status.Text = e.Message;
                status.SetStyles( $"color: {Theme.Red.Hex};" );
            }
            download.Visible = true;
            download.Text = IsInstalled( entry ) ? "Re-download" : "Download";
            cancelButton.Visible = false;
        };
        cancelButton.Clicked = () => _cancel?.Cancel();

        actionRow.Add( download );
        actionRow.Add( cancelButton );
        actionRow.Add( status );
        actionRow.AddStretchCell();
    }

    /// <summary>Per-entry install folder under the active project's assets.</summary>
    public static string InstallDir( CatalogEntry entry )
        => Path.Combine( Project.Current?.GetAssetsPath() ?? ".", "autorig_dl", entry.Id );

    public static bool IsInstalled( CatalogEntry entry )
        => entry.Files.All( f => File.Exists( Path.Combine( InstallDir( entry ), f.FileName ) ) );
}