using System.Globalization;
using System.IO;
using System.Text;

namespace Sandbox;

/// <summary>
/// Reads PLY point cloud data.
/// </summary>
public sealed class PlyReader
{
	public struct PlyData
	{
		public Vector3[] Positions;
		public Color[] Colors;
		public float[] Opacities;
		public Vector3[] Scales;
		public Vector4[] Rotations;
		public bool HasColors;
		public bool HasOpacity;
		public bool HasScaleRotation;
		public int VertexCount;
	}

	private enum PlyFormat
	{
		Ascii,
		BinaryLittleEndian,
		BinaryBigEndian
	}

	private struct PropertyInfo
	{
		public string Name;
		public string Type;
		public int ByteSize;
	}

	/// <summary>
	/// Read vertex positions from a PLY file stream.
	/// </summary>
	public static PlyData Read( Stream stream )
	{
		using var reader = new BinaryReader( stream, Encoding.ASCII, leaveOpen: true );

		// Validate magic number
		var magic = ReadAsciiLine( reader );
		if ( magic != "ply" )
			throw new Exception( "Not a valid PLY file — missing 'ply' magic number" );

		var format = PlyFormat.Ascii;
		int vertexCount = 0;
		var properties = new List<PropertyInfo>();
		bool inVertexElement = false;

		// Parse header
		while ( true )
		{
			var line = ReadAsciiLine( reader );
			if ( line is null )
				throw new Exception( "Unexpected end of file while reading PLY header" );

			if ( line == "end_header" )
				break;

			var parts = line.Split( ' ', StringSplitOptions.RemoveEmptyEntries );
			if ( parts.Length == 0 )
				continue;

			switch ( parts[0] )
			{
				case "format" when parts.Length >= 2:
					format = parts[1] switch
					{
						"ascii" => PlyFormat.Ascii,
						"binary_little_endian" => PlyFormat.BinaryLittleEndian,
						"binary_big_endian" => PlyFormat.BinaryBigEndian,
						_ => throw new Exception( $"Unsupported PLY format: {parts[1]}" )
					};
					break;

				case "element" when parts.Length >= 3:
					// When we hit a new element, stop tracking properties for vertex
					inVertexElement = parts[1] == "vertex";
					if ( inVertexElement )
						vertexCount = int.Parse( parts[2], CultureInfo.InvariantCulture );
					break;

				case "property" when inVertexElement && parts.Length >= 3:
					// Skip list properties (eg. face vertex indices)
					if ( parts[1] == "list" )
						break;

					properties.Add( new PropertyInfo
					{
						Name = parts[2],
						Type = parts[1],
						ByteSize = GetTypeByteSize( parts[1] )
					} );
					break;
			}
		}

		if ( vertexCount == 0 )
			return new PlyData { Positions = [], VertexCount = 0 };

		// Find x, y, z property indices
		int xIdx = properties.FindIndex( p => p.Name == "x" );
		int yIdx = properties.FindIndex( p => p.Name == "y" );
		int zIdx = properties.FindIndex( p => p.Name == "z" );

		if ( xIdx < 0 || yIdx < 0 || zIdx < 0 )
			throw new Exception( "PLY file missing required x, y, z vertex properties" );

		// Look for color data, either SH DC coefficients or direct RGB
		int shR = properties.FindIndex( p => p.Name == "f_dc_0" );
		int shG = properties.FindIndex( p => p.Name == "f_dc_1" );
		int shB = properties.FindIndex( p => p.Name == "f_dc_2" );
		bool hasSH = shR >= 0 && shG >= 0 && shB >= 0;

		int rgbR = properties.FindIndex( p => p.Name == "red" );
		int rgbG = properties.FindIndex( p => p.Name == "green" );
		int rgbB = properties.FindIndex( p => p.Name == "blue" );
		bool hasRGB = rgbR >= 0 && rgbG >= 0 && rgbB >= 0;

		bool hasColors = hasSH || hasRGB;

		// Look for opacity (stored as logit in 3DGS PLY files)
		int opacityIdx = properties.FindIndex( p => p.Name == "opacity" );
		bool hasOpacity = opacityIdx >= 0;

		// Look for Gaussian scale and rotation data
		int s0 = properties.FindIndex( p => p.Name == "scale_0" );
		int s1 = properties.FindIndex( p => p.Name == "scale_1" );
		int s2 = properties.FindIndex( p => p.Name == "scale_2" );
		int r0 = properties.FindIndex( p => p.Name == "rot_0" );
		int r1 = properties.FindIndex( p => p.Name == "rot_1" );
		int r2 = properties.FindIndex( p => p.Name == "rot_2" );
		int r3 = properties.FindIndex( p => p.Name == "rot_3" );
		bool hasScaleRot = s0 >= 0 && s1 >= 0 && s2 >= 0
			&& r0 >= 0 && r1 >= 0 && r2 >= 0 && r3 >= 0;

		var positions = new Vector3[vertexCount];
		var colors = hasColors ? new Color[vertexCount] : null;
		var opacities = hasOpacity ? new float[vertexCount] : null;
		var scales = hasScaleRot ? new Vector3[vertexCount] : null;
		var rotations = hasScaleRot ? new Vector4[vertexCount] : null;

		var readParams = new VertexReadParams
		{
			xIdx = xIdx,
			yIdx = yIdx,
			zIdx = zIdx,
			HasSH = hasSH,
			shR = shR,
			shG = shG,
			shB = shB,
			HasRGB = hasRGB,
			rgbR = rgbR,
			rgbG = rgbG,
			rgbB = rgbB,
			HasOpacity = hasOpacity,
			opacityIdx = opacityIdx,
			HasScaleRot = hasScaleRot,
			s0 = s0,
			s1 = s1,
			s2 = s2,
			r0 = r0,
			r1 = r1,
			r2 = r2,
			r3 = r3
		};

		if ( format == PlyFormat.Ascii )
			ReadAsciiVertices( reader, properties, positions, colors, opacities, scales, rotations, vertexCount, readParams );
		else if ( format == PlyFormat.BinaryLittleEndian )
			ReadBinaryVertices( reader, properties, positions, colors, opacities, scales, rotations, vertexCount, readParams );
		else
			throw new Exception( "Binary big-endian PLY is not supported" );

		return new PlyData
		{
			Positions = positions,
			Colors = colors,
			Opacities = opacities,
			Scales = scales,
			Rotations = rotations,
			HasColors = hasColors,
			HasOpacity = hasOpacity,
			HasScaleRotation = hasScaleRot,
			VertexCount = vertexCount
		};
	}

