using System;
using System.IO;

namespace GmodAssetImporter;

/// <summary>
/// Source1 VTF reader for common GMod formats (BGRA8888, RGBA8888, DXT1, DXT5).
/// </summary>
public sealed class VtfImage
{
	public int Width;
	public int Height;
	public byte[] Rgba; // length = Width * Height * 4
}

public static class VtfReader
{
	const int VtfSignature = 0x00465456; // "VTF\0" little-endian as int? Actually 'V','T','F','\0' = 0x00465456

	enum ImageFormat
	{
		None = -1,
		RGBA8888 = 0,
		ABGR8888 = 1,
		RGB888 = 2,
		BGR888 = 3,
		RGB565 = 4,
		I8 = 5,
		IA88 = 6,
		P8 = 7,
		A8 = 8,
		RGB888_Bluescreen = 9,
		BGR888_Bluescreen = 10,
		ARGB8888 = 11,
		BGRA8888 = 12,
		DXT1 = 13,
		DXT3 = 14,
		DXT5 = 15,
		BGRX8888 = 16,
		BGR565 = 17,
		BGRX5551 = 18,
		BGRA4444 = 19,
		DXT1_ONEBITALPHA = 20,
		BGRA5551 = 21,
		UV88 = 22,
		UVWQ8888 = 23,
		RGBA16161616F = 24,
		RGBA16161616 = 25,
		UVLX8888 = 26
	}

	public static VtfImage Load( string path )
	{
		return Load( File.ReadAllBytes( path ) );
	}

	public static VtfImage Load( byte[] data )
	{
		using var ms = new MemoryStream( data );
		using var br = new BinaryReader( ms );

		var sig = br.ReadBytes( 4 );
		if ( sig.Length < 4 || sig[0] != 'V' || sig[1] != 'T' || sig[2] != 'F' || sig[3] != 0 )
			throw new InvalidDataException( "Not a VTF file" );

		var versionMajor = br.ReadUInt32();
		var versionMinor = br.ReadUInt32();
		var headerSize = br.ReadUInt32();
		var width = br.ReadUInt16();
		var height = br.ReadUInt16();
		br.ReadUInt32(); // flags
		var frameCount = Math.Max( 1, (int)br.ReadUInt16() );
		br.ReadUInt16(); // firstFrame
		ms.Position += 4; // padding0
		ms.Position += 12; // reflectivity
		ms.Position += 4; // padding1
		br.ReadSingle(); // bumpScale
		var highFormat = (ImageFormat)br.ReadInt32();
		var mipCount = Math.Max( 1, (int)br.ReadByte() );
		var lowFormat = (ImageFormat)br.ReadInt32();
		var lowWidth = br.ReadByte();
		var lowHeight = br.ReadByte();

		// Pre-7.3 layout: header, low-res thumbnail, then high-res mips
		uint imageOffset = headerSize;
		if ( lowFormat != ImageFormat.None && lowWidth > 0 && lowHeight > 0 )
			imageOffset += (uint)ComputeFaceSize( lowWidth, lowHeight, lowFormat );

		if ( versionMajor > 7 || (versionMajor == 7 && versionMinor >= 2) )
			br.ReadUInt16(); // depth

		if ( versionMajor > 7 || (versionMajor == 7 && versionMinor >= 3) )
		{
			ms.Position += 3; // padding2
			var numResources = br.ReadInt32();
			// Resource dictionary starts at byte 80 for 7.3+
			ms.Position = 80;
			for ( var r = 0; r < numResources; r++ )
			{
				var typeBytes = br.ReadBytes( 3 );
				br.ReadByte(); // resource flags
				var offset = br.ReadUInt32();
				var type = typeBytes[0] | (typeBytes[1] << 8) | (typeBytes[2] << 16);
				if ( type == 0x30 ) // high-res image
					imageOffset = offset;
			}
		}

		ms.Position = imageOffset;

		// Skip smaller mip levels so we land on the largest mip of frame 0
		long skip = 0;
		for ( var mip = mipCount - 1; mip >= 1; mip-- )
		{
			var mw = Math.Max( 1, width >> mip );
			var mh = Math.Max( 1, height >> mip );
			skip += (long)ComputeFaceSize( mw, mh, highFormat ) * frameCount;
		}

		ms.Position = imageOffset + skip;
		var faceSize = ComputeFaceSize( width, height, highFormat );
		var face = br.ReadBytes( faceSize );
		var rgba = Decode( face, width, height, highFormat );

		return new VtfImage
		{
			Width = width,
			Height = height,
			Rgba = rgba
		};
	}

