using System;
using System.Linq;
using Sandbox;

namespace Editor.TerrainConvert;

/// <summary>
/// Which source channel to read the height value from.
/// </summary>
public enum HeightChannel
{
	/// <summary> Red channel only. Standard for grayscale heightmaps. </summary>
	Red,
	Green,
	Blue,
	Alpha,
	/// <summary> Rec.709 weighted luminance of RGB. </summary>
	Luminance,
	/// <summary> Average of the RGB channels. </summary>
	Average,
	/// <summary> The largest of the RGB channels. </summary>
	Max,
}

/// <summary>
/// Bit depth of the output .raw file. s&box terrain stores heights as 16-bit,
/// so 16-bit is recommended for the best precision.
/// </summary>
public enum HeightBitDepth
{
	/// <summary> 16-bit per height sample (recommended). 2 bytes per pixel. </summary>
	Bit16,
	/// <summary> 8-bit per height sample. 1 byte per pixel. </summary>
	Bit8,
}

/// <summary>
/// How the source values are remapped into the 0..1 height range before quantizing.
/// </summary>
public enum HeightNormalize
{
	/// <summary> Use values as-is, clamped to 0..1. </summary>
	Clamp,
	/// <summary> Stretch the min..max range of the image to fill 0..1 (good for HDR/EXR). </summary>
	MinMaxStretch,
}

/// <summary>
/// Byte order of multi-byte (16-bit) samples in the output file.
/// </summary>
public enum HeightByteOrder
{
	/// <summary> Little-endian. What the s&box terrain importer reads by default ("Windows"). </summary>
	LittleEndian,
	/// <summary> Big-endian ("Mac"). </summary>
	BigEndian,
}

/// <summary>
/// Settings for converting an image into a terrain heightmap .raw file.
/// </summary>
public class HeightmapConvertSettings
{
	/// <summary>
	/// Which channel of the source image holds the height. Most heightmaps are
	/// grayscale, where every channel is identical - <see cref="HeightChannel.Red"/> works for those.
	/// </summary>
	[Property]
	public HeightChannel Channel { get; set; } = HeightChannel.Red;

	/// <summary>
	/// 16-bit gives the smoothest terrain and matches what s&box stores internally.
	/// </summary>
	[Property]
	public HeightBitDepth BitDepth { get; set; } = HeightBitDepth.Bit16;

	/// <summary>
	/// Output heightmaps are always square. If 0, the smaller of the source
	/// dimensions is used. The image is resampled to this resolution.
	/// </summary>
	[Property, Range( 0, 8192 )]
	public int Resolution { get; set; } = 0;

	/// <summary>
	/// Round the output resolution down to the nearest power of two (e.g 1024, 2048).
	/// The terrain system expects power-of-two heightmaps, so leave this on.
	/// </summary>
	[Property]
	public bool PowerOfTwo { get; set; } = true;

	/// <summary>
	/// How source values are mapped to the height range. Use <see cref="HeightNormalize.MinMaxStretch"/>
	/// for HDR images that don't already fill 0..1.
	/// </summary>
	[Property]
	public HeightNormalize Normalize { get; set; } = HeightNormalize.Clamp;

	/// <summary>
	/// Flip the image vertically. Image and terrain coordinate origins often differ;
	/// toggle this if your terrain comes out mirrored north/south.
	/// </summary>
	[Property]
	public bool FlipVertical { get; set; } = false;

	/// <summary>
	/// Invert the heights (high becomes low). Useful for depth/inverted maps.
	/// </summary>
	[Property]
	public bool Invert { get; set; } = false;

	/// <summary>
	/// Byte order of 16-bit samples. The s&box importer reads little-endian by default.
	/// </summary>
	[Property, ShowIf( nameof( BitDepth ), HeightBitDepth.Bit16 )]
	public HeightByteOrder ByteOrder { get; set; } = HeightByteOrder.LittleEndian;
}

