using System;
using System.Collections.Generic;
using System.IO;
using System.IO.Compression;
using System.Text;

namespace Editor.TerrainConvert;

/// <summary>
/// Minimal OpenEXR reader for scanline images. Source 2's own EXR loader is internal to
/// the engine and can't be called from library code, so this decodes the float data directly.
///
/// Supports: scanline (non-tiled) images, HALF/FLOAT/UINT channels, and NONE/ZIP/ZIPS/RLE
/// compression - which covers the vast majority of heightmap exports. Tiled, deep, and the
/// lossy/wavelet codecs (PIZ, PXR24, B44, DWAA/DWAB) are not supported and throw a clear error.
/// </summary>
public sealed class ExrImage
{
	public int Width { get; private set; }
	public int Height { get; private set; }

	/// <summary> Channel name -> linear float plane (row-major, Width*Height). </summary>
	readonly Dictionary<string, float[]> channels = new( StringComparer.Ordinal );

	public IReadOnlyCollection<string> ChannelNames => channels.Keys;
	public bool HasChannel( string name ) => channels.ContainsKey( name );

	enum PixelType { Uint = 0, Half = 1, Float = 2 }

	enum Compression { None = 0, Rle = 1, Zips = 2, Zip = 3, Piz = 4, Pxr24 = 5, B44 = 6, B44A = 7, Dwaa = 8, Dwab = 9 }

	readonly record struct ChannelInfo( string Name, PixelType Type )
	{
		public int SampleSize => Type == PixelType.Half ? 2 : 4;
	}

	/// <summary>
	/// Returns true if the bytes start with the OpenEXR magic number.
	/// </summary>
	public static bool IsExr( byte[] data )
		=> data is { Length: >= 4 } && data[0] == 0x76 && data[1] == 0x2f && data[2] == 0x31 && data[3] == 0x01;

	/// <summary>
	/// Decode an EXR from a byte buffer.
	/// </summary>
	public static ExrImage Load( byte[] data )
	{
		if ( !IsExr( data ) )
			throw new InvalidDataException( "Not an OpenEXR file (bad magic number)." );

		var img = new ExrImage();
		img.Parse( data );
		return img;
	}

	/// <summary>
	/// Get a channel as a float plane, or null if absent. Common heightmap names are "Y"
	/// (luminance), "R", or a plain unnamed channel.
	/// </summary>
	public float[] GetChannel( string name ) => channels.TryGetValue( name, out var v ) ? v : null;

	/// <summary>
	/// Pick the most sensible single channel to treat as height: explicit preference first,
	/// then Y / R / the first available channel.
	/// </summary>
	public float[] GetHeightChannel( string preferred = null )
	{
		if ( preferred is not null && channels.TryGetValue( preferred, out var p ) )
			return p;

		foreach ( var name in new[] { "Y", "R", "G", "B" } )
			if ( channels.TryGetValue( name, out var v ) )
				return v;

		foreach ( var v in channels.Values )
			return v;

		return null;
	}

