using System.IO;
using System.IO.Compression;
using System.Text.Json;
using System.Text.Json.Serialization;

namespace Sandbox;

/// <summary>
/// Reads SOG v2 (Sorted Optimized Gaussians) compressed splat files.
/// SOG is a ZIP archive containing WebP-encoded textures and a JSON metadata file
/// that store Gaussian splat data in a compact, quantized format.
/// </summary>
public sealed class SogReader
{
	#region Meta JSON Schema

	private sealed class SogMeta
	{
		[JsonPropertyName( "version" )]
		public int Version { get; set; }

		[JsonPropertyName( "count" )]
		public int Count { get; set; }

		[JsonPropertyName( "antialias" )]
		public bool Antialias { get; set; }

		[JsonPropertyName( "means" )]
		public MeansMeta Means { get; set; }

		[JsonPropertyName( "scales" )]
		public CodebookMeta Scales { get; set; }

		[JsonPropertyName( "quats" )]
		public FilesMeta Quats { get; set; }

		[JsonPropertyName( "sh0" )]
		public CodebookMeta Sh0 { get; set; }
	}

	private sealed class MeansMeta
	{
		[JsonPropertyName( "mins" )]
		public float[] Mins { get; set; }

		[JsonPropertyName( "maxs" )]
		public float[] Maxs { get; set; }

		[JsonPropertyName( "files" )]
		public string[] Files { get; set; }
	}

	private sealed class CodebookMeta
	{
		[JsonPropertyName( "codebook" )]
		public float[] Codebook { get; set; }

		[JsonPropertyName( "files" )]
		public string[] Files { get; set; }
	}

	private sealed class FilesMeta
	{
		[JsonPropertyName( "files" )]
		public string[] Files { get; set; }
	}

	#endregion

	/// <summary>
	/// Read SOG compressed splat data and return it as PlyData for compatibility
	/// with the existing rendering pipeline.
	/// </summary>
	public static PlyReader.PlyData Read( Stream stream )
	{
		using var archive = new ZipArchive( stream, ZipArchiveMode.Read );

		// Parse metadata
		var metaEntry = archive.GetEntry( "meta.json" )
			?? throw new Exception( "SOG archive missing meta.json" );

		SogMeta meta;
		using ( var metaStream = metaEntry.Open() )
		{
			meta = JsonSerializer.Deserialize<SogMeta>( metaStream )
				?? throw new Exception( "Failed to parse SOG meta.json" );
		}

		if ( meta.Version != 2 )
			throw new Exception( $"Unsupported SOG version: {meta.Version}" );

		int count = meta.Count;
		if ( count == 0 )
			return new PlyReader.PlyData { VertexCount = 0 };

		// Decode all image textures from the archive into CPU-side pixel arrays.
		// Each image is stored as lossless WebP to preserve quantized values exactly.
		var meansLo = DecodeImage( archive, meta.Means.Files[0] );
		var meansHi = DecodeImage( archive, meta.Means.Files[1] );
		var quatsPx = DecodeImage( archive, meta.Quats.Files[0] );
		var scalePx = DecodeImage( archive, meta.Scales.Files[0] );
		var sh0Px = DecodeImage( archive, meta.Sh0.Files[0] );

		ValidatePixelCount( meansLo, count, "means_l" );
		ValidatePixelCount( meansHi, count, "means_u" );
		ValidatePixelCount( quatsPx, count, "quats" );
		ValidatePixelCount( scalePx, count, "scales" );
		ValidatePixelCount( sh0Px, count, "sh0" );

		var positions = new Vector3[count];
		var rotations = new Vector4[count];
		var scales = new Vector3[count];
		var colors = new Color[count];
		var opacities = new float[count];

		DecodePositions( positions, count, meansLo, meansHi, meta.Means );
		DecodeQuaternions( rotations, count, quatsPx );
		DecodeScales( scales, count, scalePx, meta.Scales.Codebook );
		DecodeColorsAndOpacity( colors, opacities, count, sh0Px, meta.Sh0.Codebook );

		return new PlyReader.PlyData
		{
			Positions = positions,
			Colors = colors,
			Opacities = opacities,
			Scales = scales,
			Rotations = rotations,
			HasColors = true,
			HasOpacity = true,
			HasScaleRotation = true,
			VertexCount = count
		};
	}

	/// <summary>
	/// Extract a ZIP entry and decode its image data (WebP/PNG) to a Color32 pixel array
	/// using s&box's Bitmap API backed by SkiaSharp.
	/// </summary>
	private static Color32[] DecodeImage( ZipArchive archive, string entryName )
	{
		var entry = archive.GetEntry( entryName )
			?? throw new Exception( $"SOG archive missing file: {entryName}" );

		byte[] imageData;
		using ( var entryStream = entry.Open() )
		using ( var ms = new MemoryStream() )
		{
			entryStream.CopyTo( ms );
			imageData = ms.ToArray();
		}

		var bitmap = Bitmap.CreateFromBytes( imageData );
		if ( bitmap is null )
			throw new Exception( $"Failed to decode image: {entryName}" );

		var pixels = bitmap.GetPixels32();
		return pixels;
	}

	private static void ValidatePixelCount( Color32[] pixels, int requiredCount, string name )
	{
		if ( pixels.Length < requiredCount )
			throw new Exception( $"SOG {name} texture has {pixels.Length} pixels, need {requiredCount}" );
	}

