using System.Runtime.InteropServices;

namespace Sandbox;

/// <summary>
/// Data container for a single Gaussian splat cloud.
/// This Does NOT render on its own, it registers with the per-SceneWorld <see cref="SceneGaussianSplatSystem"/> which merges all objects into one globally-sorted draw call.
/// </summary>
public class SceneGaussianSplatObject
{
	/// <summary>
	/// GPU-side splat position, read by both depth and billboard shaders.
	/// </summary>
	[StructLayout( LayoutKind.Sequential )]
	public struct SplatPosition
	{
		public Vector3 Position;
		/// <summary>
		/// Per-splat normalization factor for float16 covariance precision.
		/// Covariance values are divided by this before packing to keep them in a range
		/// where float16 has maximum precision, then multiplied back on the GPU.
		/// </summary>
		public float CovarianceScale;
	}

	/// <summary>
	/// GPU-side splat data, read by billboard shader.
	/// Color is packed as RGBA8 uint, covariance as 3 pairs of float16.
	/// </summary>
	[StructLayout( LayoutKind.Sequential )]
	public struct SplatData
	{
		public uint PackedColor;
		public uint CovAB0; // f16(CovA.x) | f16(CovA.y) << 16
		public uint CovAB1; // f16(CovA.z) | f16(CovB.x) << 16
		public uint CovAB2; // f16(CovB.y) | f16(CovB.z) << 16
	}

	/// <summary>
	/// Size of each splat billboard in world units (inches).
	/// </summary>
	public float SplatSize { get; set; } = 3.0f;

	/// <summary>
	/// Whether this splat cloud has covariance data for ellipsoidal rendering.
	/// </summary>
	public bool HasCovariance { get; private set; }

	/// <summary>
	/// Whether this object is actively rendered. When false, the object's splats remain
	/// in the GPU buffer but the cull shader skips them (no layout re-pack needed).
	/// Used for fast enable/disable cycling without triggering expensive data re-uploads.
	/// </summary>
	public bool IsActive { get; set; } = true;

	/// <summary>
	/// Whether splats receive scene lighting via covariance-derived normals.
	/// </summary>
	public bool ReceiveLighting { get; set; } = true;

	/// <summary>
	/// Whether splats receive shadow darkening without full lighting.
	/// Only takes effect when ReceiveLighting is false.
	/// </summary>
	public bool ReceiveShadows { get; set; }

	/// <summary>
	/// Shadow tint color for shadows-only mode. Darker = stronger shadow effect.
	/// </summary>
	public Color ShadowTint { get; set; } = new Color( 0.1f, 0.1f, 0.15f, 1.0f );

	/// <summary>
	/// Per-object color tint multiplied with each splat's color.
	/// White (default) means no tinting.
	/// </summary>
	public Color Tint { get; set; } = Color.White;

	/// <summary>
	/// Maximum LOD distance in world units. 0 = LOD disabled.
	/// </summary>
	public float LODMaxDistance { get; set; }

	/// <summary>
	/// 8-sample LUT of the LOD curve (keep fraction at evenly spaced normalized distances).
	/// Index 0 = t=0 (at camera), index 7 = t=1 (at max distance).
	/// </summary>
	public float[] LODCurveSamples { get; set; } = new float[8] { 1, 1, 1, 1, 1, 1, 1, 1 };

	/// <summary>
	/// Whether chunked LOD is enabled for this object. When true, the LOD curve is
	/// evaluated per spatial chunk rather than per object.
	/// </summary>
	public bool EnableChunkedLOD { get; set; }

	/// <summary>
	/// Size of each spatial chunk in world units (inches). Used to compute chunk assignments.
	/// </summary>
	public float ChunkSize { get; set; } = 500f;

	/// <summary>
	/// Chunks whose center Y exceeds this local height are exempt from LOD (always full density).
	/// 0 = disabled.
	/// </summary>
	public float ChunkHeightExemption { get; set; }

	/// <summary>
	/// Chunks whose center distance from object origin exceeds this are exempt from LOD.
	/// 0 = disabled.
	/// </summary>
	public float ChunkDistanceExemption { get; set; }

	/// <summary>
	/// Per-splat chunk index (into ChunkCenters). Computed at load time.
	/// null when chunked LOD is disabled or no data loaded.
	/// </summary>
	public ushort[] ChunkIds { get; private set; }

	/// <summary>
	/// Chunk center positions in local space. Index matches ChunkIds values.
	/// </summary>
	public Vector3[] ChunkCenters { get; private set; }

	/// <summary>
	/// Per-chunk flag: true if the chunk is exempt from LOD (height/distance exemption).
	/// </summary>
	public bool[] ChunkExempt { get; private set; }

	/// <summary>
	/// Number of spatial chunks for this object. 0 if chunked LOD is disabled.
	/// </summary>
	public int ChunkCount { get; private set; }

	/// <summary>
	/// Number of splats in this object.
	/// </summary>
	public int SplatCount { get; private set; }

	/// <summary>
	/// CPU-side position data, the system uploads this to the merged GPU buffer.
	/// </summary>
	public SplatPosition[] PositionData { get; private set; }

