using System;
using System.IO;
using System.Text.Json;
using AutoRig.Dl;
using AutoRig.Vast;
using Editor;
using Editor.AutoRig.Vast;
using Sandbox;

namespace AutoRig.Editor;

/// <summary>
/// "Offload…" target picker: shows the GPUs you already have running, you click
/// the one you want and press Rig, and the rig is offloaded onto that rental
/// (reusing a box that already has the model installed - no new rental). The
/// box keeps running afterwards, unless "Destroy after rigging" is checked. This
/// only ever drives / destroys the exact instance you pick.
/// </summary>
public sealed class GpuPickerDialog : Dialog
{
    static string KeyPath => Path.Combine(
        Project.Current?.GetAssetsPath() ?? ".", "autorig_dl", "vast.json" );

    readonly CatalogEntry _model;
    readonly Func<long, bool, VastRigSession, System.Threading.Tasks.Task> _offload;

    LineEdit _keyEdit;
    Label _status;
    Layout _listLayout;
    Checkbox _destroyAfter;
    Button _rigButton;

    System.Collections.Generic.List<VastInstances.InstanceSummary> _instances = new();
    long? _selectedId;

    public GpuPickerDialog(
        Widget parent, CatalogEntry model,
        Func<long, bool, VastRigSession, System.Threading.Tasks.Task> offload ) : base( parent )
    {
        _model = model;
        _offload = offload;

        Window.Title = $"Offload {model.Title} to a running GPU (experimental)";
        Window.SetWindowIcon( "dns" );
        Window.MinimumSize = new Vector2( 640, 420 );

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

        var header = new Label(
            $"Pick one of your running GPUs and press Rig to run {model.Title} on it. "
            + "The instance must be running our rig server (it auto-installs the "
            + "model if needed). The box keeps running unless you check destroy.", this );
        header.WordWrap = true;
        header.SetStyles( $"color: {Theme.TextLight.Hex};" );
        Layout.Add( header );

        var keyRow = Layout.AddRow();
        keyRow.Spacing = UiSpacing.Related;
        keyRow.Add( new Label( "API key", this ) );
        _keyEdit = new LineEdit( this ) { PlaceholderText = "vast.ai API key" };
        _keyEdit.SetStyles( "font-family: monospace;" );
        _keyEdit.Text = LoadKey();
        keyRow.Add( _keyEdit, 1 );
        var refresh = new Button( "Refresh", this ) { Icon = "refresh" };
        refresh.Clicked = () => _ = LoadInstances();
        keyRow.Add( refresh );

        _status = new Label( "", this );
        _status.WordWrap = true;
        _status.SetStyles( $"color: {Theme.TextLight.Hex};" );
        Layout.Add( _status );

        _listLayout = Layout.AddColumn();
        _listLayout.Spacing = UiSpacing.Related;
        Layout.AddStretchCell();

        var footer = Layout.AddRow();
        footer.Spacing = UiSpacing.Related;
        _destroyAfter = new Checkbox( "Destroy instance after rigging", this ) { Value = false };
        _destroyAfter.ToolTip = "Off by default. When on, this exact instance is "
            + "destroyed (verified) once the rig finishes.";
        footer.Add( _destroyAfter );
        footer.AddStretchCell();
        _rigButton = new Button( "Rig", "engineering", this ) { Enabled = false };
        _rigButton.Clicked = () => _ = OffloadSelected();
        footer.Add( _rigButton );
        var close = new Button( "Close", this );
        close.Clicked = () => Window.Close();
        footer.Add( close );

        _ = LoadInstances();
    }

    async System.Threading.Tasks.Task LoadInstances()
    {
        SaveKey( _keyEdit.Text );
        _status.Text = "loading your instances…";
        _selectedId = null;
        _rigButton.Enabled = false;
        try
        {
            using var client = new VastClient( _keyEdit.Text );
            _instances = await client.ListInstances();
            if ( _instances.Count == 0 )
                _status.Text = "no rented instances found - rent one first, or check your key.";
            else
                _status.Text = $"{_instances.Count} instance(s). Click one, then press Rig.";
        }
        catch ( Exception e )
        {
            _status.Text = $"could not list instances: {e.Message}";
            _instances = new();
        }
        Render();
    }

    void Render()
    {
        _listLayout.Clear( true );
        foreach ( var instance in _instances )
        {
            var captured = instance;
            var selected = _selectedId == instance.Id;
            var detail = $"{instance.GpuName}  ·  {instance.ActualStatus}  ·  id {instance.Id}";
            if ( instance.DollarsPerHour > 0 )
                detail += $"  ·  ${instance.DollarsPerHour:0.000}/hr";
            detail += instance.HasRigServer ? "  ·  rig server ✓" : "  ·  no rig server";

            var button = new Button( detail, this ) { Enabled = instance.IsRunning };
            button.SetStyles( selected
                ? $"background-color: {Theme.Blue.Hex}; color: #ffffff; text-align: left;"
                : $"background-color: {Theme.ControlBackground.Hex}; text-align: left;" );
            button.ToolTip = instance.IsRunning
                ? (instance.HasRigServer
                    ? "Select this GPU."
                    : "Select (no rig server detected - offload may fail unless it was started by auto-rig).")
                : "Instance is not running yet.";
            button.Clicked = () =>
            {
                _selectedId = captured.Id;
                _rigButton.Enabled = true;
                Render();
            };
            _listLayout.Add( button );
        }
    }

    async System.Threading.Tasks.Task OffloadSelected()
    {
        if ( _selectedId is not long id )
            return;
        _rigButton.Enabled = false;
        try
        {
            using var client = new VastClient( _keyEdit.Text );
            var session = new VastRigSession( client, s => _status.Text = $"[offload] {s}" );
            await _offload( id, _destroyAfter.Value, session );
            _status.Text += "  - done.";
        }
        catch ( Exception e )
        {
            _status.Text = $"offload failed: {e.Message}";
        }
        finally
        {
            _rigButton.Enabled = true;
        }
    }

    static string LoadKey()
    {
        try
        {
            if ( File.Exists( KeyPath ) )
            {
                using var doc = JsonDocument.Parse( File.ReadAllText( KeyPath ) );
                if ( doc.RootElement.TryGetProperty( "api_key", out var key ) )
                    return key.GetString() ?? "";
            }
        }
        catch { }
        return "";
    }

    static void SaveKey( string key )
    {
        try
        {
            Directory.CreateDirectory( Path.GetDirectoryName( KeyPath )! );
            File.WriteAllText( KeyPath, JsonSerializer.Serialize(
                new { api_key = key?.Trim() ?? "" } ) );
        }
        catch { }
    }
}