using System.Collections.Generic;
using UnityEngine;

namespace MazeGen
{
    /// <summary>
    /// The full set of generated rooms so far, plus O(1)-ish id lookup, spatial queries and
    /// exit-usage bookkeeping. Everything else (agent tracking, visualization, collision/merge
    /// logic) reads from here instead of walking the scene hierarchy.
    ///
    /// Also tracks connected components (union-find over rooms, merged whenever two exits get
    /// connected) so "does an exit already exist" and "are there still unused exits" can be
    /// asked per-component instead of globally. That distinction matters once there are
    /// multiple start rooms: a global answer would let one isolated branch's safety valves get
    /// silently satisfied by an exit room that lives in a totally different, disconnected part
    /// of the maze - which is exactly how a branch can end up with no way out at all.
    /// </summary>
    public class MazeGraph
    {
        public readonly Dictionary<int, Room> Rooms = new Dictionary<int, Room>();
        public readonly List<Room> StartRooms = new List<Room>();
        public readonly List<Room> ExitRooms = new List<Room>();

        private readonly SpatialRoomIndex _spatialIndex;

        // Union-find over room ids, keyed by root id.
        private readonly Dictionary<int, int> _componentParent = new Dictionary<int, int>();
        private readonly Dictionary<int, bool> _componentHasExit = new Dictionary<int, bool>();
        private readonly Dictionary<int, int> _componentUnusedExitCount = new Dictionary<int, int>();

        private int _nextRoomId;
        private int _nextExitId;

        public int RoomCount => Rooms.Count;

        /// <summary>Total exits across the whole graph that are not yet used (informational only - generation decisions should use ComponentUnusedExitCount instead).</summary>
        public int UnusedExitCount { get; private set; }

        public MazeGraph(float spatialCellSize)
        {
            _spatialIndex = new SpatialRoomIndex(spatialCellSize);
        }

        public int AllocateRoomId() => _nextRoomId++;
        public int AllocateExitId() => _nextExitId++;

        public void Register(Room room)
        {
            Rooms[room.Id] = room;
            if (room.Kind == RoomKind.Start) StartRooms.Add(room);
            if (room.Kind == RoomKind.Exit)
            {
                ExitRooms.Add(room);
                _componentHasExit[Find(room.Id)] = true;
            }
            _spatialIndex.Add(room);
        }

        public Room Get(int id) => Rooms.TryGetValue(id, out var room) ? room : null;

        public IEnumerable<Room> QueryRoomsNear(Vector2 position, float radius) =>
            _spatialIndex.QueryNearby(position, radius);

        /// <summary>Adds an exit to a room and keeps the unused-exit counters (global and per-component) in sync.</summary>
        public void AddExit(Room room, Exit exit)
        {
            room.Exits.Add(exit);
            if (!exit.IsUsed)
            {
                UnusedExitCount++;
                _componentUnusedExitCount[Find(room.Id)]++;
            }
        }

        /// <summary>Marks an exit as used (idempotent) and keeps the unused-exit counters in sync.</summary>
        public void MarkExitUsed(Exit exit)
        {
            if (exit.IsUsed) return;
            exit.IsUsed = true;
            UnusedExitCount--;
            _componentUnusedExitCount[Find(exit.OwnerRoomId)]--;
        }

        /// <summary>
        /// Undoes a speculative AddExit for an exit that turned out to be unusable (e.g. a
        /// door added on-the-fly to close a loop, where the path to it was later found to be
        /// blocked). Only valid for still-unused exits.
        /// </summary>
        public void RemoveUnusedExit(Room room, Exit exit)
        {
            if (exit.IsUsed) return;
            if (room.Exits.Remove(exit))
            {
                UnusedExitCount--;
                _componentUnusedExitCount[Find(room.Id)]--;
            }
        }

        /// <summary>Merges the connected components containing the two rooms - call whenever a connection is made between them.</summary>
        public void Union(int roomIdA, int roomIdB)
        {
            int a = Find(roomIdA);
            int b = Find(roomIdB);
            if (a == b) return;

            _componentParent[b] = a;
            _componentHasExit[a] = _componentHasExit[a] || _componentHasExit[b];
            _componentUnusedExitCount[a] += _componentUnusedExitCount[b];
            _componentHasExit.Remove(b);
            _componentUnusedExitCount.Remove(b);
        }

        public bool AreSameComponent(int roomIdA, int roomIdB) => Find(roomIdA) == Find(roomIdB);

        /// <summary>True if an Exit room already exists anywhere in the same connected component as this room.</summary>
        public bool ComponentHasExit(int roomId) => _componentHasExit[Find(roomId)];

        /// <summary>Unused exits remaining within this room's connected component only.</summary>
        public int ComponentUnusedExitCount(int roomId) => _componentUnusedExitCount[Find(roomId)];

        private int Find(int roomId)
        {
            if (!_componentParent.ContainsKey(roomId))
            {
                _componentParent[roomId] = roomId;
                _componentHasExit[roomId] = false;
                _componentUnusedExitCount[roomId] = 0;
                return roomId;
            }

            int parent = _componentParent[roomId];
            if (parent != roomId)
            {
                parent = Find(parent); // path compression
                _componentParent[roomId] = parent;
            }
            return parent;
        }
    }
}