using System;
using System.Collections.Generic;
using System.Collections.Concurrent;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;

/// <summary>
/// Thread-safe min-heap (binary heap) priority queue for A* open sets.
/// Lower fScore = higher priority.
/// </summary>
public class MinHeap<T>
{
    private readonly List<(float priority, T item)> _heap = new();

    public int Count => _heap.Count;

    public void Enqueue(T item, float priority)
    {
        _heap.Add((priority, item));
        BubbleUp(_heap.Count - 1);
    }

    public T Dequeue()
    {
        var top = _heap[0].item;
        int last = _heap.Count - 1;
        _heap[0] = _heap[last];
        _heap.RemoveAt(last);
        if (_heap.Count > 0) SiftDown(0);
        return top;
    }

    public T Peek() => _heap[0].item;

    public bool Contains(T item, out float priority)
    {
        for (int i = 0; i < _heap.Count; i++)
        {
            if (EqualityComparer<T>.Default.Equals(_heap[i].item, item))
            {
                priority = _heap[i].priority;
                return true;
            }
        }
        priority = float.MaxValue;
        return false;
    }

    public void UpdatePriority(T item, float newPriority)
    {
        for (int i = 0; i < _heap.Count; i++)
        {
            if (EqualityComparer<T>.Default.Equals(_heap[i].item, item))
            {
                _heap[i] = (newPriority, item);
                BubbleUp(i);
                SiftDown(i);
                return;
            }
        }
    }

    private void BubbleUp(int i)
    {
        while (i > 0)
        {
            int parent = (i - 1) / 2;
            if (_heap[parent].priority <= _heap[i].priority) break;
            (_heap[parent], _heap[i]) = (_heap[i], _heap[parent]);
            i = parent;
        }
    }

    private void SiftDown(int i)
    {
        int n = _heap.Count;
        while (true)
        {
            int smallest = i;
            int left = 2 * i + 1, right = 2 * i + 2;
            if (left < n && _heap[left].priority < _heap[smallest].priority) smallest = left;
            if (right < n && _heap[right].priority < _heap[smallest].priority) smallest = right;
            if (smallest == i) break;
            (_heap[smallest], _heap[i]) = (_heap[i], _heap[smallest]);
            i = smallest;
        }
    }
}

/// <summary>
/// A single pathfinding request queued by an AIAgent.
/// </summary>
public class PathRequest
{
    public int AgentId;
    public Vector2Int Start;
    public Vector2Int Goal;
    public MazeRoom RoomContext;   // null = open hallway A*
    public bool IsHallwayMode;     // true = FindPathToNearestRoom style
    public Action<List<Vector2Int>> Callback; // Called on main thread with result
}

/// <summary>
/// Schedules A* pathfinding requests on worker threads and delivers results back
/// to the main thread on the next frame.
///
/// Usage: instead of calling FindPathInRoom() directly in Update(), agents call
///   PathfindingScheduler.Instance.RequestPath(request)
/// and supply a callback that will fire on the main thread.
///
/// Worker threads process up to <see cref="MaxConcurrentJobs"/> requests in parallel.
/// This keeps Unity's main thread free for rendering and physics while the CPU
/// cores that were sitting idle (the user observed CPU at ~1%) do the heavy lifting.
/// </summary>
public class PathfindingScheduler : MonoBehaviour
{
    public static PathfindingScheduler Instance { get; private set; }

    [Header("Threading")]
    [Tooltip("How many pathfinding jobs may run simultaneously. Match to logical CPU cores - 2.")]
    [SerializeField] private int maxConcurrentJobs = 6;

    [Header("Budget")]
    [Tooltip("Maximum path results to dispatch to agents per Unity frame. Prevents main-thread spikes.")]
    [SerializeField] private int resultsPerFrame = 50;

    // Pending requests waiting for a worker
    private readonly ConcurrentQueue<PathRequest> _pendingQueue = new();

    // Completed results ready to dispatch on the main thread
    private readonly ConcurrentQueue<(Action<List<Vector2Int>> callback, List<Vector2Int> result)> _resultQueue = new();

    // Counts active background jobs so we don't exceed the limit
    private int _activeJobs = 0;