/// <summary>
/// Converts loaded images into raw single-channel heightmaps compatible with the
/// s&box terrain importer (square, 8 or 16 bit, raw samples with no header).
/// </summary>
public static class HeightmapConverter
{
	/// <summary>
	/// Convert a bitmap into a raw heightmap byte buffer.
	/// </summary>
	/// <param name="bitmap">Source image.</param>
	/// <param name="settings">Conversion settings.</param>
	/// <param name="resolution">The square resolution of the produced heightmap.</param>
	/// <returns>Raw heightmap bytes ready to write to a .raw file.</returns>
	public static byte[] Convert( Bitmap bitmap, HeightmapConvertSettings settings, out int resolution )
	{
		ArgumentNullException.ThrowIfNull( bitmap );
		ArgumentNullException.ThrowIfNull( settings );

		var pixels = bitmap.GetPixels();
		int count = bitmap.Width * bitmap.Height;

		// Extract the chosen channel as a float per pixel, then run the shared pipeline.
		var values = new float[count];
		for ( int i = 0; i < count; i++ )
			values[i] = SampleChannel( pixels[i], settings.Channel );

		return BuildRaw( values, bitmap.Width, bitmap.Height, settings, out resolution );
	}

	/// <summary>
	/// Convert a decoded EXR into a raw heightmap byte buffer, reading height from the
	/// chosen channel (falling back to Y/R/first channel if the exact one is absent).
	/// </summary>
	public static byte[] Convert( ExrImage exr, HeightmapConvertSettings settings, out int resolution )
	{
		ArgumentNullException.ThrowIfNull( exr );
		ArgumentNullException.ThrowIfNull( settings );

		var values = SampleChannel( exr, settings.Channel );
		return BuildRaw( values, exr.Width, exr.Height, settings, out resolution );
	}

	/// <summary>
	/// Shared pipeline: resample a single-channel float plane to a square power-of-two,
	/// normalize into 0..1, optionally invert, and quantize to raw bytes.
	/// </summary>
	static byte[] BuildRaw( float[] values, int srcWidth, int srcHeight, HeightmapConvertSettings settings, out int resolution )
	{
		resolution = ResolveResolution( srcWidth, srcHeight, settings );

		// Bilinear resample to a square of the target resolution, working entirely in float
		// so we don't lose precision on high bit-depth / HDR sources. Resampling produces a
		// fresh array; otherwise clone so the in-place normalize/invert never mutates a
		// caller-owned buffer (e.g. an EXR's cached channel plane).
		values = srcWidth != resolution || srcHeight != resolution
			? ResampleBilinear( values, srcWidth, srcHeight, resolution, resolution )
			: (float[])values.Clone();

		ApplyNormalize( values, settings.Normalize );

		if ( settings.Invert )
		{
			for ( int i = 0; i < values.Length; i++ )
				values[i] = 1f - values[i];
		}

		return Quantize( values, resolution, settings );
	}

	static float[] ResampleBilinear( float[] src, int srcW, int srcH, int dstW, int dstH )
	{
		var dst = new float[dstW * dstH];

		for ( int y = 0; y < dstH; y++ )
		{
			float fy = dstH > 1 ? (float)y / (dstH - 1) * (srcH - 1) : 0f;
			int y0 = (int)fy;
			int y1 = Math.Min( y0 + 1, srcH - 1 );
			float ty = fy - y0;

			for ( int x = 0; x < dstW; x++ )
			{
				float fx = dstW > 1 ? (float)x / (dstW - 1) * (srcW - 1) : 0f;
				int x0 = (int)fx;
				int x1 = Math.Min( x0 + 1, srcW - 1 );
				float tx = fx - x0;

				float top = MathX.Lerp( src[y0 * srcW + x0], src[y0 * srcW + x1], tx );
				float bottom = MathX.Lerp( src[y1 * srcW + x0], src[y1 * srcW + x1], tx );
				dst[y * dstW + x] = MathX.Lerp( top, bottom, ty );
			}
		}

		return dst;
	}