	static int ComputeFaceSize( int w, int h, ImageFormat format )
	{
		return format switch
		{
			ImageFormat.DXT1 or ImageFormat.DXT1_ONEBITALPHA => Math.Max( 1, (w + 3) / 4 ) * Math.Max( 1, (h + 3) / 4 ) * 8,
			ImageFormat.DXT3 or ImageFormat.DXT5 => Math.Max( 1, (w + 3) / 4 ) * Math.Max( 1, (h + 3) / 4 ) * 16,
			ImageFormat.RGBA8888 or ImageFormat.ABGR8888 or ImageFormat.ARGB8888 or ImageFormat.BGRA8888 or ImageFormat.BGRX8888 => w * h * 4,
			ImageFormat.RGB888 or ImageFormat.BGR888 => w * h * 3,
			ImageFormat.I8 or ImageFormat.A8 => w * h,
			ImageFormat.IA88 => w * h * 2,
			_ => w * h * 4
		};
	}

	static byte[] Decode( byte[] src, int width, int height, ImageFormat format )
	{
		return format switch
		{
			ImageFormat.RGBA8888 => DecodeRgba8888( src, width, height ),
			ImageFormat.BGRA8888 or ImageFormat.BGRX8888 => DecodeBgra8888( src, width, height ),
			ImageFormat.ARGB8888 => DecodeArgb8888( src, width, height ),
			ImageFormat.ABGR8888 => DecodeAbgr8888( src, width, height ),
			ImageFormat.RGB888 => DecodeRgb888( src, width, height, bgr: false ),
			ImageFormat.BGR888 => DecodeRgb888( src, width, height, bgr: true ),
			ImageFormat.DXT1 or ImageFormat.DXT1_ONEBITALPHA => DecodeDxt1( src, width, height ),
			ImageFormat.DXT5 => DecodeDxt5( src, width, height ),
			ImageFormat.DXT3 => DecodeDxt3( src, width, height ),
			ImageFormat.I8 => DecodeI8( src, width, height ),
			_ => throw new NotSupportedException( $"Unsupported VTF format: {format}" )
		};
	}

	static byte[] DecodeRgba8888( byte[] src, int w, int h )
	{
		var dst = new byte[w * h * 4];
		Buffer.BlockCopy( src, 0, dst, 0, Math.Min( src.Length, dst.Length ) );
		return dst;
	}

	static byte[] DecodeBgra8888( byte[] src, int w, int h )
	{
		var dst = new byte[w * h * 4];
		var n = Math.Min( src.Length / 4, w * h );
		for ( var i = 0; i < n; i++ )
		{
			var o = i * 4;
			dst[o] = src[o + 2];
			dst[o + 1] = src[o + 1];
			dst[o + 2] = src[o];
			dst[o + 3] = src[o + 3];
		}
		return dst;
	}

	static byte[] DecodeArgb8888( byte[] src, int w, int h )
	{
		var dst = new byte[w * h * 4];
		var n = Math.Min( src.Length / 4, w * h );
		for ( var i = 0; i < n; i++ )
		{
			var o = i * 4;
			dst[o] = src[o + 1];
			dst[o + 1] = src[o + 2];
			dst[o + 2] = src[o + 3];
			dst[o + 3] = src[o];
		}
		return dst;
	}

