using System;
using System.Collections.Generic;
using System.Linq;

namespace BlockParty;

/// <summary>
/// Pure C# GIF89a encoder with LZW compression and animation support. Writes animated GIFs with per-frame
/// local color tables and transparency.
///
/// Frames are delta-encoded: after the first frame, only the bounding rectangle of pixels that changed
/// since the previous frame is written, with unchanged pixels inside it marked transparent ("keep the
/// previous pixel", disposal method 1). Static backgrounds therefore cost almost nothing per frame —
/// the same trick gif optimizers (gifsicle etc.) apply, ~10x smaller on gameplay footage. Source frames
/// that themselves contain transparent pixels can't be delta-encoded (transparency already means
/// "keep previous" in a delta) and fall back to whole-frame writes.
///
/// Incremental: each frame is quantized + LZW-compressed and appended the moment it's handed to
/// <see cref="AddFrame"/>, so raw frames are never buffered — memory stays ~the size of the growing
/// compressed output plus one retained copy of the previous frame for delta comparison.
/// (Construct, AddFrame per frame, Finish.)
///
/// Ported from the PixelShitter editor's encoder, but rewritten to build into a plain <c>byte[]</c> buffer
/// instead of a <see cref="System.IO.Stream"/>/<c>BinaryWriter</c> — the game runs in the sbox sandbox,
/// which keeps gameplay code off System.IO streams (see <see cref="RunInputCodec"/>).
/// </summary>
public sealed class GifEncoder
{
	private readonly ByteBuffer _w;
	private readonly int _width;
	private readonly int _height;

	// Previous frame's pixels, kept for delta encoding. Null until an opaque frame has been written
	// (and reset by whole-frame fallbacks, which invalidate what's on the canvas for diffing).
	private Color32[] _prev;

	// 1×1 transparent frame body: emitted when a frame is identical to the previous one, purely to carry its delay.
	private static readonly Color32[] NoChangePixel = { new Color32( 0, 0, 0, 0 ) };

	/// <summary>Begin an animated GIF of the given size. <paramref name="loopCount"/>: 0 = infinite loop,
	/// 1 = play once, N = loop N times.</summary>
	public GifEncoder( int width, int height, int loopCount = 0 )
	{
		_width = width;
		_height = height;
		_w = new ByteBuffer( 1 << 16 );

		// ── Header ──
		_w.WriteAscii( "GIF89a" );

		// ── Logical Screen Descriptor ── (no global color table — we use local tables per frame)
		_w.WriteUInt16( (ushort)width );
		_w.WriteUInt16( (ushort)height );
		_w.WriteByte( 0x70 ); // packed: no GCT; color resolution 8-bit (bits 4-6 — tools report this as the file's bit depth)
		_w.WriteByte( 0x00 ); // background color index
		_w.WriteByte( 0x00 ); // pixel aspect ratio

		// ── Netscape Application Extension (looping) ──
		if ( loopCount != 1 )
		{
			_w.WriteByte( 0x21 ); // extension introducer
			_w.WriteByte( 0xFF ); // application extension
			_w.WriteByte( 11 );   // block size
			_w.WriteAscii( "NETSCAPE2.0" );
			_w.WriteByte( 3 );    // sub-block size
			_w.WriteByte( 1 );    // sub-block ID
			_w.WriteUInt16( (ushort)(loopCount == 0 ? 0 : loopCount - 1) );
			_w.WriteByte( 0 );    // block terminator
		}
	}