    // Read-only copy of maze data shared across threads (set once after maze generation)
    private MazeData _maze;

    // ---- Snapshot arrays copied from MazeData for thread-safe read access ----
    private bool[,] _walkable;   // [x, y] → is walkable
    private int _mazeWidth, _mazeHeight;
    // Room tile lookup (room ID per tile, -1 = no room)
    private int[,] _roomIdMap;
    // Rooms snapshot (read-only after bake)
    private MazeRoom[] _rooms;

    void Awake()
    {
        if (Instance != null && Instance != this) { Destroy(gameObject); return; }
        Instance = this;
    }

    /// <summary>
    /// Must be called once the maze is fully generated before any agents run.
    /// Bakes thread-safe copies of maze data.
    /// </summary>
    public void BakeMaze(MazeData maze)
    {
        _maze = maze;
        _mazeWidth = maze.Width;
        _mazeHeight = maze.Height;

        _walkable = new bool[_mazeWidth, _mazeHeight];
        _roomIdMap = new int[_mazeWidth, _mazeHeight];

        for (int x = 0; x < _mazeWidth; x++)
            for (int y = 0; y < _mazeHeight; y++)
            {
                _walkable[x, y] = maze.IsWalkable(x, y);
                _roomIdMap[x, y] = -1;
            }

        // Build room-id map
        _rooms = maze.Rooms.ToArray();
        foreach (var room in _rooms)
            for (int x = room.MinX; x <= room.MaxX; x++)
                for (int y = room.MinY; y <= room.MaxY; y++)
                    if (x >= 0 && y >= 0 && x < _mazeWidth && y < _mazeHeight)
                        _roomIdMap[x, y] = room.Id;

        Debug.Log($"[PathfindingScheduler] Baked {_mazeWidth}x{_mazeHeight} maze, {_rooms.Length} rooms, maxJobs={maxConcurrentJobs}");
    }

    /// <summary>
    /// Enqueue a pathfinding request. The callback fires on the main thread.
    /// </summary>
    public void RequestPath(PathRequest request)
    {
        _pendingQueue.Enqueue(request);
    }

    void Update()
    {
        if (_maze == null) return;

        // Dispatch pending requests to worker threads
        while (_activeJobs < maxConcurrentJobs && _pendingQueue.TryDequeue(out var req))
        {
            Interlocked.Increment(ref _activeJobs);
            Task.Run(() => ProcessRequest(req));
        }

        // Deliver completed results back on the main thread
        int dispatched = 0;
        while (dispatched < resultsPerFrame && _resultQueue.TryDequeue(out var result))
        {
            try { result.callback?.Invoke(result.result); }
            catch (Exception e) { Debug.LogException(e); }
            dispatched++;
        }
    }

    // -----------------------------------------------------------------------
    // Worker thread methods — NO Unity API calls allowed here
    // -----------------------------------------------------------------------

    private void ProcessRequest(PathRequest req)
    {
        try
        {
            List<Vector2Int> path;

            if (req.IsHallwayMode)
                path = FindPathHallway(req.Start, req.Goal);
            else if (req.RoomContext != null)
                path = FindPathInRoom(req.Start, req.Goal, req.RoomContext);
            else
                path = FindPathHallway(req.Start, req.Goal);

            _resultQueue.Enqueue((req.Callback, path));
        }
        catch (Exception e)
        {
            // Log safely; Unity's Debug.Log is thread-safe for logging
            Debug.LogException(e);
            _resultQueue.Enqueue((req.Callback, new List<Vector2Int>()));
        }
        finally
        {
            Interlocked.Decrement(ref _activeJobs);
        }
    }