	/// <summary>
	/// Decode positions from two 8-bit images that together form 16-bit quantized values per axis.
	/// Values are in a log domain between per-axis min/max bounds; invLogTransform recovers world positions.
	/// </summary>
	private static void DecodePositions( Vector3[] positions, int count,
		Color32[] lo, Color32[] hi, MeansMeta means )
	{
		float xMin = means.Mins[0], xRange = means.Maxs[0] - means.Mins[0];
		float yMin = means.Mins[1], yRange = means.Maxs[1] - means.Mins[1];
		float zMin = means.Mins[2], zRange = means.Maxs[2] - means.Mins[2];

		// Avoid division by zero for degenerate single-point ranges
		if ( xRange == 0 ) xRange = 1f;
		if ( yRange == 0 ) yRange = 1f;
		if ( zRange == 0 ) zRange = 1f;

		const float inv65535 = 1f / 65535f;

		for ( int i = 0; i < count; i++ )
		{
			// Combine upper/lower bytes into 16-bit per-axis quantized values
			ushort qx = (ushort)((hi[i].r << 8) | lo[i].r);
			ushort qy = (ushort)((hi[i].g << 8) | lo[i].g);
			ushort qz = (ushort)((hi[i].b << 8) | lo[i].b);

			// Dequantize to the log-domain range then invert the log transform
			float nx = xMin + xRange * (qx * inv65535);
			float ny = yMin + yRange * (qy * inv65535);
			float nz = zMin + zRange * (qz * inv65535);

			// SOG stores positions in the same coordinate convention as the source training
			// (typically COLMAP: x-right, y-down, z-forward), matching PLY output directly.
			positions[i] = new Vector3(
				InvLogTransform( nx ),
				InvLogTransform( ny ),
				InvLogTransform( nz )
			);
		}
	}

	/// <summary>
	/// Decode quaternions using smallest-three encoding.
	/// RGB stores three quantized components; alpha (252-255) indicates which component was omitted.
	/// </summary>
	private static void DecodeQuaternions( Vector4[] rotations, int count, Color32[] pixels )
	{
		const float sqrt2 = 1.41421356237f;

		for ( int i = 0; i < count; i++ )
		{
			byte tag = pixels[i].a;

			// Alpha encodes the mode: 252=w largest, 253=x, 254=y, 255=z
			if ( tag < 252 || tag > 255 )
			{
				rotations[i] = new Vector4( 1, 0, 0, 0 );
				continue;
			}

			int mode = tag - 252;

			// Dequantize stored components from [0,255] to [-1/√2, 1/√2]
			float a = (pixels[i].r / 255f - 0.5f) * sqrt2;
			float b = (pixels[i].g / 255f - 0.5f) * sqrt2;
			float c = (pixels[i].b / 255f - 0.5f) * sqrt2;

			// Reconstruct the omitted (largest) component from the unit constraint
			float d = MathF.Sqrt( MathF.Max( 0f, 1f - a * a - b * b - c * c ) );

			// Re-assemble quaternion as (w, x, y, z) matching PLY convention
			float qw, qx, qy, qz;
			switch ( mode )
			{
				case 0: qw = d; qx = a; qy = b; qz = c; break;
				case 1: qw = a; qx = d; qy = b; qz = c; break;
				case 2: qw = a; qx = b; qy = d; qz = c; break;
				case 3: qw = a; qx = b; qy = c; qz = d; break;
				default: qw = 1; qx = 0; qy = 0; qz = 0; break;
			}

			// Normalize to correct accumulated quantization error
			float len = MathF.Sqrt( qw * qw + qx * qx + qy * qy + qz * qz );
			if ( len > 0f )
				rotations[i] = new Vector4( qw / len, qx / len, qy / len, qz / len );
			else
				rotations[i] = new Vector4( 1, 0, 0, 0 );
		}
	}

	/// <summary>
	/// Decode scales via codebook lookup. The codebook stores log-space values
	/// (matching 3DGS training output); exp() converts to linear scale magnitudes
	/// to match PlyReader's output convention.
	/// </summary>
	private static void DecodeScales( Vector3[] scales, int count, Color32[] pixels, float[] codebook )
	{
		for ( int i = 0; i < count; i++ )
		{
			float sx = MathF.Exp( codebook[pixels[i].r] );
			float sy = MathF.Exp( codebook[pixels[i].g] );
			float sz = MathF.Exp( codebook[pixels[i].b] );
			scales[i] = new Vector3( sx, sy, sz );
		}
	}

	/// <summary>
	/// Decode colors and opacity from the sh0 image. RGB indices map into a codebook of
	/// SH degree-0 DC coefficients; alpha stores sigmoid-space opacity directly.
	/// </summary>
	private static void DecodeColorsAndOpacity( Color[] colors, float[] opacities, int count,
		Color32[] pixels, float[] codebook )
	{
		const float SH_C0 = 0.28209479177387814f;

		for ( int i = 0; i < count; i++ )
		{
			// SH DC coefficient → linear color via same formula as PlyReader.SHToColor
			float dc0 = codebook[pixels[i].r];
			float dc1 = codebook[pixels[i].g];
			float dc2 = codebook[pixels[i].b];

			float r = Math.Clamp( SH_C0 * dc0 + 0.5f, 0f, 1f );
			float g = Math.Clamp( SH_C0 * dc1 + 0.5f, 0f, 1f );
			float b = Math.Clamp( SH_C0 * dc2 + 0.5f, 0f, 1f );
			colors[i] = new Color( r, g, b, 1f );

			// Alpha byte is already in sigmoid space (0-255 mapped to 0-1),
			// matching PlyReader's output where opacities are sigmoid(logit).
			opacities[i] = pixels[i].a / 255f;
		}
	}

	/// <summary>
	/// Inverse of logTransform(x) = sign(x) * ln(|x| + 1).
	/// Recovers the original world-space coordinate from the log-encoded value.
	/// </summary>
	private static float InvLogTransform( float v )
	{
		float a = MathF.Abs( v );
		float e = MathF.Exp( a ) - 1f;
		return v < 0f ? -e : e;
	}
}