using System;
using System.Collections.Generic;
using Sandbox;

namespace RedSnail.FloraTool;

/// <summary>
/// Painted flora instances, grouped by the model they draw. Unlike grass, flora is placed rather
/// than generated - each instance is a deliberate transform an artist put somewhere - so the
/// transforms are stored rather than derived from a density field. Serialized as a binary blob, so
/// a few tens of thousands of trees cost a couple of megabytes instead of that many GameObjects.
/// </summary>
public sealed class FloraStorage : BlobData
{
	public override int Version => 1;

	public record struct Instance( Vector3 Position, Rotation Rotation, float Scale )
	{
		public readonly Transform ToTransform() => new( Position, Rotation, Scale );
	}

	private readonly Dictionary<string, List<Instance>> _instances = [];

	/// <summary>Bumped on every mutation so the renderer knows to rebuild its GPU buffers.</summary>
	public int Revision { get; private set; }

	public IReadOnlyDictionary<string, List<Instance>> Instances => _instances;

	public int ModelCount => _instances.Count;

	public int TotalCount
	{
		get
		{
			var count = 0;
			foreach ( var list in _instances.Values )
				count += list.Count;
			return count;
		}
	}

	public IReadOnlyList<Instance> GetInstances( string modelPath ) =>
		_instances.TryGetValue( modelPath, out var list ) ? list : [];

	public void AddInstance( string modelPath, Vector3 position, Rotation rotation, float scale )
	{
		if ( string.IsNullOrEmpty( modelPath ) )
			return;

		if ( !_instances.TryGetValue( modelPath, out var list ) )
		{
			list = [];
			_instances[modelPath] = list;
		}

		list.Add( new Instance( position, rotation, scale ) );
		Revision++;
	}

	/// <summary>
	/// True when nothing is planted within <paramref name="spacing"/> of the position. Painting uses
	/// this to keep instances off each other rather than piling several trees into one spot.
	/// </summary>
	public bool IsClear( Vector3 position, float spacing )
	{
		var spacingSquared = spacing * spacing;

		foreach ( var list in _instances.Values )
		{
			for ( var i = 0; i < list.Count; i++ )
			{
				if ( list[i].Position.DistanceSquared( position ) < spacingSquared )
					return false;
			}
		}

		return true;
	}

	/// <summary>
	/// Removes every instance within a radius, ignoring height so the brush erases what is under the
	/// cursor rather than only what sits at the traced elevation.
	/// </summary>
	public int Erase( Vector3 center, float radius )
	{
		var radiusSquared = radius * radius;
		var removed = 0;

		foreach ( var list in _instances.Values )
		{
			removed += list.RemoveAll( instance =>
			{
				var dx = instance.Position.x - center.x;
				var dy = instance.Position.y - center.y;
				return dx * dx + dy * dy <= radiusSquared;
			} );
		}

		if ( removed > 0 )
		{
			PruneEmptyModels();
			Revision++;
		}

		return removed;
	}

	public void ClearAll()
	{
		if ( _instances.Count == 0 )
			return;

		_instances.Clear();
		Revision++;
	}

	/// <summary>
	/// Drops instances whose model no longer resolves, so a deleted asset doesn't keep a dead batch
	/// alive in the renderer.
	/// </summary>
	public int RemoveMissingModels( Func<string, bool> exists )
	{
		List<string> missing = null;

		foreach ( var modelPath in _instances.Keys )
		{
			if ( exists( modelPath ) ) continue;

			missing ??= [];
			missing.Add( modelPath );
		}

		if ( missing is null )
			return 0;

		foreach ( var modelPath in missing )
			_instances.Remove( modelPath );

		Revision++;
		return missing.Count;
	}

	private void PruneEmptyModels()
	{
		List<string> empty = null;

		foreach ( var (modelPath, list) in _instances )
		{
			if ( list.Count > 0 ) continue;

			empty ??= [];
			empty.Add( modelPath );
		}

		if ( empty is null )
			return;

		foreach ( var modelPath in empty )
			_instances.Remove( modelPath );
	}

	public override void Serialize( ref Writer writer )
	{
		writer.Stream.Write( _instances.Count );

		foreach ( var (modelPath, list) in _instances )
		{
			writer.Stream.Write( modelPath );
			writer.Stream.Write( list.Count );

			for ( var i = 0; i < list.Count; i++ )
			{
				writer.Stream.Write( list[i].Position );
				writer.Stream.Write( list[i].Rotation );
				writer.Stream.Write( list[i].Scale );
			}
		}
	}

	public override void Deserialize( ref Reader reader )
	{
		_instances.Clear();

		var modelCount = reader.Stream.Read<int>();

		for ( var m = 0; m < modelCount; m++ )
		{
			var modelPath = reader.Stream.Read<string>();
			var instanceCount = reader.Stream.Read<int>();

			var list = new List<Instance>( instanceCount );

			for ( var i = 0; i < instanceCount; i++ )
			{
				var position = reader.Stream.Read<Vector3>();
				var rotation = reader.Stream.Read<Rotation>();
				var scale = reader.Stream.Read<float>();

				list.Add( new Instance( position, rotation, scale ) );
			}

			if ( !string.IsNullOrEmpty( modelPath ) && list.Count > 0 )
				_instances[modelPath] = list;
		}

		Revision++;
	}
}