AI/GridNavigation.cs
using HC3.Terrain;
using System;

namespace HC3;

public partial class GridNavigation : Component
{
	// allowance for being just under but not quite on the right Z level
	public const float Z_BIAS = 4;

	// Pooled A* data structures to avoid per-pathfind allocations
	private readonly PriorityQueue<Vector3Int, int> _openSet = new();
	private readonly Dictionary<Vector3Int, Vector3Int> _cameFrom = new();
	private readonly Dictionary<Vector3Int, int> _gScore = new();
	private readonly Dictionary<Vector3Int, int> _fScore = new();
	private readonly HashSet<Vector3Int> _closedSet = new();

	// Pooled path result lists to avoid per-pathfind heap allocations
	private readonly Stack<List<Vector3Int>> _gridPathPool = new();
	private readonly Stack<List<Vector3>> _worldPathPool = new();

	private List<Vector3Int> RentGridPath()
	{
		var list = _gridPathPool.Count > 0 ? _gridPathPool.Pop() : new List<Vector3Int>( 64 );
		list.Clear();
		return list;
	}

	/// <summary>
	/// Return a grid path list to the pool when no longer needed.
	/// </summary>
	public void ReturnGridPath( List<Vector3Int> path )
	{
		if ( path is null ) return;
		path.Clear();
		_gridPathPool.Push( path );
	}

	private List<Vector3> RentWorldPath()
	{
		var list = _worldPathPool.Count > 0 ? _worldPathPool.Pop() : new List<Vector3>( 64 );
		list.Clear();
		return list;
	}

	/// <summary>
	/// Singleton
	/// </summary>
	public static GridNavigation Instance { get; private set; }

	/// <summary>
	/// How big is the grid
	/// </summary>
	public int GridSize => GridManager.GridSize;

	/// <summary>
	/// Ref to the legacy grid manager for now
	/// </summary>
	[RequireComponent]
	public GridManager GridManager { get; set; }

	/// <summary>
	/// Ref to the terrain
	/// </summary>
	[RequireComponent]
	public ParkTerrain Terrain { get; private set; }

	/// <summary>
	/// Gets the world height of a cell using the terrain system
	/// </summary>
	/// <param name="cell"></param>
	/// <returns></returns>
	public float GetWorldHeight( GridCell cell )
	{
		var gridPos = GridManager.GridToWorldPosition( cell.Position );
		var tile = Terrain[cell.Position];
		var height = tile.MaxHeight * Terrain.TileHeight;

		return height;
	}

	/// <summary>
	/// Quick list of paths we can definitely get to (right now)
	/// </summary>
	public IEnumerable<Vector3Int> GetNavablePaths( int regionId, NavFlags flags = NavFlags.Default ) => GridManager.Walkable
		.Where( x => x.RegionId == regionId )
		.Select( x => GridManager.WorldToGridPosition3D( x.WorldPosition ) )
		.Where( gridPos => CanEnter( new Vector2Int( gridPos ), flags ) );

	public bool IsWalkable( Vector3Int pos, TileEdge edge, out Vector3Int gridPos )
	{
		gridPos = Vector3Int.Zero;

		var cell = GridManager.Instance?.GetCell( new Vector2Int( pos.x, pos.y ) );
		if ( cell is null ) return false;

		// First check if there's a path to connect from
		var current = cell.GetComponents<IPathConnector>( pos.z ).FirstOrDefault();
		if ( current is not null )
		{
			// Get height offset and target position
			pos += ((Vector3Int)edge.GetDirection()).WithZ( current.GetHeightOffset( edge ) );
		}

		var neighbour = GridManager.Instance?.GetCell( new Vector2Int( pos.x, pos.y ) );
		if ( neighbour is null ) return false;

		// Now check if there's a valid connection to walk to
		if ( neighbour.GetConnection( pos, edge.GetOpposite() ) is not Component connection )
		{
			return false;  // No valid path connection in target cell
		}

		// Check if any decorations block navigation in this direction
		if ( IsBlockedByDecorations( pos, edge ) )
			return false;

		// Check decorations in target cell for navigation blocking in opposite direction
		if ( IsBlockedByDecorations( pos, edge.GetOpposite() ) )
			return false;

		gridPos = GridManager.WorldToGridPosition3D( connection.WorldPosition );
		return true;
	}