	static float[] SampleChannel( ExrImage exr, HeightChannel channel )
	{
		// For multi-channel EXRs, blend RGB the same way the bitmap path does. For the
		// common single-channel (Y) heightmap, every option resolves to that one plane.
		switch ( channel )
		{
			case HeightChannel.Red: return exr.GetHeightChannel( "R" );
			case HeightChannel.Green: return exr.GetHeightChannel( "G" );
			case HeightChannel.Blue: return exr.GetHeightChannel( "B" );
			case HeightChannel.Alpha: return exr.GetHeightChannel( "A" );
		}

		var r = exr.GetChannel( "R" );
		var g = exr.GetChannel( "G" );
		var b = exr.GetChannel( "B" );

		// No RGB set - it's a luminance/single-channel image, use it directly.
		if ( r is null || g is null || b is null )
			return exr.GetHeightChannel();

		var outv = new float[r.Length];
		for ( int i = 0; i < outv.Length; i++ )
		{
			outv[i] = channel switch
			{
				HeightChannel.Luminance => 0.2126f * r[i] + 0.7152f * g[i] + 0.0722f * b[i],
				HeightChannel.Average => (r[i] + g[i] + b[i]) / 3f,
				HeightChannel.Max => Math.Max( r[i], Math.Max( g[i], b[i] ) ),
				_ => r[i],
			};
		}
		return outv;
	}

	static int ResolveResolution( int width, int height, HeightmapConvertSettings settings )
	{
		int res = settings.Resolution > 0 ? settings.Resolution : Math.Min( width, height );

		if ( settings.PowerOfTwo )
			res = RoundDownToPowerOfTwo( res );

		return Math.Clamp( res, 4, 16384 );
	}

	static float SampleChannel( Color c, HeightChannel channel ) => channel switch
	{
		HeightChannel.Red => c.r,
		HeightChannel.Green => c.g,
		HeightChannel.Blue => c.b,
		HeightChannel.Alpha => c.a,
		HeightChannel.Luminance => 0.2126f * c.r + 0.7152f * c.g + 0.0722f * c.b,
		HeightChannel.Average => (c.r + c.g + c.b) / 3f,
		HeightChannel.Max => Math.Max( c.r, Math.Max( c.g, c.b ) ),
		_ => c.r,
	};

	static void ApplyNormalize( float[] values, HeightNormalize mode )
	{
		if ( mode == HeightNormalize.MinMaxStretch )
		{
			float min = values.Min();
			float max = values.Max();
			float range = max - min;

			if ( range > 1e-6f )
			{
				for ( int i = 0; i < values.Length; i++ )
					values[i] = (values[i] - min) / range;
				return;
			}
			// Flat image - fall through to clamp.
		}

		for ( int i = 0; i < values.Length; i++ )
			values[i] = Math.Clamp( values[i], 0f, 1f );
	}

	static byte[] Quantize( float[] values, int resolution, HeightmapConvertSettings settings )
	{
		bool flip = settings.FlipVertical;

		if ( settings.BitDepth == HeightBitDepth.Bit8 )
		{
			var bytes = new byte[resolution * resolution];
			for ( int y = 0; y < resolution; y++ )
			{
				int srcY = flip ? resolution - 1 - y : y;
				for ( int x = 0; x < resolution; x++ )
				{
					float v = Math.Clamp( values[srcY * resolution + x], 0f, 1f );
					bytes[y * resolution + x] = (byte)MathF.Round( v * byte.MaxValue );
				}
			}
			return bytes;
		}
		else
		{
			bool little = settings.ByteOrder == HeightByteOrder.LittleEndian;
			var bytes = new byte[resolution * resolution * 2];
			for ( int y = 0; y < resolution; y++ )
			{
				int srcY = flip ? resolution - 1 - y : y;
				for ( int x = 0; x < resolution; x++ )
				{
					float v = Math.Clamp( values[srcY * resolution + x], 0f, 1f );
					ushort h = (ushort)MathF.Round( v * ushort.MaxValue );

					int o = (y * resolution + x) * 2;
					if ( little )
					{
						bytes[o] = (byte)(h & 0xFF);
						bytes[o + 1] = (byte)(h >> 8);
					}
					else
					{
						bytes[o] = (byte)(h >> 8);
						bytes[o + 1] = (byte)(h & 0xFF);
					}
				}
			}
			return bytes;
		}
	}

	/// <summary>
	/// Rounds a value down to the nearest power of two.
	/// </summary>
	public static int RoundDownToPowerOfTwo( int value )
	{
		if ( value < 1 ) return 1;
		value |= value >> 1;
		value |= value >> 2;
		value |= value >> 4;
		value |= value >> 8;
		value |= value >> 16;
		return value - (value >> 1);
	}
}