using System;

namespace HC3;

public partial class AgentCellCache : GameObjectSystem<AgentCellCache>
{
	private readonly Dictionary<GridCell, HashSet<AgentController>> _agentsInCells = new();

	public AgentCellCache( Scene scene ) : base( scene )
	{
		_agentsInCells = new();
	}

	public static void Register( AgentController agent, GridCell cell )
	{
		if ( !Current._agentsInCells.TryGetValue( cell, out var agents ) )
		{
			agents = new HashSet<AgentController>();
			Current._agentsInCells[cell] = agents;
		}

		agents.Add( agent );
	}

	public static void Unregister( AgentController agent, GridCell cell )
	{
		if ( Current._agentsInCells.TryGetValue( cell, out var agents ) )
		{
			agents.Remove( agent );
			if ( agents.Count == 0 )
				Current._agentsInCells.Remove( cell );
		}
	}

	private static HashSet<AgentController> _neighbourResult;

	public static IReadOnlyCollection<AgentController> GetNeighbouring( GridCell center )
	{
		if ( center is null ) return null;

		_neighbourResult ??= new HashSet<AgentController>();
		_neighbourResult.Clear();

		var neighbors = GridManager.Instance.GetNeighbors( center.Position );
		if ( neighbors is null ) return null;

		foreach ( var cell in neighbors )
		{
			if ( Current._agentsInCells.TryGetValue( cell, out var agents ) )
			{
				foreach ( var a in agents )
					_neighbourResult.Add( a );
			}
		}

		// Include the center cell itself
		if ( Current._agentsInCells.TryGetValue( center, out var centerAgents ) )
		{
			foreach ( var a in centerAgents )
				_neighbourResult.Add( a );
		}

		return _neighbourResult;
	}

	public static IReadOnlyCollection<AgentController> GetAgents( GridCell cell )
	{
		if ( Current._agentsInCells.TryGetValue( cell, out var agents ) )
			return agents;
		return System.Array.Empty<AgentController>();
	}
}