	private bool IsBlockedByDecorations( Vector3Int pos, TileEdge edge )
	{
		var cell = GridManager.Instance?.GetCell( new Vector2Int( pos.x, pos.y ) );
		if ( cell is null ) return false;

		var mask = edge.ToPathMask();
		return (cell.GetBlockedNavigation( pos.z ) & mask) != 0;
	}

	/// <summary>
	/// Get a list of enterable neighbouring cells from this cell.
	/// </summary>
	public IEnumerable<Vector3Int> GetEnterableNeighbors( Vector3Int cell, NavFlags flags = NavFlags.Default )
	{
		foreach ( var edge in GridManager.AllEdges )
		{
			var neighborPos = cell + edge.GetDirection();
			var neighbor = GridManager.GetCell( new Vector2Int( neighborPos ) );
			if ( neighbor is null )
				continue;

			if ( !CanEnter( neighbor.Position, flags ) )
				continue;

			// Skip walkable checks
			if ( flags.HasFlag( NavFlags.IncludeUnwalkable ) )
				yield return neighborPos;

			var oppositeEdge = edge.GetOpposite();
			if ( IsWalkable( cell, edge, out var pos ) && IsWalkable( pos, oppositeEdge, out _ ) )
			{
				yield return pos;
			}
		}
	}

	public int GetDistance2d( Vector3Int a, Vector3Int b ) => Math.Abs( a.x - b.x ) + Math.Abs( a.y - b.y );

	/// <summary>
	/// Retrieves a path between <paramref name="start"/> and <paramref name="goal"/>
	/// </summary>
	/// <param name="start"></param>
	/// <param name="goal"></param>
	/// <param name="flags"></param>
	/// <returns></returns>
	public List<Vector3Int> FindPath( Vector3Int start, Vector3Int goal, NavFlags flags = NavFlags.Default )
	{
		_openSet.Clear();
		_cameFrom.Clear();
		_gScore.Clear();
		_fScore.Clear();
		_closedSet.Clear();

		_openSet.Enqueue( start, 0 );

		_gScore[start] = 0;
		_fScore[start] = GetDistance2d( start, goal );

		while ( _openSet.Count > 0 )
		{
			var current = _openSet.Dequeue();

			if ( current == goal )
				return ReconstructPath( _cameFrom, current );

			_closedSet.Add( current );

			// Use precomputed nav edges when possible (standard navigation).
			// Fall back to dynamic GetEnterableNeighbors for IncludeUnwalkable flag.
			IEnumerable<Vector3Int> neighbors;
			if ( flags.HasFlag( NavFlags.IncludeUnwalkable ) )
			{
				neighbors = GetEnterableNeighbors( current, flags );
			}
			else
			{
				var cell = GridManager.GetCell( new Vector2Int( current ) );
				var edges = cell?.GetNavEdges( current.z );
				neighbors = edges ?? Enumerable.Empty<Vector3Int>();
			}

			foreach ( var neighbor in neighbors )
			{
				if ( _closedSet.Contains( neighbor ) )
					continue;

				if ( !CanEnter( new Vector2Int( neighbor ), flags ) )
					continue;

				var tentativeGScore = _gScore[current] + GetDistance2d( current, neighbor );

				if ( tentativeGScore < _gScore.GetValueOrDefault( neighbor, int.MaxValue ) )
				{
					_cameFrom[neighbor] = current;
					_gScore[neighbor] = tentativeGScore;
					_fScore[neighbor] = tentativeGScore + GetDistance2d( neighbor, goal );
					_openSet.Enqueue( neighbor, _fScore[neighbor] );
				}
			}
		}

		return null;
	}

	/// <summary>
	/// Can we enter a cell? This.. seems odd to have in this class? should be done outside maybe
	/// </summary>
	/// <param name="cell"></param>
	/// <param name="flags"></param>
	/// <returns></returns>
	public bool CanEnter( Vector2Int cell, NavFlags flags )
	{
		if ( flags.HasFlag( NavFlags.IncludeUnwalkable ) ) return true;

		if ( flags.HasFlag( NavFlags.Owned ) )
		{
			var buildingZone = BuildingZone.Instance;
			if ( !buildingZone.IsOwned( cell ) ) return false;
		}

		return true;
	}