	/// <summary>Quantize and append one frame. <paramref name="pixels"/> is RGBA, row-major top-down, length
	/// = width * height. The pixel array isn't retained past this call.</summary>
	public void AddFrame( Color32[] pixels, int delayCentiseconds )
	{
		// Source transparency collides with the delta's keep-previous transparency — write those whole,
		// with the pre-delta disposal (restore to background), and restart diffing on the next opaque frame.
		if ( HasTransparentPixels( pixels ) )
		{
			WriteFrame( _w, 0, 0, _width, _height, pixels, delayCentiseconds, disposal: 2 );
			_prev = null;
			return;
		}

		if ( _prev is null )
		{
			WriteFrame( _w, 0, 0, _width, _height, pixels, delayCentiseconds, disposal: 1 );
			_prev = new Color32[pixels.Length];
			Array.Copy( pixels, _prev, pixels.Length );
			return;
		}

		if ( !ComputeDiffRect( pixels, _prev, _width, _height, out int left, out int top, out int rectW, out int rectH ) )
		{
			// Identical frame — emit a 1×1 transparent frame just to carry the delay.
			WriteFrame( _w, 0, 0, 1, 1, NoChangePixel, delayCentiseconds, disposal: 1 );
			return;
		}

		// Delta frame: the changed bounding rect, with still-matching pixels inside it transparent
		// so they keep the previous frame's color (and LZW-compress to almost nothing).
		var rect = new Color32[rectW * rectH];
		for ( int y = 0; y < rectH; y++ )
		{
			int src = (top + y) * _width + left;
			int dst = y * rectW;
			for ( int x = 0; x < rectW; x++, src++, dst++ )
			{
				var px = pixels[src];
				var pv = _prev[src];
				rect[dst] = (px.r == pv.r && px.g == pv.g && px.b == pv.b) ? default : px;
			}
		}

		WriteFrame( _w, left, top, rectW, rectH, rect, delayCentiseconds, disposal: 1 );
		Array.Copy( pixels, _prev, pixels.Length );
	}

	private static bool HasTransparentPixels( Color32[] pixels )
	{
		foreach ( var px in pixels )
			if ( px.a < 128 )
				return true;
		return false;
	}

	/// <summary>Bounding rectangle of pixels whose RGB differs between the two frames.
	/// Returns false when the frames are identical.</summary>
	private static bool ComputeDiffRect( Color32[] cur, Color32[] prev, int width, int height,
		out int left, out int top, out int rectW, out int rectH )
	{
		int minX = width, minY = height, maxX = -1, maxY = -1;
		int i = 0;
		for ( int y = 0; y < height; y++ )
		{
			for ( int x = 0; x < width; x++, i++ )
			{
				if ( cur[i].r == prev[i].r && cur[i].g == prev[i].g && cur[i].b == prev[i].b )
					continue;
				if ( x < minX ) minX = x;
				if ( x > maxX ) maxX = x;
				if ( y < minY ) minY = y;
				if ( y > maxY ) maxY = y;
			}
		}

		left = minX;
		top = minY;
		rectW = maxX - minX + 1;
		rectH = maxY - minY + 1;
		return maxX >= 0;
	}

	/// <summary>Bytes written so far (header + frames appended to date). Lets callers meter how much each
	/// AddFrame cost — the size estimator samples a few frames and extrapolates from the growth.</summary>
	public int Length => _w.Count;

	/// <summary>Write the trailer and return the finished GIF bytes.</summary>
	public byte[] Finish()
	{
		_w.WriteByte( 0x3B );
		return _w.ToArray();
	}

	// ─────────────────────────────────────────────────────────────────
	// Frame writing
	// ─────────────────────────────────────────────────────────────────

	private static void WriteFrame( ByteBuffer w, int left, int top, int width, int height, Color32[] pixels,
		int delayCentiseconds, int disposal )
	{
		// Quantize RGBA → palette + indices
		Quantize( pixels, out var palette, out var indices, out int transparentIndex );

		int paletteBits = PaletteBits( palette.Length );
		int paletteSize = 1 << paletteBits;

		// ── Graphic Control Extension ──
		w.WriteByte( 0x21 ); // extension introducer
		w.WriteByte( 0xF9 ); // graphic control label
		w.WriteByte( 4 );    // block size

		// Packed: disposal method (1 = do not dispose for delta frames, 2 = restore to bg), transparent flag
		byte packed = (byte)((disposal & 0x7) << 2);
		if ( transparentIndex >= 0 )
			packed |= 0x01; // transparent color flag
		w.WriteByte( packed );

		int delay = Math.Max( delayCentiseconds, 2 ); // browsers ignore delay < 2cs
		w.WriteUInt16( (ushort)delay );
		w.WriteByte( (byte)(transparentIndex >= 0 ? transparentIndex : 0) );
		w.WriteByte( 0 ); // block terminator

		// ── Image Descriptor ──
		w.WriteByte( 0x2C ); // image separator
		w.WriteUInt16( (ushort)left );
		w.WriteUInt16( (ushort)top );
		w.WriteUInt16( (ushort)width );
		w.WriteUInt16( (ushort)height );

		// Packed: local color table flag + size
		w.WriteByte( (byte)(0x80 | (paletteBits - 1)) ); // LCT flag + LCT size

		// ── Local Color Table ──
		for ( int i = 0; i < paletteSize; i++ )
		{
			if ( i < palette.Length )
			{
				w.WriteByte( palette[i].r );
				w.WriteByte( palette[i].g );
				w.WriteByte( palette[i].b );
			}
			else
			{
				w.WriteByte( 0 );
				w.WriteByte( 0 );
				w.WriteByte( 0 );
			}
		}

		// ── LZW Image Data ──
		int minCodeSize = Math.Max( paletteBits, 2 ); // GIF minimum is 2
		w.WriteByte( (byte)minCodeSize );
		WriteLzwData( w, indices, minCodeSize );
		w.WriteByte( 0 ); // block terminator
	}

