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

namespace AutoRig.Editor;

/// <summary>
/// "Connect Vast.ai" - two tabs. RENT A GPU: pick an offer, rent, rig, and the
/// instance is destroyed (verified) the moment results are back - and only that
/// one. MY RENTALS: the boxes you already have running - PROVISION a model onto
/// one (install it so a rig can be offloaded there later; you can provision
/// several) or CANCEL one you are done with (destroys exactly that id, verified).
/// Offloading a rig onto a running box lives on each model's row in the main
/// window ("Offload…"), not here. Spending money is never automatic: every
/// rent/provision/cancel is an explicit click that states what it does.
/// </summary>
public sealed class VastDialog : Dialog
{
    static string KeyPath => Path.Combine(
        Project.Current?.GetAssetsPath() ?? ".", "autorig_dl", "vast.json" );

    readonly CatalogEntry _model;
    readonly Func<VastOffer, VastRigSession, System.Threading.Tasks.Task> _rig;
    readonly Func<long, CatalogEntry, VastRigSession, System.Threading.Tasks.Task> _provision;
    readonly List<DlModelRegistry.ModelChoice> _choices;

    LineEdit _keyEdit;
    Label _status;
    Button _rentTabButton, _rentalsTabButton;
    Widget _rentPage, _rentalsPage;
    Layout _offersLayout, _instancesLayout;
    ComboBox _provisionModelCombo;

    public VastDialog(
        Widget parent, CatalogEntry model,
        Func<VastOffer, VastRigSession, System.Threading.Tasks.Task> rig,
        Func<long, CatalogEntry, VastRigSession, System.Threading.Tasks.Task> provision = null ) : base( parent )
    {
        _model = model;
        _rig = rig;
        _provision = provision;
        _choices = DlModelRegistry.EnabledChoices();

        Window.Title = $"Rig in the cloud (experimental) - {model.Title}";
        Window.SetWindowIcon( "cloud" );
        Window.MinimumSize = new Vector2( 720, 460 );

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

        var experimental = new Label(
            "EXPERIMENTAL. Every model is wired for the cloud and the instance "
            + "installs its dependencies automatically (Blender, voxelizer, "
            + "checkpoints, etc.), but the pipelines are authored from each repo "
            + "and not yet validated on a live GPU - some may fail to download, "
            + "install, or run. If a rig fails you see the remote log, and the "
            + "instance is still destroyed so you are not billed for idle time. "
            + "Rig locally as the reliable path.", this );
        experimental.WordWrap = true;
        experimental.SetStyles(
            $"color: {Theme.Yellow.Hex}; font-weight: 600; "
            + $"background-color: {Theme.Yellow.WithAlpha( 0.12f ).Rgba}; "
            + "border-radius: 5px; padding: 8px 10px;" );
        Layout.Add( experimental );

        // ---- shared API key ----
        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;" );   // no Password mode on LineEdit
        _keyEdit.Text = LoadKey();
        keyRow.Add( _keyEdit, 1 );

        // ---- tab strip ----
        var tabs = Layout.AddRow();
        tabs.Spacing = UiSpacing.Related;
        _rentTabButton = new Button( "Rent a GPU", "bolt", this );
        _rentTabButton.Clicked = () => ShowTab( rentals: false );
        tabs.Add( _rentTabButton );
        _rentalsTabButton = new Button( "My rentals", "dns", this );
        _rentalsTabButton.ToolTip = "GPUs you already have running: provision a "
            + "model onto one, or cancel one you are done with.";
        _rentalsTabButton.Clicked = () => ShowTab( rentals: true );
        tabs.Add( _rentalsTabButton );
        tabs.AddStretchCell();

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

        BuildRentPage();
        BuildRentalsPage();
        Layout.AddStretchCell();

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

        ShowTab( rentals: false );
    }

    // ============================================================ tab plumbing

    bool _offersLoaded;

    void ShowTab( bool rentals )
    {
        _rentPage.Visible = !rentals;
        _rentalsPage.Visible = rentals;
        StyleTab( _rentTabButton, !rentals );
        StyleTab( _rentalsTabButton, rentals );
        if ( rentals )
            _ = LoadRentals();
        else if ( !_offersLoaded && !string.IsNullOrWhiteSpace( _keyEdit.Text ) )
            _ = LoadOffers();   // the "Rent a GPU" tab searches on its own - no separate button
    }

    void StyleTab( Button button, bool active )
        => button.SetStyles( active
            ? $"background-color: {Theme.Blue.Hex}; color: #ffffff; font-weight: 600;"
            : $"background-color: {Theme.ControlBackground.Hex};" );

    // ================================================================ rent tab