	/// <summary>
	/// Creates a path in worldspace. If the end target is not walkable, we'll find the nearest walkable neighbouring cell.
	/// </summary>
	/// <param name="worldStart"></param>
	/// <param name="worldTarget"></param>
	/// <param name="flags"></param>
	/// <returns></returns>
	public List<Vector3> FindWorldPath( Vector3 worldStart, Vector3 worldTarget, NavFlags flags )
	{
		var startGrid = GridManager.WorldToGridPosition3D( worldStart );
		var targetGrid = GridManager.WorldToGridPosition3D( worldTarget );

		var cellPath = FindPath( startGrid, targetGrid, flags );

		if ( cellPath is null )
		{
			Log.Warning( $"Cannot path from {startGrid} to {targetGrid}" );
			return null;
		}

		var worldPath = new List<Vector3>( cellPath.Count );

		for ( int i = 1; i < cellPath.Count - 1; i++ )
		{
			var pos = cellPath[i];
			var worldPos = GridManager.GridToWorldPosition( pos ) + GridManager.CentreOffset;

			var cell = GridManager.Instance.GetCell( new Vector2Int( pos ) );
			var path = cell?.GetComponents<Path>( pos.z ).FirstOrDefault();

			// if it's stairs, we want to add nodes to the top and bottom, rather than the centre
			// that way everything else just works?? (delete when proven untrue)
			if ( path.IsValid() && path.StairDirection != 0 )
			{
				Vector2 stairDir = new Vector2( path.StairDirection.x, path.StairDirection.y );

				float dist = (GridManager.GridSize / 2.0f) - Z_BIAS;

				Vector2Int approachDir = 0;
				if ( i > 0 )
					approachDir = new Vector2Int( pos - cellPath[i - 1] );
				else if ( i < cellPath.Count - 1 )
					approachDir = new Vector2Int( cellPath[i + 1] - pos );

				// if we climb the full HeightStep, we'll actually go up a level (bad)
				float maxHeight = GridManager.HeightStep - Z_BIAS;

				bool ascending = Vector2.Dot( approachDir, stairDir ) > 0;
				if ( ascending )
				{
					worldPath.Add( new Vector3( worldPos.x - stairDir.x * dist, worldPos.y - stairDir.y * dist, worldPos.z ) );
					worldPath.Add( new Vector3( worldPos.x + stairDir.x * dist, worldPos.y + stairDir.y * dist, worldPos.z + maxHeight ) );
				}
				else
				{
					worldPath.Add( new Vector3( worldPos.x + stairDir.x * dist, worldPos.y + stairDir.y * dist, worldPos.z + maxHeight ) );
					worldPath.Add( new Vector3( worldPos.x - stairDir.x * dist, worldPos.y - stairDir.y * dist, worldPos.z ) );
				}
			}
			else
			{
				worldPath.Add( worldPos );
			}
		}

		// Return grid path to pool — no longer needed
		ReturnGridPath( cellPath );

		worldPath.Add( worldTarget );

		OffsetPathToRight( worldPath );

		return worldPath;
	}

	private void OffsetPathToRight( List<Vector3> path )
	{
		var offsetLength = Game.Random.Float( GridSize / 10.0f, GridSize / 5.0f );

		if ( path.Count < 2 ) return;

		// Offset path to the right, using the tangent between prev & next point
		for ( var i = 0; i < path.Count; i++ )
		{
			var prev = path[Math.Max( i - 1, 0 )];
			var next = path[Math.Min( i + 1, path.Count - 1 )];

			var tangent = (next - prev).Normal;
			var normal = tangent.Cross( Vector3.Up );
			path[i] += normal * offsetLength;
		}
	}

	private List<Vector3Int> ReconstructPath( Dictionary<Vector3Int, Vector3Int> cameFrom, Vector3Int current )
	{
		var totalPath = RentGridPath();
		totalPath.Add( current );
		while ( cameFrom.TryGetValue( current, out var prev ) )
		{
			current = prev;
			totalPath.Add( current );
		}
		totalPath.Reverse();
		return totalPath;
	}

	protected override void OnAwake()
	{
		Instance = this;
	}
}