	// ─────────────────────────────────────────────────────────────────
	// Color quantization
	// ─────────────────────────────────────────────────────────────────

	/// <summary>
	/// Quantize RGBA pixels to a palette of up to 256 colors. Index 0 is reserved for transparency if any
	/// pixel has alpha below 128. Uses a simple popularity-based approach (most frequent colors first), which
	/// works very well for pixel art with few unique colors; falls back to median-cut when there are many.
	/// </summary>
	private static void Quantize( Color32[] pixels, out Color32[] palette, out byte[] indices, out int transparentIndex )
	{
		transparentIndex = -1;
		bool hasTransparency = false;

		// Count unique opaque colors
		var colorCounts = new Dictionary<uint, int>();
		foreach ( var px in pixels )
		{
			if ( px.a < 128 )
			{
				hasTransparency = true;
				continue;
			}
			uint key = PackRgb( px );
			colorCounts.TryGetValue( key, out int count );
			colorCounts[key] = count + 1;
		}

		int maxColors = hasTransparency ? 255 : 256;
		transparentIndex = hasTransparency ? 0 : -1;
		int paletteOffset = hasTransparency ? 1 : 0;

		// Build palette from most popular colors
		Color32[] uniqueColors;
		if ( colorCounts.Count <= maxColors )
		{
			// All colors fit — perfect quantization
			uniqueColors = colorCounts.Keys.Select( UnpackRgb ).ToArray();
		}
		else
		{
			// Too many colors — use median-cut
			uniqueColors = MedianCut( colorCounts, maxColors );
		}

		// Build final palette
		palette = new Color32[paletteOffset + uniqueColors.Length];
		if ( hasTransparency )
			palette[0] = new Color32( 0, 0, 0, 0 );
		for ( int i = 0; i < uniqueColors.Length; i++ )
			palette[paletteOffset + i] = uniqueColors[i];

		// Build a fast lookup map
		var colorToIndex = new Dictionary<uint, byte>( palette.Length );
		for ( int i = paletteOffset; i < palette.Length; i++ )
			colorToIndex[PackRgb( palette[i] )] = (byte)i;

		// Map pixels to palette indices
		indices = new byte[pixels.Length];
		for ( int i = 0; i < pixels.Length; i++ )
		{
			if ( pixels[i].a < 128 )
			{
				indices[i] = (byte)transparentIndex;
				continue;
			}

			uint key = PackRgb( pixels[i] );
			if ( colorToIndex.TryGetValue( key, out byte idx ) )
				indices[i] = idx;
			else
				indices[i] = FindNearest( palette, pixels[i], paletteOffset ); // nearest palette color
		}
	}

	private static uint PackRgb( Color32 c ) => (uint)(c.r << 16 | c.g << 8 | c.b);
	private static Color32 UnpackRgb( uint v ) => new Color32( (byte)(v >> 16), (byte)(v >> 8 & 0xFF), (byte)(v & 0xFF), 255 );

	private static byte FindNearest( Color32[] palette, Color32 target, int startIndex )
	{
		int bestDist = int.MaxValue;
		byte bestIdx = (byte)startIndex;
		for ( int i = startIndex; i < palette.Length; i++ )
		{
			int dr = palette[i].r - target.r;
			int dg = palette[i].g - target.g;
			int db = palette[i].b - target.b;
			int dist = dr * dr + dg * dg + db * db;
			if ( dist < bestDist )
			{
				bestDist = dist;
				bestIdx = (byte)i;
				if ( dist == 0 ) break;
			}
		}
		return bestIdx;
	}

