using System.Collections.Generic;
using UnityEngine;

namespace MazeGen
{
    /// <summary>
    /// Uniform grid bucket index over Room positions, so "which rooms are near this point"
    /// stays cheap even once the maze has thousands of rooms, instead of scanning all of them.
    /// </summary>
    public class SpatialRoomIndex
    {
        private readonly float _cellSize;
        private readonly Dictionary<(int x, int y), List<Room>> _cells = new Dictionary<(int, int), List<Room>>();

        public SpatialRoomIndex(float cellSize)
        {
            _cellSize = Mathf.Max(cellSize, 0.01f);
        }

        private (int, int) CellOf(Vector2 position) =>
            (Mathf.FloorToInt(position.x / _cellSize), Mathf.FloorToInt(position.y / _cellSize));

        public void Add(Room room)
        {
            var cell = CellOf(room.Position);
            if (!_cells.TryGetValue(cell, out var list))
            {
                list = new List<Room>();
                _cells[cell] = list;
            }
            list.Add(room);
        }

        /// <summary>
        /// Yields every registered room whose bucket falls within `radius` of `position`.
        /// This is a broad-phase filter - callers still need an exact distance check.
        /// </summary>
        public IEnumerable<Room> QueryNearby(Vector2 position, float radius)
        {
            var (cx, cy) = CellOf(position);
            int cellRadius = Mathf.CeilToInt(radius / _cellSize) + 1;

            for (int dx = -cellRadius; dx <= cellRadius; dx++)
            {
                for (int dy = -cellRadius; dy <= cellRadius; dy++)
                {
                    if (_cells.TryGetValue((cx + dx, cy + dy), out var list))
                    {
                        foreach (var room in list)
                            yield return room;
                    }
                }
            }
        }
    }
}