using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Sandbox;
using TextToAnimation.Animation;
using TextToAnimation.Editor.Engine;
using TextToAnimation.Editor.Generation;
using TextToAnimation.Generation;

namespace TextToAnimation.Editor.Session;

/// <summary>
/// Runs model downloads and generations for the editor without ever blocking it: work happens on worker
/// threads, progress lines are marshalled to the main thread, everything can be cancelled, and failures leave
/// the workspace untouched. Results become new clips, or replace the active clip as one undoable edit.
/// </summary>
public sealed class GenerationFlow
{
	readonly EditorSession _session;
	CancellationTokenSource _cts;

	public GeneratorService Service => GeneratorService.Instance;

	/// <summary>Progress line for the loading indicator ("Downloading … · n% · … left" draws a bar).</summary>
	public event Action<string> Progress;

	/// <summary>True while a download or generation runs.</summary>
	public bool Running => _cts is not null;

	public GenerationFlow( EditorSession session ) => _session = session;

	public void Cancel() => _cts?.Cancel();

	void Report( string line ) => Progress?.Invoke( line );

	/// <summary>Downloads the model (the Download button).</summary>
	public async Task<bool> DownloadAsync()
	{
		if ( Running ) return false;
		_cts = new CancellationTokenSource();
		_session.SetBusy( true, "Downloading UniMate" );
		try
		{
			var progress = new EngineThread.MainThreadProgress<string>( Report );
			var ok = await Service.InstallAsync( progress, _cts.Token );
			await EngineThread.SwitchToMainThread();
			_session.SetStatus( ok ? "UniMate is installed. Describe an animation and press Enter." : $"Download failed: {Service.LastError}", ok ? UI.Tone.Accent : UI.Tone.Red );
			return ok;
		}
		catch ( OperationCanceledException )
		{
			await EngineThread.SwitchToMainThread();
			_session.SetStatus( "Download paused. It resumes where it stopped next time.", UI.Tone.Amber );
			return false;
		}
		catch ( Exception e )
		{
			await EngineThread.SwitchToMainThread();
			_session.SetStatus( $"Download failed: {e.Message}", UI.Tone.Red );
			return false;
		}
		finally
		{
			await EngineThread.SwitchToMainThread();
			_cts?.Dispose();
			_cts = null;
			_session.SetBusy( false );
			Service.Refresh();
		}
	}

	/// <summary>
	/// Generates and applies the result. <paramref name="replace"/> overwrites the active clip (Regenerate);
	/// otherwise every take becomes a new clip next to it (originals are never touched).
	/// </summary>
	public async Task<bool> GenerateAsync( GenerationRequest request, bool replace, string baseName )
	{
		if ( Running || !_session.HasModel ) return false;
		var rig = _session.Rig;
		var target = _session.ActiveClip;
		if ( replace && target is null ) return false;
		_cts = new CancellationTokenSource();
		var started = DateTime.UtcNow;
		_session.SetBusy( true, "Generating" );
		try
		{
			var load = new EngineThread.MainThreadProgress<string>( Report );
			var progress = new EngineThread.MainThreadProgress<GenerationProgress>( p => Report( $"{p.Stage} · {p.Fraction * 100:0}%" ) );
			var before = _session.Workspace.Clips.ToList();
			var results = await Service.GenerateAsync( rig, request, load, progress, _cts.Token );
			await EngineThread.SwitchToMainThread();
			var record = new GenerationRecord
			{
				Mode = request.Mode.ToString(), Prompts = request.Prompts.ToList(), Seed = request.Seed,
				DurationSeconds = request.DurationSeconds, Guidance = request.Guidance, Generator = Service.Backend.Name,
			};
			if ( replace )
			{
				var first = results[0];
				_session.ReplaceClipFrames( target, first.Frames, first.Fps, "Regenerate", c =>
				{
					c.Generation = record;
					if ( request.Mode != GenerationMode.InBetween ) c.PinnedFrames.Clear();
					ClipCleanup.GenerateFootsteps( c, rig );
				} );
				for ( var i = 1; i < results.Count; i++ ) AddResult( results[i], record, $"{target.Name} take {i + 1}", false );
			}
			else
			{
				for ( var i = 0; i < results.Count; i++ )
				{
					var (result, name, select) = (results[i], results.Count > 1 ? $"{baseName} {i + 1}" : baseName, i == 0);
					UI.FrameProbe.Time( "result: add clip", () => AddResult( result, record, name, select ) );
				}
			}
			var made = _session.Workspace.Clips.Except( before ).ToList();
			if ( replace ) made.Insert( 0, target );
			_session.RecordPrompt( string.Join( " → ", request.Prompts.Where( p => !string.IsNullOrWhiteSpace( p ) ) ), request.Mode.ToString(), request.Seed, made );
			var notes = results.SelectMany( r => r.Notes ).Distinct().ToList();
			// a GPU the network couldn't run on (or that disagreed with the CPU) leaves generation on the CPU: slow, and
			// the editor shares its cores - say so instead of just being slow
			if ( Inference.Onnx.GpuAcceleration.Status is { } gpu )
			{
				if ( gpu != _gpuLogged ) Log.Warning( $"[text-to-animation] generating on the CPU: {gpu}" );
				_gpuLogged = gpu;
				notes.Add( $"Generated on the CPU ({gpu})" );
			}
			var seconds = (DateTime.UtcNow - started).TotalSeconds;
			_session.SetStatus( notes.Count > 0 ? string.Join( " ", notes ) : $"Generated in {seconds:0} s.", notes.Count > 0 ? UI.Tone.Amber : UI.Tone.Accent );
			return true;
		}
		catch ( OperationCanceledException )
		{
			await EngineThread.SwitchToMainThread();
			_session.SetStatus( "Generation cancelled.", UI.Tone.Neutral );
			return false;
		}
		catch ( Exception e )
		{
			await EngineThread.SwitchToMainThread();
			Log.Warning( $"[text-to-animation] generation failed: {e}" );
			_session.SetStatus( $"Generation failed: {e.Message}", UI.Tone.Red );
			return false;
		}
		finally
		{
			await EngineThread.SwitchToMainThread();
			_cts?.Dispose();
			_cts = null;
			_session.SetBusy( false );
		}
	}