    /// <summary>
    /// Thread-safe A* within a room (limited to room tiles + 1-tile boundary).
    /// Uses a min-heap open set instead of a linear list → O(n log n) vs O(n²).
    /// </summary>
    private List<Vector2Int> FindPathInRoom(Vector2Int start, Vector2Int goal, MazeRoom room)
    {
        var openSet = new MinHeap<Vector2Int>();
        var cameFrom = new Dictionary<Vector2Int, Vector2Int>();
        var gScore = new Dictionary<Vector2Int, float> { [start] = 0f };

        float h0 = Heuristic(start, goal);
        openSet.Enqueue(start, h0);

        const int maxIter = 1000;
        int iter = 0;

        while (openSet.Count > 0 && iter++ < maxIter)
        {
            Vector2Int current = openSet.Dequeue();

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

            foreach (var nb in GetRoomNeighbors(current, room))
            {
                float tentG = gScore[current] + 1f;
                if (!gScore.TryGetValue(nb, out float existingG) || tentG < existingG)
                {
                    cameFrom[nb] = current;
                    gScore[nb] = tentG;
                    float f = tentG + Heuristic(nb, goal);
                    if (openSet.Contains(nb, out _))
                        openSet.UpdatePriority(nb, f);
                    else
                        openSet.Enqueue(nb, f);
                }
            }
        }

        return new List<Vector2Int>();
    }

    /// <summary>
    /// Thread-safe A* through hallways (not room-constrained).
    /// Used when agent is navigating open corridors toward a target tile.
    /// </summary>
    private List<Vector2Int> FindPathHallway(Vector2Int start, Vector2Int goal)
    {
        var openSet = new MinHeap<Vector2Int>();
        var cameFrom = new Dictionary<Vector2Int, Vector2Int>();
        var gScore = new Dictionary<Vector2Int, float> { [start] = 0f };

        openSet.Enqueue(start, Heuristic(start, goal));

        const int maxIter = 5000;
        int iter = 0;

        while (openSet.Count > 0 && iter++ < maxIter)
        {
            Vector2Int current = openSet.Dequeue();

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

            foreach (var nb in GetWalkableNeighbors(current))
            {
                float tentG = gScore[current] + 1f;
                if (!gScore.TryGetValue(nb, out float existingG) || tentG < existingG)
                {
                    cameFrom[nb] = current;
                    gScore[nb] = tentG;
                    float f = tentG + Heuristic(nb, goal);
                    if (openSet.Contains(nb, out _))
                        openSet.UpdatePriority(nb, f);
                    else
                        openSet.Enqueue(nb, f);
                }
            }
        }

        return new List<Vector2Int>();
    }

    // ---- Thread-safe helpers (access only _walkable / _roomIdMap arrays) ----

    private static readonly Vector2Int[] _dirs = {
        new(1,0), new(-1,0), new(0,1), new(0,-1)
    };

    private List<Vector2Int> GetWalkableNeighbors(Vector2Int pos)
    {
        var result = new List<Vector2Int>(4);
        foreach (var d in _dirs)
        {
            int nx = pos.x + d.x, ny = pos.y + d.y;
            if (nx >= 0 && ny >= 0 && nx < _mazeWidth && ny < _mazeHeight && _walkable[nx, ny])
                result.Add(new Vector2Int(nx, ny));
        }
        return result;
    }

    private List<Vector2Int> GetRoomNeighbors(Vector2Int pos, MazeRoom room)
    {
        var result = new List<Vector2Int>(4);
        foreach (var d in _dirs)
        {
            int nx = pos.x + d.x, ny = pos.y + d.y;
            if (nx < 0 || ny < 0 || nx >= _mazeWidth || ny >= _mazeHeight) continue;
            if (!_walkable[nx, ny]) continue;

            bool inRoom = room.Contains(nx, ny);
            bool nearBoundary = nx == room.MinX - 1 || nx == room.MaxX + 1 ||
                                ny == room.MinY - 1 || ny == room.MaxY + 1;
            if (inRoom || nearBoundary)
                result.Add(new Vector2Int(nx, ny));
        }
        return result;
    }

    private static float Heuristic(Vector2Int a, Vector2Int b)
        => Mathf.Abs(a.x - b.x) + Mathf.Abs(a.y - b.y);

    private static List<Vector2Int> ReconstructPath(Dictionary<Vector2Int, Vector2Int> cameFrom, Vector2Int current)
    {
        var path = new List<Vector2Int>();
        while (cameFrom.TryGetValue(current, out var prev))
        {
            path.Add(current);
            current = prev;
        }
        path.Add(current); // start node
        path.Reverse();
        return path;
    }
}