	/// <summary>
	/// Indices into the property list for position and color fields.
	/// </summary>
	private struct VertexReadParams
	{
		public int xIdx, yIdx, zIdx;
		public bool HasSH;
		public int shR, shG, shB;
		public bool HasRGB;
		public int rgbR, rgbG, rgbB;
		public bool HasOpacity;
		public int opacityIdx;
		public bool HasScaleRot;
		public int s0, s1, s2;
		public int r0, r1, r2, r3;
	}

	/// <summary>
	/// Convert SH degree-0 coefficient to linear color.
	/// SH_C0 = 0.28209479177387814 = 1 / (2 * sqrt(pi))
	/// </summary>
	private static Color SHToColor( float dc0, float dc1, float dc2 )
	{
		const float SH_C0 = 0.28209479177387814f;
		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 );
		return new Color( r, g, b, 1f );
	}

	/// <summary>
	/// Convert 0-255 integer RGB to Color.
	/// </summary>
	private static Color RGBToColor( float r, float g, float b )
	{
		return new Color( r / 255f, g / 255f, b / 255f, 1f );
	}

	/// <summary>
	/// Sigmoid function: converts logit-space opacity to 0-1 range.
	/// </summary>
	private static float Sigmoid( float x ) => 1f / (1f + MathF.Exp( -x ));

	private static void ReadAsciiVertices( BinaryReader reader, List<PropertyInfo> properties,
		Vector3[] positions, Color[] colors, float[] opacities, Vector3[] scales, Vector4[] rotations,
		int vertexCount, VertexReadParams p )
	{
		for ( int i = 0; i < vertexCount; i++ )
		{
			var line = ReadAsciiLine( reader );
			if ( line is null )
				throw new Exception( $"Unexpected end of file at vertex {i}" );

			var values = line.Split( ' ', StringSplitOptions.RemoveEmptyEntries );

			float x = float.Parse( values[p.xIdx], CultureInfo.InvariantCulture );
			float y = float.Parse( values[p.yIdx], CultureInfo.InvariantCulture );
			float z = float.Parse( values[p.zIdx], CultureInfo.InvariantCulture );
			positions[i] = new Vector3( x, y, z );

			if ( colors is not null )
			{
				if ( p.HasSH )
				{
					float dc0 = float.Parse( values[p.shR], CultureInfo.InvariantCulture );
					float dc1 = float.Parse( values[p.shG], CultureInfo.InvariantCulture );
					float dc2 = float.Parse( values[p.shB], CultureInfo.InvariantCulture );
					colors[i] = SHToColor( dc0, dc1, dc2 );
				}
				else if ( p.HasRGB )
				{
					float r = float.Parse( values[p.rgbR], CultureInfo.InvariantCulture );
					float g = float.Parse( values[p.rgbG], CultureInfo.InvariantCulture );
					float b = float.Parse( values[p.rgbB], CultureInfo.InvariantCulture );
					colors[i] = RGBToColor( r, g, b );
				}
			}

			if ( opacities is not null )
			{
				float raw = float.Parse( values[p.opacityIdx], CultureInfo.InvariantCulture );
				opacities[i] = Sigmoid( raw );
			}

			if ( scales is not null )
			{
				float sv0 = float.Parse( values[p.s0], CultureInfo.InvariantCulture );
				float sv1 = float.Parse( values[p.s1], CultureInfo.InvariantCulture );
				float sv2 = float.Parse( values[p.s2], CultureInfo.InvariantCulture );
				scales[i] = new Vector3( MathF.Exp( sv0 ), MathF.Exp( sv1 ), MathF.Exp( sv2 ) );

				float rv0 = float.Parse( values[p.r0], CultureInfo.InvariantCulture );
				float rv1 = float.Parse( values[p.r1], CultureInfo.InvariantCulture );
				float rv2 = float.Parse( values[p.r2], CultureInfo.InvariantCulture );
				float rv3 = float.Parse( values[p.r3], CultureInfo.InvariantCulture );
				float rlen = MathF.Sqrt( rv0 * rv0 + rv1 * rv1 + rv2 * rv2 + rv3 * rv3 );
				if ( rlen > 0f )
					rotations[i] = new Vector4( rv0 / rlen, rv1 / rlen, rv2 / rlen, rv3 / rlen );
				else
					rotations[i] = new Vector4( 1, 0, 0, 0 );
			}
		}
	}

	private static void ReadBinaryVertices( BinaryReader reader, List<PropertyInfo> properties,
		Vector3[] positions, Color[] colors, float[] opacities, Vector3[] scales, Vector4[] rotations,
		int vertexCount, VertexReadParams p )
	{
		int vertexStride = 0;
		var offsets = new int[properties.Count];
		for ( int j = 0; j < properties.Count; j++ )
		{
			offsets[j] = vertexStride;
			vertexStride += properties[j].ByteSize;
		}

		var vertexBytes = new byte[vertexStride];

		for ( int i = 0; i < vertexCount; i++ )
		{
			int read = reader.Read( vertexBytes, 0, vertexStride );
			if ( read < vertexStride )
				throw new Exception( $"Unexpected end of file at vertex {i}" );

			float x = ReadPropertyAsFloat( vertexBytes, offsets[p.xIdx], properties[p.xIdx] );
			float y = ReadPropertyAsFloat( vertexBytes, offsets[p.yIdx], properties[p.yIdx] );
			float z = ReadPropertyAsFloat( vertexBytes, offsets[p.zIdx], properties[p.zIdx] );
			positions[i] = new Vector3( x, y, z );

			if ( colors is not null )
			{
				if ( p.HasSH )
				{
					float dc0 = ReadPropertyAsFloat( vertexBytes, offsets[p.shR], properties[p.shR] );
					float dc1 = ReadPropertyAsFloat( vertexBytes, offsets[p.shG], properties[p.shG] );
					float dc2 = ReadPropertyAsFloat( vertexBytes, offsets[p.shB], properties[p.shB] );
					colors[i] = SHToColor( dc0, dc1, dc2 );
				}
				else if ( p.HasRGB )
				{
					float r = ReadPropertyAsFloat( vertexBytes, offsets[p.rgbR], properties[p.rgbR] );
					float g = ReadPropertyAsFloat( vertexBytes, offsets[p.rgbG], properties[p.rgbG] );
					float b = ReadPropertyAsFloat( vertexBytes, offsets[p.rgbB], properties[p.rgbB] );
					colors[i] = RGBToColor( r, g, b );
				}
			}

			if ( opacities is not null )
			{
				float raw = ReadPropertyAsFloat( vertexBytes, offsets[p.opacityIdx], properties[p.opacityIdx] );
				opacities[i] = Sigmoid( raw );
			}

			if ( scales is not null )
			{
				float sv0 = ReadPropertyAsFloat( vertexBytes, offsets[p.s0], properties[p.s0] );
				float sv1 = ReadPropertyAsFloat( vertexBytes, offsets[p.s1], properties[p.s1] );
				float sv2 = ReadPropertyAsFloat( vertexBytes, offsets[p.s2], properties[p.s2] );
				scales[i] = new Vector3( MathF.Exp( sv0 ), MathF.Exp( sv1 ), MathF.Exp( sv2 ) );

				float rv0 = ReadPropertyAsFloat( vertexBytes, offsets[p.r0], properties[p.r0] );
				float rv1 = ReadPropertyAsFloat( vertexBytes, offsets[p.r1], properties[p.r1] );
				float rv2 = ReadPropertyAsFloat( vertexBytes, offsets[p.r2], properties[p.r2] );
				float rv3 = ReadPropertyAsFloat( vertexBytes, offsets[p.r3], properties[p.r3] );
				float rlen = MathF.Sqrt( rv0 * rv0 + rv1 * rv1 + rv2 * rv2 + rv3 * rv3 );
				if ( rlen > 0f )
					rotations[i] = new Vector4( rv0 / rlen, rv1 / rlen, rv2 / rlen, rv3 / rlen );
				else
					rotations[i] = new Vector4( 1, 0, 0, 0 );
			}
		}
	}

	private static float ReadPropertyAsFloat( byte[] buffer, int offset, PropertyInfo prop )
	{
		return prop.Type switch
		{
			"float" or "float32" => BitConverter.ToSingle( buffer, offset ),
			"double" or "float64" => (float)BitConverter.ToDouble( buffer, offset ),
			"int" or "int32" => BitConverter.ToInt32( buffer, offset ),
			"uint" or "uint32" => BitConverter.ToUInt32( buffer, offset ),
			"short" or "int16" => BitConverter.ToInt16( buffer, offset ),
			"ushort" or "uint16" => BitConverter.ToUInt16( buffer, offset ),
			"char" or "int8" => (sbyte)buffer[offset],
			"uchar" or "uint8" => buffer[offset],
			_ => throw new Exception( $"Unsupported PLY property type: {prop.Type}" )
		};
	}

	private static int GetTypeByteSize( string type )
	{
		return type switch
		{
			"char" or "int8" or "uchar" or "uint8" => 1,
			"short" or "int16" or "ushort" or "uint16" => 2,
			"int" or "int32" or "uint" or "uint32" or "float" or "float32" => 4,
			"double" or "float64" => 8,
			_ => throw new Exception( $"Unknown PLY type: {type}" )
		};
	}

	/// <summary>
	/// Read a line of ASCII text from a binary reader, handling both \n and \r\n.
	/// </summary>
	private static string ReadAsciiLine( BinaryReader reader )
	{
		var sb = new StringBuilder();
		try
		{
			while ( true )
			{
				byte b = reader.ReadByte();
				if ( b == '\n' )
					break;
				if ( b == '\r' )
					continue;
				sb.Append( (char)b );
			}
		}
		catch ( EndOfStreamException )
		{
			if ( sb.Length == 0 )
				return null;
		}
		return sb.ToString();
	}
}