	// ─────────────────────────────────────────────────────────────────
	// Median-cut quantization
	// ─────────────────────────────────────────────────────────────────

	private static Color32[] MedianCut( Dictionary<uint, int> colorCounts, int maxColors )
	{
		// Build initial box of all colors with their counts
		var allColors = colorCounts.Select( kv => (color: UnpackRgb( kv.Key ), count: kv.Value) ).ToList();
		var boxes = new List<List<(Color32 color, int count)>> { allColors };

		while ( boxes.Count < maxColors )
		{
			// Find the box with the largest color range
			int bestBox = 0;
			int bestRange = -1;
			for ( int b = 0; b < boxes.Count; b++ )
			{
				if ( boxes[b].Count <= 1 ) continue;
				int range = BoxRange( boxes[b] );
				if ( range > bestRange )
				{
					bestRange = range;
					bestBox = b;
				}
			}

			if ( bestRange <= 0 ) break;

			// Split the box along its widest channel
			var box = boxes[bestBox];
			int channel = WidestChannel( box );
			box.Sort( ( a, b ) => ChannelValue( a.color, channel ).CompareTo( ChannelValue( b.color, channel ) ) );

			int totalCount = box.Sum( c => c.count );
			int halfCount = totalCount / 2;
			int cumulative = 0;
			int splitAt = 0;
			for ( int i = 0; i < box.Count; i++ )
			{
				cumulative += box[i].count;
				if ( cumulative >= halfCount )
				{
					splitAt = Math.Max( 1, i );
					break;
				}
			}

			boxes[bestBox] = box.GetRange( 0, splitAt );
			boxes.Add( box.GetRange( splitAt, box.Count - splitAt ) );
		}

		// Average each box to produce a palette color
		return boxes.Select( box =>
		{
			long tr = 0, tg = 0, tb = 0, total = 0;
			foreach ( var (color, count) in box )
			{
				tr += color.r * count;
				tg += color.g * count;
				tb += color.b * count;
				total += count;
			}
			return new Color32( (byte)(tr / total), (byte)(tg / total), (byte)(tb / total), 255 );
		} ).ToArray();
	}

	private static int BoxRange( List<(Color32 color, int count)> box )
	{
		int rMin = 255, rMax = 0, gMin = 255, gMax = 0, bMin = 255, bMax = 0;
		foreach ( var (c, _) in box )
		{
			if ( c.r < rMin ) rMin = c.r; if ( c.r > rMax ) rMax = c.r;
			if ( c.g < gMin ) gMin = c.g; if ( c.g > gMax ) gMax = c.g;
			if ( c.b < bMin ) bMin = c.b; if ( c.b > bMax ) bMax = c.b;
		}
		return Math.Max( rMax - rMin, Math.Max( gMax - gMin, bMax - bMin ) );
	}

	private static int WidestChannel( List<(Color32 color, int count)> box )
	{
		int rMin = 255, rMax = 0, gMin = 255, gMax = 0, bMin = 255, bMax = 0;
		foreach ( var (c, _) in box )
		{
			if ( c.r < rMin ) rMin = c.r; if ( c.r > rMax ) rMax = c.r;
			if ( c.g < gMin ) gMin = c.g; if ( c.g > gMax ) gMax = c.g;
			if ( c.b < bMin ) bMin = c.b; if ( c.b > bMax ) bMax = c.b;
		}
		int rRange = rMax - rMin, gRange = gMax - gMin, bRange = bMax - bMin;
		if ( rRange >= gRange && rRange >= bRange ) return 0;
		if ( gRange >= bRange ) return 1;
		return 2;
	}

	private static int ChannelValue( Color32 c, int channel ) => channel switch
	{
		0 => c.r,
		1 => c.g,
		_ => c.b,
	};

	// ─────────────────────────────────────────────────────────────────
	// LZW compression
	// ─────────────────────────────────────────────────────────────────

