using System;
using System.Collections.Generic;

namespace BlockParty;

/// <summary>
/// Estimates the byte size a GIF export of a replay segment would produce, by sampling frames through the
/// real encoder: the first frame is encoded at full cost, then a few short runs of consecutive ticks spread
/// across the span measure the average delta-frame cost, which is extrapolated to the full frame count.
/// Runs (not single pairs) because per-tick delta cost swings hard — a single sampled pair can land on a
/// hit-stop freeze and measure ~nothing. 5 sites × 4 deltas stayed within ±8% of the real size on test clips.
///
/// Incremental for the same reason <see cref="GifExporter"/> is: sprite batches upload once per main-loop
/// update, so each sampled tick needs its own update (seek now, capture next Step). Drive it with
/// <see cref="Step"/> once per UI tick; ~25 updates total. The shrunk preview briefly shows the sample
/// frames flick past (the screen renders each seeked tick) — it reads as a quick scrub across the selected
/// span. The replay is seeked back to where it was when sampling completes or the job is
/// <see cref="Cancel"/>led; the camera is only touched inside a capture, never across updates.
///
/// The native 240px sample frames are kept (~25 × 230KB), so <see cref="Retarget"/> to a different output
/// scale re-encodes them one per Step with no seeks and no preview scrub.
/// </summary>
public sealed class GifSizeEstimator
{
	/// <summary>Sampling seeks + captures (and encodes at the requested scale as it goes); Encoding
	/// re-encodes the cached frames after a <see cref="Retarget"/>.</summary>
	public enum Phase { Sampling, Encoding, Done, Failed }

	// Sample-run shape: Sites runs of RunLen consecutive-tick deltas, spread evenly across the span.
	private const int Sites = 5;
	private const int RunLen = 4;

	public Phase State { get; private set; } = Phase.Sampling;

	/// <summary>True while <see cref="Step"/> still has work to do.</summary>
	public bool Running => State is Phase.Sampling or Phase.Encoding;

	/// <summary>Estimated GIF size in bytes; valid once <see cref="State"/> is Done.</summary>
	public long ResultBytes { get; private set; } = -1;

	private readonly GameManager _mgr;
	private readonly int _frameCount;
	private readonly int _savedFrame;
	// Capture schedule: ticks to seek+capture in order. Measured captures feed the delta average;
	// unmeasured ones are the first frame (full-frame cost, counted separately) and re-bases between sites.
	private readonly List<(int Tick, bool Measure)> _ticks = new();
	// Native 240px pixels of every captured tick, parallel to _ticks — the input to (re-)encoding.
	private readonly List<Color32[]> _frames = new();

	private int _scale;
	private int _pendingScale; // scale requested mid-sampling; applied once every frame is captured
	private GifEncoder _enc;
	private int _cursor;       // next tick to seek (Sampling) / next cached frame to encode (Encoding)
	private bool _pendingCapture;
	private long _firstFrameBytes; // header + first frame
	private long _deltaTotal;
	private int _deltaCount;

	public GifSizeEstimator( GameManager mgr, int startTick, int endTick, int scale )
	{
		_mgr = mgr;
		_scale = _pendingScale = scale;

		if ( mgr is null || !mgr.IsReplaying || mgr.Camera is null || mgr.ReplayLength <= 0 )
		{
			State = Phase.Failed;
			return;
		}

		int last = Math.Max( 0, mgr.ReplayLength - 1 );
		startTick = Math.Clamp( startTick, 0, last );
		endTick = Math.Clamp( endTick, startTick, last );
		_frameCount = endTick - startTick + 1;
		_savedFrame = mgr.ReplayFrame;
		ResetEncoder();

		// First frame: full-frame cost.
		_ticks.Add( (startTick, false) );

		if ( _frameCount > 1 )
		{
			int sites = Math.Min( Sites, Math.Max( 1, (_frameCount - 1) / RunLen ) );
			int cur = startTick;

			for ( int i = 0; i < sites; i++ )
			{
				// Run start ticks spread over the span, ascending so every seek is a cheap forward re-sim.
				// Site 0 starts at startTick, where the first frame leaves both the replay and the
				// encoder's delta base — no re-base capture needed there.
				int t = sites == 1 ? startTick : startTick + (int)((long)i * (_frameCount - 1 - RunLen) / (sites - 1));
				t = Math.Max( t, cur ); // short spans can make sites overlap — never rewind

				if ( t != cur )
					_ticks.Add( (t, false) ); // re-base the delta; growth not measured

				for ( int r = 0; r < RunLen && t + r + 1 <= endTick; r++ )
					_ticks.Add( (t + r + 1, true) );

				cur = Math.Min( t + RunLen, endTick );
			}
		}
	}