	/// <summary>
	/// CPU-side packed splat data (color + covariance)
	/// </summary>
	public SplatData[] SplatDataArray { get; private set; }

	/// <summary>
	/// Set to true when splat data has been updated and needs re-uploading.
	/// The system clears this after upload.
	/// </summary>
	public bool DataChanged { get; set; }

	/// <summary>
	/// World-space transform for this splat cloud. Updated each frame by the component.
	/// </summary>
	public Transform Transform { get; set; } = Transform.Zero;

	/// <summary>
	/// Tags for this object, used for rendering flags etc.
	/// </summary>
	public TagSet Tags { get; } = new();

	private readonly SceneGaussianSplatSystem _system;
	private bool _destroyed;

	public SceneGaussianSplatSystem System => _system;

	public SceneGaussianSplatObject( SceneGaussianSplatSystem system )
	{
		_system = system;
		_system.Register( this );
	}

	/// <summary>
	/// Set splat positions and data. The unified render system will pick these up.
	/// If the count changes from a previous non-zero value, marks the system layout dirty
	/// because subsequent objects' buffer offsets have shifted.
	/// </summary>
	public void SetSplatData( SplatPosition[] positions, SplatData[] data, int count, bool hasCovariance = false )
	{
		// Detect offset shift: if this object already had splats and the count changed,
		// all subsequent objects in the unified buffer have wrong offsets.
		if ( SplatCount > 0 && count != SplatCount )
			_system?.MarkLayoutDirty();

		PositionData = positions;
		SplatDataArray = data;
		SplatCount = count;
		HasCovariance = hasCovariance;
		DataChanged = true;
	}

	/// <summary>
	/// Check if this object is still valid (not destroyed)
	/// </summary>
	public bool IsValid() => !_destroyed && _system is not null;

	/// <summary>
	/// Compute spatial chunk assignments for all splats based on the current ChunkSize.
	/// Should be called after SetSplatData when chunked LOD is enabled.
	/// </summary>
	public void ComputeChunks()
	{
		if ( PositionData is null || SplatCount == 0 || ChunkSize <= 0f )
		{
			ChunkIds = null;
			ChunkCenters = null;
			ChunkExempt = null;
			ChunkCount = 0;
			return;
		}

		float invChunkSize = 1f / ChunkSize;

		// Assign each splat to a grid cell, collecting unique cells.
		// We use a dictionary to map grid coordinates → chunk index.
		var cellToChunk = new Dictionary<(int, int, int), int>();
		var chunkIds = new ushort[SplatCount];
		var chunkAccum = new List<(Vector3 Sum, int Count)>();

		for ( int i = 0; i < SplatCount; i++ )
		{
			var pos = PositionData[i].Position;
			int cx = (int)MathF.Floor( pos.x * invChunkSize );
			int cy = (int)MathF.Floor( pos.y * invChunkSize );
			int cz = (int)MathF.Floor( pos.z * invChunkSize );
			var cell = (cx, cy, cz);

			if ( !cellToChunk.TryGetValue( cell, out int chunkIdx ) )
			{
				chunkIdx = chunkAccum.Count;

				// Limit to ushort max — if we exceed 65535 chunks, clamp to last.
				if ( chunkIdx >= ushort.MaxValue )
				{
					chunkIdx = ushort.MaxValue - 1;
				}
				else
				{
					cellToChunk[cell] = chunkIdx;
					chunkAccum.Add( (Vector3.Zero, 0) );
				}
			}

			chunkIds[i] = (ushort)chunkIdx;

			// Accumulate for center computation
			var (sum, count) = chunkAccum[chunkIdx];
			chunkAccum[chunkIdx] = (sum + pos, count + 1);
		}

		// Compute chunk centers (average position of all splats in the chunk)
		int chunkCount = chunkAccum.Count;
		var centers = new Vector3[chunkCount];
		var exempt = new bool[chunkCount];

		for ( int i = 0; i < chunkCount; i++ )
		{
			var (sum, count) = chunkAccum[i];
			centers[i] = sum / count;

			// Evaluate exemptions
			bool isExempt = false;
			if ( ChunkHeightExemption > 0f && centers[i].z > ChunkHeightExemption )
				isExempt = true;
			if ( ChunkDistanceExemption > 0f && centers[i].Length > ChunkDistanceExemption )
				isExempt = true;

			exempt[i] = isExempt;
		}

		ChunkIds = chunkIds;
		ChunkCenters = centers;
		ChunkExempt = exempt;
		ChunkCount = chunkCount;
		DataChanged = true;
	}

	/// <summary>
	/// Clear chunk data (when chunked LOD is disabled).
	/// </summary>
	public void ClearChunks()
	{
		ChunkIds = null;
		ChunkCenters = null;
		ChunkExempt = null;
		ChunkCount = 0;
		DataChanged = true;
	}

	/// <summary>
	/// Clean up and unregister from the system.
	/// </summary>
	public void Destroy()
	{
		if ( _destroyed ) return;
		_destroyed = true;

		_system?.Unregister( this );

		PositionData = null;
		SplatDataArray = null;
		SplatCount = 0;
	}
}