    void BuildRentPage()
    {
        _rentPage = new Widget( this ) { Layout = Layout.Column() };
        _rentPage.Layout.Spacing = UiSpacing.Group;
        Layout.Add( _rentPage, 1 );

        var header = new Label(
            $"Rent a GPU on vast.ai to run {_model.Title} remotely. You pay vast.ai "
            + "directly; the instance is destroyed automatically after the rig "
            + "(only the one rented here - your other instances are never touched).", this );
        header.WordWrap = true;
        header.SetStyles( $"color: {Theme.TextLight.Hex};" );
        _rentPage.Layout.Add( header );

        var actions = _rentPage.Layout.AddRow();
        actions.Spacing = UiSpacing.Related;
        var refresh = new Button( "Refresh offers", this ) { Icon = "refresh" };
        refresh.ToolTip = "Search vast.ai again for the cheapest rentable GPUs.";
        refresh.Clicked = () => _ = LoadOffers();
        actions.Add( refresh );
        actions.AddStretchCell();

        _offersLayout = _rentPage.Layout.AddColumn();
        _offersLayout.Spacing = UiSpacing.Related;

        CheckStaleRental( _rentPage.Layout );
    }

    async System.Threading.Tasks.Task LoadOffers()
    {
        _offersLoaded = true;
        SaveKey( _keyEdit.Text );
        _status.Text = "searching offers…";
        _offersLayout.Clear( true );
        try
        {
            using var client = new VastClient( _keyEdit.Text );
            var minVramMb = (int)(_model.BaseRamBytes / (1024L * 1024) / 2);   // model VRAM ≈ half host RAM floor
            var offers = await client.SearchOffers( Math.Max( minVramMb, 8192 ), VastProtocol.DiskGb );
            _status.Text = offers.Count == 0
                ? "no rentable GPUs matched - try again later."
                : $"{offers.Count} rentable GPUs (cheapest first). Renting requires the click below - never automatic.";
            foreach ( var offer in offers.Take( 10 ) )
                BuildOfferRow( offer );
        }
        catch ( Exception e )
        {
            _status.Text = $"search failed: {e.Message}";
        }
    }

    void BuildOfferRow( VastOffer offer )
    {
        var row = _offersLayout.AddRow();
        row.Spacing = UiSpacing.Related;
        var label = new Label(
            $"{offer.GpuName} ×{offer.NumGpus}  ·  {offer.GpuRamMb / 1024f:0.#} GB  ·  "
            + $"${offer.DollarsPerHour:0.000}/hr  ·  reliability {offer.Reliability:P0}", this );
        row.Add( label, 1 );

        var rent = new Button(
            $"Rent (≤ ${offer.EstimateCost( 45 ):0.00} for 45 min)", this )
        { Icon = "bolt" };
        rent.ToolTip = "Rents this exact offer, rigs, then destroys the instance "
            + "(verified). You are charged by vast.ai only while it runs.";
        rent.Clicked = () => _ = Rent( offer, rent );
        row.Add( rent );
    }

    async System.Threading.Tasks.Task Rent( VastOffer offer, Button button )
    {
        button.Enabled = false;
        try
        {
            using var client = new VastClient( _keyEdit.Text );
            var session = new VastRigSession( client, s => _status.Text = $"[cloud] {s}" );
            await _rig( offer, session );
            _status.Text += "  - done.";
        }
        catch ( Exception e )
        {
            _status.Text = $"cloud rig failed: {e.Message}";
        }
        finally
        {
            button.Enabled = true;
        }
    }

    // ============================================================= rentals tab

    void BuildRentalsPage()
    {
        _rentalsPage = new Widget( this ) { Layout = Layout.Column() };
        _rentalsPage.Layout.Spacing = UiSpacing.Group;
        Layout.Add( _rentalsPage, 1 );

        var note = new Label(
            "GPUs you currently have running on vast.ai. Provision installs the "
            + "chosen model onto a box (do it for as many models as you like) so "
            + "you can offload rigs there from a model's row. Cancel destroys "
            + "exactly that instance (verified).", this );
        note.WordWrap = true;
        note.SetStyles( $"color: {Theme.TextLight.Hex};" );
        _rentalsPage.Layout.Add( note );

        var controls = _rentalsPage.Layout.AddRow();
        controls.Spacing = UiSpacing.Related;
        controls.Add( new Label( "Model to provision", this ) );
        _provisionModelCombo = new ComboBox( this );
        _provisionModelCombo.MinimumWidth = 220;
        foreach ( var choice in _choices )
            _provisionModelCombo.AddItem( choice.Title, "model_training" );
        _provisionModelCombo.TrySelectNamed( _model.Title );
        controls.Add( _provisionModelCombo );
        var refresh = new Button( "Refresh", this ) { Icon = "refresh" };
        refresh.Clicked = () => _ = LoadRentals();
        controls.Add( refresh );
        controls.AddStretchCell();

        _instancesLayout = _rentalsPage.Layout.AddColumn();
        _instancesLayout.Spacing = UiSpacing.Related;
    }