	string _gpuLogged;

	void AddResult( GeneratedMotion motion, GenerationRecord record, string name, bool select )
	{
		var clip = new AnimClip
		{
			Name = name,
			Fps = motion.Fps,
			Frames = motion.Frames,
			Origin = ClipOrigin.Generated,
			Generation = record.Clone(),
			Looping = false,
		};
		clip.Generation.Seed = motion.Seed;
		UI.FrameProbe.Time( "result: footsteps", () => ClipCleanup.GenerateFootsteps( clip, _session.Rig ) );
		_session.AddClip( clip, select );
	}

	/// <summary>
	/// The request for <paramref name="mode"/> on the active clip: its final frames, pinned frames and locked
	/// bones. Empty prompts are dropped; variations and in-betweens fall back to the clip's original prompt.
	/// </summary>
	public static GenerationRequest BuildRequest( EditorSession session, GenerationMode mode, IEnumerable<string> prompts,
		float durationSeconds, int seed, int count, float guidance, int steps, float variationStrength,
		IReadOnlyCollection<int> keepBones = null, bool cleanUp = true )
	{
		var clip = session.ActiveClip;
		var list = (prompts ?? Enumerable.Empty<string>()).Select( p => (p ?? "").Trim() ).ToList();
		if ( mode != GenerationMode.Expansion ) list = new List<string> { list.FirstOrDefault() ?? "" };
		else list = list.Where( p => p.Length > 0 ).ToList();
		if ( mode is GenerationMode.Variation or GenerationMode.InBetween && list.Count > 0 && list[0].Length == 0
			&& clip?.Generation?.Prompts.FirstOrDefault() is { } original )
			list[0] = original;
		var usesClip = clip is not null && mode is GenerationMode.InBetween or GenerationMode.TextEdit or GenerationMode.Variation;
		return new GenerationRequest
		{
			Mode = mode,
			Prompts = list,
			DurationSeconds = mode == GenerationMode.TextToMotion ? durationSeconds : 0,
			OutputFps = clip?.Fps ?? session.Workspace?.DefaultFps ?? 30f,
			Seed = seed,
			Count = Math.Clamp( count, 1, 8 ),
			Guidance = guidance,
			Steps = steps,
			CleanUp = cleanUp,
			SourceFrames = usesClip ? session.ActiveFrames : null,
			SourceFps = clip?.Fps ?? 30f,
			KeepFrames = usesClip ? clip.PinnedFrames.ToList() : new List<int>(),
			KeepBones = !usesClip ? Array.Empty<int>()
				: keepBones?.ToArray() ?? clip.LockedBones.Select( session.Rig.Skeleton.IndexOf ).Where( i => i >= 0 ).ToArray(),
			VariationStrength = variationStrength,
		};
	}

	/// <summary>What is missing before the request can run (null = ready), phrased for the status bar.</summary>
	public static string Validate( GenerationRequest request )
	{
		var needsClip = request.Mode is GenerationMode.InBetween or GenerationMode.TextEdit or GenerationMode.Variation;
		if ( needsClip && request.SourceFrames is not { Count: > 1 } ) return "Open an animation first.";
		if ( request.Mode is GenerationMode.TextToMotion or GenerationMode.TextEdit && string.IsNullOrWhiteSpace( request.Prompts.FirstOrDefault() ) )
			return "Describe the motion first.";
		if ( request.Mode == GenerationMode.Expansion && request.Prompts.Count == 0 ) return "Add at least one step.";
		if ( request.Mode == GenerationMode.InBetween && request.KeepFrames.Count < 2 )
			return "Pin at least two frames on the timeline first (double click the Pinned lane).";
		return null;
	}

	/// <summary>A clip name from a prompt (see <see cref="Inference.UniMate.UniMatePrompt.ClipName"/>).</summary>
	public static string NameFromPrompt( string prompt ) => Inference.UniMate.UniMatePrompt.ClipName( prompt );
}