using Microsoft.Xna.Framework;

namespace MrGameEng.Pathfinding;

/// <summary>
/// A flow field: per-cell distance to the nearest goal and a normalized direction to follow.
/// Built once per goal change by <see cref="FlowFieldBuilder"/>, then any number of agents
/// steer by an O(1) lookup per frame — the tool of choice for crowds heading to shared targets.
/// </summary>
public sealed class FlowField
{
    /// <summary>Grid width in cells.</summary>
    public int Width { get; private set; }

    /// <summary>Grid height in cells.</summary>
    public int Height { get; private set; }

    internal float[] Distances = [];
    internal Vector2[] Directions = [];

    /// <summary>Cost-weighted distance to the nearest goal; <see cref="float.PositiveInfinity"/> when unreachable.</summary>
    public float DistanceAt(int x, int y) => Distances[y * Width + x];

    /// <summary>Normalized direction toward the nearest goal; <see cref="Vector2.Zero"/> at goals and unreachable cells.</summary>
    public Vector2 DirectionAt(int x, int y) => Directions[y * Width + x];

    /// <summary>True when a path to a goal exists from this cell.</summary>
    public bool IsReachable(int x, int y) => !float.IsPositiveInfinity(Distances[y * Width + x]);

    internal void EnsureSize(int width, int height)
    {
        Width = width;
        Height = height;
        var size = width * height;
        if (Distances.Length < size)
        {
            Distances = new float[size];
            Directions = new Vector2[size];
        }
    }
}

/// <summary>
/// Builds <see cref="FlowField"/>s with a multi-source Dijkstra over an <see cref="IPathGrid"/>.
/// Buffers are reused between builds (no allocations after warm-up). One instance per system.
/// </summary>
public sealed class FlowFieldBuilder
{
    private static readonly int[] OffsetX = [1, -1, 0, 0, 1, 1, -1, -1];
    private static readonly int[] OffsetY = [0, 0, 1, -1, 1, -1, 1, -1];
    private const float DiagonalCost = 1.4142135f;

    private readonly IPathGrid _grid;
    private readonly GridConnectivity _connectivity;
    private readonly int _width;
    private readonly int _height;
    private readonly int[] _closedStamp;
    private int[] _heapNodes;
    private float[] _heapPriorities;
    private int _generation;
    private int _heapCount;

    /// <summary>Creates a builder bound to <paramref name="grid"/>.</summary>
    public FlowFieldBuilder(IPathGrid grid, GridConnectivity connectivity = GridConnectivity.Eight)
    {
        _grid = grid;
        _connectivity = connectivity;
        _width = grid.Width;
        _height = grid.Height;
        var size = _width * _height;
        _closedStamp = new int[size];
        _heapNodes = new int[size + 1];
        _heapPriorities = new float[size + 1];
    }

    /// <summary>
    /// Fills <paramref name="field"/> with distances and directions toward the nearest of
    /// <paramref name="goals"/>. Impassable goals are ignored; with no valid goal the whole
    /// field is unreachable.
    /// </summary>
    public void Build(ReadOnlySpan<Point> goals, FlowField field)
    {
        field.EnsureSize(_width, _height);
        var distances = field.Distances;
        var directions = field.Directions;
        Array.Fill(distances, float.PositiveInfinity, 0, _width * _height);

        _generation++;
        _heapCount = 0;

        foreach (var goal in goals)
        {
            if (goal.X >= 0 && goal.X < _width && goal.Y >= 0 && goal.Y < _height &&
                _grid.IsPassable(goal.X, goal.Y))
            {
                var index = goal.Y * _width + goal.X;
                distances[index] = 0f;
                HeapPush(index, 0f);
            }
        }

        var directionCount = (int)_connectivity;
        while (_heapCount > 0)
        {
            var current = HeapPop();
            if (_closedStamp[current] == _generation)
            {
                continue;
            }

            _closedStamp[current] = _generation;
            var cx = current % _width;
            var cy = current / _width;

            for (var d = 0; d < directionCount; d++)
            {
                var nx = cx + OffsetX[d];
                var ny = cy + OffsetY[d];
                if (!Walkable(cx, cy, nx, ny, d))
                {
                    continue;
                }

                var neighbor = ny * _width + nx;
                var tentative = distances[current] + (d < 4 ? 1f : DiagonalCost) * _grid.Cost(nx, ny);
                if (tentative < distances[neighbor])
                {
                    distances[neighbor] = tentative;
                    HeapPush(neighbor, tentative);
                }
            }
        }

        // Направление — к соседу с минимальной дистанцией (с учётом запрета срезать углы).
        for (var y = 0; y < _height; y++)
        {
            for (var x = 0; x < _width; x++)
            {
                var index = y * _width + x;
                directions[index] = Vector2.Zero;
                if (float.IsPositiveInfinity(distances[index]) || distances[index] == 0f)
                {
                    continue;
                }

                var best = distances[index];
                var bestDx = 0;
                var bestDy = 0;
                for (var d = 0; d < directionCount; d++)
                {
                    var nx = x + OffsetX[d];
                    var ny = y + OffsetY[d];
                    if (!Walkable(x, y, nx, ny, d))
                    {
                        continue;
                    }

                    var distance = distances[ny * _width + nx];
                    if (distance < best)
                    {
                        best = distance;
                        bestDx = OffsetX[d];
                        bestDy = OffsetY[d];
                    }
                }

                if (bestDx != 0 || bestDy != 0)
                {
                    directions[index] = Vector2.Normalize(new Vector2(bestDx, bestDy));
                }
            }
        }
    }

    private bool Walkable(int fromX, int fromY, int toX, int toY, int direction)
    {
        if (toX < 0 || toX >= _width || toY < 0 || toY >= _height || !_grid.IsPassable(toX, toY))
        {
            return false;
        }

        return direction < 4 || (_grid.IsPassable(toX, fromY) && _grid.IsPassable(fromX, toY));
    }

    private void HeapPush(int node, float priority)
    {
        if (_heapCount + 1 == _heapNodes.Length)
        {
            Array.Resize(ref _heapNodes, _heapNodes.Length * 2);
            Array.Resize(ref _heapPriorities, _heapPriorities.Length * 2);
        }

        var i = ++_heapCount;
        while (i > 1 && _heapPriorities[i >> 1] > priority)
        {
            _heapNodes[i] = _heapNodes[i >> 1];
            _heapPriorities[i] = _heapPriorities[i >> 1];
            i >>= 1;
        }

        _heapNodes[i] = node;
        _heapPriorities[i] = priority;
    }

    private int HeapPop()
    {
        var top = _heapNodes[1];
        var lastNode = _heapNodes[_heapCount];
        var lastPriority = _heapPriorities[_heapCount];
        _heapCount--;

        var i = 1;
        while (true)
        {
            var child = i << 1;
            if (child > _heapCount)
            {
                break;
            }

            if (child < _heapCount && _heapPriorities[child + 1] < _heapPriorities[child])
            {
                child++;
            }

            if (_heapPriorities[child] >= lastPriority)
            {
                break;
            }

            _heapNodes[i] = _heapNodes[child];
            _heapPriorities[i] = _heapPriorities[child];
            i = child;
        }

        if (_heapCount > 0)
        {
            _heapNodes[i] = lastNode;
            _heapPriorities[i] = lastPriority;
        }

        return top;
    }
}