	private static void WriteLzwData( ByteBuffer w, byte[] indices, int minCodeSize )
	{
		int clearCode = 1 << minCodeSize;
		int eoiCode = clearCode + 1;
		int nextCode = eoiCode + 1;
		int codeSize = minCodeSize + 1;
		int maxCode = (1 << codeSize) - 1;

		// String table: maps (prefix code, suffix byte) → code.
		var table = new Dictionary<long, int>( 4096 );

		var bitBuffer = new LzwBitBuffer();

		// Emit clear code
		bitBuffer.WriteBits( clearCode, codeSize );

		void ResetTable()
		{
			table.Clear();
			nextCode = eoiCode + 1;
			codeSize = minCodeSize + 1;
			maxCode = (1 << codeSize) - 1;
		}

		if ( indices.Length == 0 )
		{
			bitBuffer.WriteBits( eoiCode, codeSize );
			bitBuffer.FlushToSubBlocks( w );
			return;
		}

		int prefix = indices[0];

		for ( int i = 1; i < indices.Length; i++ )
		{
			byte suffix = indices[i];
			long key = ((long)prefix << 8) | suffix;

			if ( table.TryGetValue( key, out int existingCode ) )
			{
				prefix = existingCode;
			}
			else
			{
				// Emit the code for the prefix
				bitBuffer.WriteBits( prefix, codeSize );

				// Add new entry if table isn't full
				if ( nextCode <= 4095 )
				{
					table[key] = nextCode++;
					if ( nextCode > maxCode + 1 && codeSize < 12 )
					{
						codeSize++;
						maxCode = (1 << codeSize) - 1;
					}
				}
				else
				{
					// Table is full — emit clear code and reset
					bitBuffer.WriteBits( clearCode, codeSize );
					ResetTable();
				}

				prefix = suffix;
			}
		}

		// Emit final prefix code and EOI
		bitBuffer.WriteBits( prefix, codeSize );
		bitBuffer.WriteBits( eoiCode, codeSize );

		bitBuffer.FlushToSubBlocks( w );
	}

	/// <summary>
	/// Accumulates LZW codes as a bit stream and flushes them as GIF sub-blocks (max 255 bytes each).
	/// </summary>
	private sealed class LzwBitBuffer
	{
		private readonly List<byte> _bytes = new( 4096 );
		private int _currentBits;
		private int _bitCount;

		public void WriteBits( int code, int codeSize )
		{
			_currentBits |= code << _bitCount;
			_bitCount += codeSize;

			while ( _bitCount >= 8 )
			{
				_bytes.Add( (byte)(_currentBits & 0xFF) );
				_currentBits >>= 8;
				_bitCount -= 8;
			}
		}

		public void FlushToSubBlocks( ByteBuffer w )
		{
			// Flush remaining bits
			if ( _bitCount > 0 )
				_bytes.Add( (byte)(_currentBits & 0xFF) );

			// Write as GIF sub-blocks (max 255 bytes per block)
			int offset = 0;
			while ( offset < _bytes.Count )
			{
				int blockSize = Math.Min( 255, _bytes.Count - offset );
				w.WriteByte( (byte)blockSize );
				for ( int i = 0; i < blockSize; i++ )
					w.WriteByte( _bytes[offset + i] );
				offset += blockSize;
			}
		}
	}

	// ─────────────────────────────────────────────────────────────────
	// Helpers
	// ─────────────────────────────────────────────────────────────────

	private static int PaletteBits( int colorCount )
	{
		// GIF requires palette size to be a power of 2 (min 2 = 2^1)
		int bits = 1;
		while ( (1 << bits) < colorCount )
			bits++;
		return Math.Min( bits, 8 );
	}

	/// <summary>A growable little-endian byte sink — the sandbox-safe stand-in for BinaryWriter/Stream.</summary>
	private sealed class ByteBuffer
	{
		private readonly List<byte> _bytes;

		public ByteBuffer( int capacity ) => _bytes = new List<byte>( capacity );

		public int Count => _bytes.Count;

		public void WriteByte( byte b ) => _bytes.Add( b );

		public void WriteUInt16( ushort value )
		{
			_bytes.Add( (byte)(value & 0xFF) );
			_bytes.Add( (byte)(value >> 8) );
		}

		public void WriteAscii( string s )
		{
			foreach ( char c in s )
				_bytes.Add( (byte)c );
		}

		public byte[] ToArray() => _bytes.ToArray();
	}
}