    async System.Threading.Tasks.Task LoadRentals()
    {
        SaveKey( _keyEdit.Text );
        _status.Text = "loading your instances…";
        _instancesLayout.Clear( true );
        try
        {
            using var client = new VastClient( _keyEdit.Text );
            var instances = await client.ListInstances();
            if ( instances.Count == 0 )
            {
                _status.Text = "no rented instances found on this account.";
                return;
            }
            _status.Text = $"{instances.Count} instance(s).";
            foreach ( var instance in instances )
                BuildInstanceRow( instance );
        }
        catch ( Exception e )
        {
            _status.Text = $"could not list instances: {e.Message}";
        }
    }

    void BuildInstanceRow( VastInstances.InstanceSummary instance )
    {
        var row = _instancesLayout.AddRow();
        row.Spacing = UiSpacing.Related;

        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";
        if ( !string.IsNullOrEmpty( instance.Label ) )
            detail += $"  ·  {instance.Label}";
        var label = new Label( detail, this );
        label.SetStyles( instance.HasRigServer ? "" : $"color: {Theme.TextLight.Hex};" );
        row.Add( label, 1 );

        if ( _provision is not null )
        {
            var provision = new Button( "Provision", this ) { Icon = "download" };
            provision.Enabled = instance.IsRunning;
            provision.ToolTip = instance.IsRunning
                ? "Install the selected model onto this instance so rigs can be offloaded here."
                : "Instance is not running yet.";
            provision.Clicked = () => _ = ProvisionInstance( instance, provision );
            row.Add( provision );
        }

        // Cancel is a two-click guard: destroying a box is irreversible.
        var cancel = new Button( "Cancel", this ) { Icon = "delete_forever" };
        var armed = false;
        cancel.Clicked = () =>
        {
            if ( !armed )
            {
                armed = true;
                cancel.Text = "Confirm cancel";
                cancel.SetStyles( $"background-color: {Theme.Red.Hex}; color: #ffffff;" );
                return;
            }
            _ = CancelInstance( instance, cancel );
        };
        row.Add( cancel );
    }

    async System.Threading.Tasks.Task ProvisionInstance( VastInstances.InstanceSummary instance, Button button )
    {
        var model = SelectedProvisionModel();
        button.Enabled = false;
        try
        {
            using var client = new VastClient( _keyEdit.Text );
            var session = new VastRigSession( client, s => _status.Text = $"[provision] {s}" );
            await _provision( instance.Id, model, session );
            _status.Text += "  - provisioned.";
        }
        catch ( Exception e )
        {
            _status.Text = $"provision failed: {e.Message}";
        }
        finally
        {
            button.Enabled = true;
        }
    }

    async System.Threading.Tasks.Task CancelInstance( VastInstances.InstanceSummary instance, Button button )
    {
        button.Enabled = false;
        _status.Text = $"destroying instance {instance.Id}…";
        try
        {
            using var client = new VastClient( _keyEdit.Text );
            var gone = await client.DestroyVerified( instance.Id );
            _status.Text = gone
                ? $"instance {instance.Id} destroyed (verified)."
                : $"could not verify destruction of {instance.Id} - check console.vast.ai.";
            await LoadRentals();
        }
        catch ( Exception e )
        {
            _status.Text = $"cancel failed: {e.Message}";
            button.Enabled = true;
        }
    }

    CatalogEntry SelectedProvisionModel()
    {
        var pick = _choices.FirstOrDefault( c => c.Title == _provisionModelCombo?.CurrentText );
        return pick?.Entry ?? _model;
    }

    // ================================================================== shared

    void CheckStaleRental( Layout target )
    {
        var stale = VastRigSession.StaleRental();
        if ( stale is null )
            return;
        var row = target.AddRow();
        row.Spacing = UiSpacing.Related;
        var warning = new Label(
            $"⚠ A previous session left instance {stale.InstanceId} ({stale.Label}) "
            + "possibly running. Destroy it now to stop charges?", this );
        warning.WordWrap = true;
        warning.SetStyles( $"color: {Theme.Yellow.Hex};" );
        row.Add( warning, 1 );
        var destroy = new Button( "Destroy it", this ) { Icon = "delete_forever" };
        destroy.Clicked = () => _ = DestroyStale( stale, destroy, warning );
        row.Add( destroy );
    }

    async System.Threading.Tasks.Task DestroyStale( VastRental stale, Button button, Label warning )
    {
        button.Enabled = false;
        try
        {
            using var client = new VastClient( _keyEdit.Text );
            var session = new VastRigSession( client, s => _status.Text = s );
            warning.Text = await session.DestroyStale( stale )
                ? $"Instance {stale.InstanceId} destroyed (verified)."
                : $"Could not verify destruction of {stale.InstanceId} - check console.vast.ai.";
        }
        catch ( Exception e )
        {
            warning.Text = $"Destroy failed: {e.Message}";
            button.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 { }
    }
}