	/// <summary>Advance by one frame. Call once per UI tick; does nothing once finished.</summary>
	public void Step()
	{
		if ( State == Phase.Encoding )
		{
			EncodeFrame( _cursor++ );
			if ( _cursor >= _frames.Count )
				FinishEstimate();
			return;
		}

		if ( State != Phase.Sampling )
			return;

		if ( _mgr is null || !_mgr.IsReplaying || _mgr.Camera is null )
		{
			State = Phase.Failed;
			return;
		}

		try
		{
			// Same one-frame-per-update pipeline as GifExporter.Step: capture the tick seeked last update
			// (its sprite batch has uploaded by now), then seek the next scheduled tick.
			if ( _pendingCapture )
			{
				_pendingCapture = false;
				_frames.Add( CaptureTight() );
				EncodeFrame( _cursor );

				_cursor++;
				if ( _cursor >= _ticks.Count )
				{
					RestoreReplayFrame();
					// A scale picked while sampling ran: the captures are scale-independent, so just
					// re-encode them.
					if ( _pendingScale != _scale )
						BeginEncoding( _pendingScale );
					else
						FinishEstimate();
					return;
				}
			}

			_mgr.SeekReplayToFrame( _ticks[_cursor].Tick, playSfx: false );
			_pendingCapture = true;
		}
		catch ( Exception e )
		{
			Log.Warning( $"BlockParty GIF size estimate failed: {e.Message}" );
			State = Phase.Failed;
			RestoreReplayFrame();
		}
	}

	/// <summary>Re-estimate at a different output scale using the cached sample frames — no seeks, so
	/// the preview doesn't scrub. Takes effect after sampling if that's still running; restarts the
	/// encode if it was Encoding or Done. Ignored once Failed.</summary>
	public void Retarget( int scale )
	{
		if ( State == Phase.Failed || scale == _pendingScale )
			return;

		_pendingScale = scale;
		if ( State != Phase.Sampling )
			BeginEncoding( scale );
	}

	/// <summary>Abandon the job (heads moved, export started, overlay closed) and, if it was still
	/// seeking, put the replay back on the frame it was showing.</summary>
	public void Cancel()
	{
		if ( !Running )
			return;

		bool wasSampling = State == Phase.Sampling;
		State = Phase.Failed;
		if ( wasSampling )
			RestoreReplayFrame();
	}

	private void BeginEncoding( int scale )
	{
		_scale = _pendingScale = scale;
		ResetEncoder();
		_cursor = 0;
		ResultBytes = -1;
		State = Phase.Encoding;
	}

	private void ResetEncoder()
	{
		int size = GifExporter.OutputSizeFor( _scale );
		_enc = new GifEncoder( size, size, 0 );
		_firstFrameBytes = 0;
		_deltaTotal = 0;
		_deltaCount = 0;
	}

	// Feed cached frame i through the encoder at the current scale, booking its cost.
	private void EncodeFrame( int i )
	{
		int before = _enc.Length;
		_enc.AddFrame( GifExporter.Upscale( _frames[i], _scale ), GifExporter.FrameDelayCs );

		if ( i == 0 )
			_firstFrameBytes = _enc.Length; // header + first frame
		else if ( _ticks[i].Measure )
		{
			_deltaTotal += _enc.Length - before;
			_deltaCount++;
		}
	}

	private void FinishEstimate()
	{
		long estimate = _firstFrameBytes + 1; // + trailer byte
		if ( _frameCount > 1 && _deltaCount > 0 )
			estimate += _deltaTotal / _deltaCount * (_frameCount - 1);
		ResultBytes = estimate;
		State = Phase.Done;
	}

	private void RestoreReplayFrame()
	{
		if ( _mgr is not null && _mgr.IsReplaying )
			_mgr.SeekReplayToFrame( _savedFrame, playSfx: false );
	}

	// Capture with the export's tight arena framing, restoring the preview framing before returning so the
	// on-screen camera never shows it (RenderToBitmap reads the camera transform at call time; only the
	// sprite batches are frozen per update).
	private Color32[] CaptureTight()
	{
		var cam = _mgr.Camera;
		var savedPos = cam.WorldPosition;
		float savedOrtho = cam.OrthographicHeight;

		try
		{
			cam.WorldPosition = new Vector3( Arena.WIDTH / 2f, Arena.HEIGHT / 2f, savedPos.z );
			cam.OrthographicHeight = Arena.HEIGHT;
			return GifExporter.CaptureNative( _mgr );
		}
		finally
		{
			cam.WorldPosition = savedPos;
			cam.OrthographicHeight = savedOrtho;
		}
	}
}