	static byte[] DecodeAbgr8888( byte[] src, int w, int h )
	{
		var dst = new byte[w * h * 4];
		var n = Math.Min( src.Length / 4, w * h );
		for ( var i = 0; i < n; i++ )
		{
			var o = i * 4;
			dst[o] = src[o + 3];
			dst[o + 1] = src[o + 2];
			dst[o + 2] = src[o + 1];
			dst[o + 3] = src[o];
		}
		return dst;
	}

	static byte[] DecodeRgb888( byte[] src, int w, int h, bool bgr )
	{
		var dst = new byte[w * h * 4];
		var n = Math.Min( src.Length / 3, w * h );
		for ( var i = 0; i < n; i++ )
		{
			var s = i * 3;
			var d = i * 4;
			var r = bgr ? src[s + 2] : src[s];
			var g = src[s + 1];
			var b = bgr ? src[s] : src[s + 2];
			dst[d] = r;
			dst[d + 1] = g;
			dst[d + 2] = b;
			dst[d + 3] = 255;
		}
		return dst;
	}

	static byte[] DecodeI8( byte[] src, int w, int h )
	{
		var dst = new byte[w * h * 4];
		var n = Math.Min( src.Length, w * h );
		for ( var i = 0; i < n; i++ )
		{
			var v = src[i];
			var d = i * 4;
			dst[d] = v;
			dst[d + 1] = v;
			dst[d + 2] = v;
			dst[d + 3] = 255;
		}
		return dst;
	}

	static byte[] DecodeDxt1( byte[] src, int w, int h )
	{
		var dst = new byte[w * h * 4];
		var blocksX = Math.Max( 1, (w + 3) / 4 );
		var blocksY = Math.Max( 1, (h + 3) / 4 );
		var offset = 0;
		for ( var by = 0; by < blocksY; by++ )
		for ( var bx = 0; bx < blocksX; bx++ )
		{
			if ( offset + 8 > src.Length ) return dst;
			DecodeDxt1Block( src, offset, bx * 4, by * 4, w, h, dst, oneBitAlpha: false );
			offset += 8;
		}
		return dst;
	}

	static byte[] DecodeDxt3( byte[] src, int w, int h )
	{
		var dst = new byte[w * h * 4];
		var blocksX = Math.Max( 1, (w + 3) / 4 );
		var blocksY = Math.Max( 1, (h + 3) / 4 );
		var offset = 0;
		for ( var by = 0; by < blocksY; by++ )
		for ( var bx = 0; bx < blocksX; bx++ )
		{
			if ( offset + 16 > src.Length ) return dst;
			var alphaBits = BitConverter.ToUInt64( src, offset );
			DecodeDxt1Block( src, offset + 8, bx * 4, by * 4, w, h, dst, oneBitAlpha: false, alwaysFourColor: true );
			for ( var py = 0; py < 4; py++ )
			for ( var px = 0; px < 4; px++ )
			{
				var x = bx * 4 + px;
				var y = by * 4 + py;
				if ( x >= w || y >= h ) continue;
				var aIdx = py * 4 + px;
				var a = (int)((alphaBits >> (aIdx * 4)) & 0xF);
				dst[(y * w + x) * 4 + 3] = (byte)((a * 255) / 15);
			}
			offset += 16;
		}
		return dst;
	}