	void Parse( byte[] data )
	{
		int pos = 4; // skip magic

		int version = ReadInt32( data, ref pos );
		bool tiled = (version & 0x200) != 0;
		bool deep = (version & 0x800) != 0;
		bool multipart = (version & 0x1000) != 0;

		if ( tiled ) throw new NotSupportedException( "Tiled EXR images are not supported - re-export as a scanline image." );
		if ( deep ) throw new NotSupportedException( "Deep EXR images are not supported." );
		if ( multipart ) throw new NotSupportedException( "Multi-part EXR images are not supported." );

		var chans = new List<ChannelInfo>();
		var compression = Compression.None;
		int xMin = 0, yMin = 0, xMax = 0, yMax = 0;

		// Header attributes, terminated by an empty name.
		while ( true )
		{
			string name = ReadNullString( data, ref pos );
			if ( name.Length == 0 ) break;

			string type = ReadNullString( data, ref pos );
			int size = ReadInt32( data, ref pos );
			int valueStart = pos;

			switch ( name )
			{
				case "channels":
					ParseChannels( data, valueStart, chans );
					break;
				case "compression":
					compression = (Compression)data[valueStart];
					break;
				case "dataWindow":
					int p = valueStart;
					xMin = ReadInt32( data, ref p );
					yMin = ReadInt32( data, ref p );
					xMax = ReadInt32( data, ref p );
					yMax = ReadInt32( data, ref p );
					break;
			}

			pos = valueStart + size; // skip to next attribute
		}

		if ( chans.Count == 0 )
			throw new InvalidDataException( "EXR has no channels." );

		Width = xMax - xMin + 1;
		Height = yMax - yMin + 1;

		if ( Width <= 0 || Height <= 0 )
			throw new InvalidDataException( $"EXR has invalid data window ({Width}x{Height})." );

		int linesPerBlock = compression switch
		{
			Compression.None or Compression.Rle or Compression.Zips => 1,
			Compression.Zip or Compression.Pxr24 => 16,
			Compression.Piz or Compression.B44 or Compression.B44A or Compression.Dwaa => 32,
			Compression.Dwab => 256,
			_ => 1,
		};

		if ( compression is Compression.Piz or Compression.Pxr24 or Compression.B44 or Compression.B44A or Compression.Dwaa or Compression.Dwab )
			throw new NotSupportedException( $"EXR compression '{compression}' is not supported. Re-export the heightmap as Uncompressed, ZIP, or RLE." );

		// Allocate channel planes.
		foreach ( var c in chans )
			channels[c.Name] = new float[Width * Height];

		// Scanline offset table: one ulong per block.
		int blockCount = (Height + linesPerBlock - 1) / linesPerBlock;
		var offsets = new long[blockCount];
		for ( int i = 0; i < blockCount; i++ )
			offsets[i] = (long)ReadUInt64( data, ref pos );

		int rowBytes = 0;
		foreach ( var c in chans )
			rowBytes += Width * c.SampleSize;

		// Each block: int32 yStart, int32 dataSize, then (compressed) pixel data.
		foreach ( var off in offsets )
		{
			int bp = (int)off;
			int yStart = ReadInt32( data, ref bp );
			int dataSize = ReadInt32( data, ref bp );

			int lines = Math.Min( linesPerBlock, yMax - yStart + 1 );
			int uncompressedSize = rowBytes * lines;

			byte[] block;
			if ( compression == Compression.None || dataSize >= uncompressedSize )
			{
				// Stored uncompressed (NONE, or a block that didn't compress smaller).
				block = new byte[uncompressedSize];
				Array.Copy( data, bp, block, 0, uncompressedSize );
			}
			else
			{
				block = Decompress( compression, data, bp, dataSize, uncompressedSize );
			}

			ScatterBlock( block, chans, rowBytes, Width, yStart - yMin, lines );
		}
	}

	void ScatterBlock( byte[] block, List<ChannelInfo> chans, int rowBytes, int width, int rowOffset, int lines )
	{
		for ( int i = 0; i < lines; i++ )
		{
			int rowBase = i * rowBytes;
			int channelOffset = 0;
			int destRow = rowOffset + i;

			foreach ( var c in chans )
			{
				int src = rowBase + channelOffset;
				var plane = channels[c.Name];
				int destBase = destRow * width;

				for ( int x = 0; x < width; x++ )
				{
					plane[destBase + x] = ReadSample( block, src, c.Type );
					src += c.SampleSize;
				}

				channelOffset += width * c.SampleSize;
			}
		}
	}

	static float ReadSample( byte[] b, int offset, PixelType type ) => type switch
	{
		PixelType.Half => (float)BitConverter.UInt16BitsToHalf( (ushort)(b[offset] | (b[offset + 1] << 8)) ),
		PixelType.Float => BitConverter.Int32BitsToSingle( b[offset] | (b[offset + 1] << 8) | (b[offset + 2] << 16) | (b[offset + 3] << 24) ),
		PixelType.Uint => (uint)(b[offset] | (b[offset + 1] << 8) | (b[offset + 2] << 16) | (b[offset + 3] << 24)),
		_ => 0f,
	};