	static byte[] DecodeDxt5( byte[] src, int w, int h )
	{
		var dst = new byte[w * h * 4];
		var blocksX = Math.Max( 1, (w + 3) / 4 );
		var blocksY = Math.Max( 1, (h + 3) / 4 );
		var offset = 0;
		for ( var by = 0; by < blocksY; by++ )
		for ( var bx = 0; bx < blocksX; bx++ )
		{
			if ( offset + 16 > src.Length ) return dst;
			var a0 = src[offset];
			var a1 = src[offset + 1];
			ulong aBits = src[offset + 2]
				| ((ulong)src[offset + 3] << 8)
				| ((ulong)src[offset + 4] << 16)
				| ((ulong)src[offset + 5] << 24)
				| ((ulong)src[offset + 6] << 32)
				| ((ulong)src[offset + 7] << 40);

			DecodeDxt1Block( src, offset + 8, bx * 4, by * 4, w, h, dst, oneBitAlpha: false, alwaysFourColor: true );

			var alphas = BuildDxt5AlphaPalette( a0, a1 );
			for ( var py = 0; py < 4; py++ )
			for ( var px = 0; px < 4; px++ )
			{
				var x = bx * 4 + px;
				var y = by * 4 + py;
				if ( x >= w || y >= h ) continue;
				var idx = py * 4 + px;
				var code = (int)((aBits >> (idx * 3)) & 0x7);
				dst[(y * w + x) * 4 + 3] = alphas[code];
			}
			offset += 16;
		}
		return dst;
	}

	static byte[] BuildDxt5AlphaPalette( byte a0, byte a1 )
	{
		var p = new byte[8];
		p[0] = a0;
		p[1] = a1;
		if ( a0 > a1 )
		{
			for ( var i = 1; i <= 6; i++ )
				p[i + 1] = (byte)(((7 - i) * a0 + i * a1) / 7);
		}
		else
		{
			for ( var i = 1; i <= 4; i++ )
				p[i + 1] = (byte)(((5 - i) * a0 + i * a1) / 5);
			p[6] = 0;
			p[7] = 255;
		}
		return p;
	}

	/// <param name="alwaysFourColor">DXT3/DXT5 colour blocks ignore the c0 &lt;= c1 three-colour mode.</param>
	static void DecodeDxt1Block( byte[] src, int offset, int blockX, int blockY, int w, int h, byte[] dst, bool oneBitAlpha, bool alwaysFourColor = false )
	{
		var c0 = BitConverter.ToUInt16( src, offset );
		var c1 = BitConverter.ToUInt16( src, offset + 2 );
		var indices = BitConverter.ToUInt32( src, offset + 4 );
		var colors = new byte[4 * 4];
		Rgb565( c0, out colors[0], out colors[1], out colors[2] );
		colors[3] = 255;
		Rgb565( c1, out colors[4], out colors[5], out colors[6] );
		colors[7] = 255;

		if ( c0 > c1 || alwaysFourColor )
		{
			for ( var i = 0; i < 3; i++ )
			{
				colors[8 + i] = (byte)((2 * colors[i] + colors[4 + i]) / 3);
				colors[12 + i] = (byte)((colors[i] + 2 * colors[4 + i]) / 3);
			}
			colors[11] = 255;
			colors[15] = 255;
		}
		else
		{
			for ( var i = 0; i < 3; i++ )
			{
				colors[8 + i] = (byte)((colors[i] + colors[4 + i]) / 2);
				colors[12 + i] = 0;
			}
			colors[11] = 255;
			colors[15] = oneBitAlpha ? (byte)0 : (byte)255;
		}

		for ( var py = 0; py < 4; py++ )
		for ( var px = 0; px < 4; px++ )
		{
			var x = blockX + px;
			var y = blockY + py;
			if ( x >= w || y >= h ) continue;
			var code = (int)((indices >> (2 * (py * 4 + px))) & 0x3);
			var di = (y * w + x) * 4;
			var ci = code * 4;
			dst[di] = colors[ci];
			dst[di + 1] = colors[ci + 1];
			dst[di + 2] = colors[ci + 2];
			dst[di + 3] = colors[ci + 3];
		}
	}

	static void Rgb565( ushort c, out byte r, out byte g, out byte b )
	{
		var ri = (c >> 11) & 0x1F;
		var gi = (c >> 5) & 0x3F;
		var bi = c & 0x1F;
		r = (byte)((ri << 3) | (ri >> 2));
		g = (byte)((gi << 2) | (gi >> 4));
		b = (byte)((bi << 3) | (bi >> 2));
	}
}