	static void ParseChannels( byte[] data, int pos, List<ChannelInfo> chans )
	{
		while ( true )
		{
			string name = ReadNullString( data, ref pos );
			if ( name.Length == 0 ) break;

			int ptype = ReadInt32( data, ref pos );
			pos += 1;      // pLinear
			pos += 3;      // reserved
			ReadInt32( data, ref pos ); // xSampling
			ReadInt32( data, ref pos ); // ySampling

			chans.Add( new ChannelInfo( name, (PixelType)ptype ) );
		}
	}

	static byte[] Decompress( Compression compression, byte[] data, int offset, int size, int uncompressedSize )
	{
		// Step 1: codec-specific decompression into a temp buffer.
		byte[] tmp = compression switch
		{
			Compression.Zip or Compression.Zips => Inflate( data, offset, size, uncompressedSize ),
			Compression.Rle => RleDecode( data, offset, size, uncompressedSize ),
			_ => throw new NotSupportedException( $"EXR compression '{compression}' is not supported." ),
		};

		// Step 2: undo EXR's byte predictor + interleave (shared by ZIP and RLE).
		Predictor( tmp );
		return Interleave( tmp );
	}

	static byte[] Inflate( byte[] data, int offset, int size, int expected )
	{
		using var ms = new MemoryStream( data, offset, size );
		using var z = new ZLibStream( ms, CompressionMode.Decompress );
		var outBuf = new byte[expected];
		int read = 0;
		while ( read < expected )
		{
			int n = z.Read( outBuf, read, expected - read );
			if ( n == 0 ) break;
			read += n;
		}
		return outBuf;
	}

	static byte[] RleDecode( byte[] data, int offset, int size, int expected )
	{
		var outBuf = new byte[expected];
		int o = 0;
		int i = offset;
		int end = offset + size;

		while ( i < end && o < expected )
		{
			sbyte count = (sbyte)data[i++];
			if ( count < 0 )
			{
				int n = -count;
				while ( n-- > 0 && i < end && o < expected )
					outBuf[o++] = data[i++];
			}
			else
			{
				int n = count + 1;
				byte v = data[i++];
				while ( n-- > 0 && o < expected )
					outBuf[o++] = v;
			}
		}
		return outBuf;
	}

	// EXR delta predictor: each byte is reconstructed from the running difference.
	static void Predictor( byte[] b )
	{
		for ( int i = 1; i < b.Length; i++ )
		{
			int d = b[i - 1] + b[i] - 128;
			b[i] = (byte)d;
		}
	}

	// EXR de-interleave: data is split into two halves that must be zippered back together.
	static byte[] Interleave( byte[] src )
	{
		int len = src.Length;
		var outB = new byte[len];
		int t1 = 0;
		int t2 = (len + 1) / 2;
		int s = 0;

		while ( true )
		{
			if ( s < len ) outB[s++] = src[t1++]; else break;
			if ( s < len ) outB[s++] = src[t2++]; else break;
		}
		return outB;
	}

	static int ReadInt32( byte[] b, ref int pos )
	{
		int v = b[pos] | (b[pos + 1] << 8) | (b[pos + 2] << 16) | (b[pos + 3] << 24);
		pos += 4;
		return v;
	}

	static ulong ReadUInt64( byte[] b, ref int pos )
	{
		ulong v = 0;
		for ( int i = 0; i < 8; i++ )
			v |= (ulong)b[pos + i] << (i * 8);
		pos += 8;
		return v;
	}

	static string ReadNullString( byte[] b, ref int pos )
	{
		int start = pos;
		while ( pos < b.Length && b[pos] != 0 ) pos++;
		string s = Encoding.ASCII.GetString( b, start, pos - start );
		pos++; // skip null
		return s;
	}
}