"""
Spatial Grid (AOI - Area of Interest) Partitioning System for FreeExile.
Efficiently manages entity locations in 2.5D space and computes visibility subsets (AOI)
to scale to hundreds of thousands of concurrent entities per zone without O(N^2) broadcast overhead.
"""

import math
from typing import Dict, Set, Tuple, List, Optional
from dataclasses import dataclass


@dataclass
class Entity:
    entity_id: int
    x: float
    y: float
    is_active: bool = True


class SpatialGrid:
    def __init__(self, cell_size: float = 64.0):
        self.cell_size = cell_size
        self.cells: Dict[Tuple[int, int], Set[int]] = {}  # (cell_x, cell_y) -> Set of entity_ids
        self.entities: Dict[int, Entity] = {}
        self.entity_cell_map: Dict[int, Tuple[int, int]] = {}

    def get_cell_coords(self, x: float, y: float) -> Tuple[int, int]:
        cell_x = int(math.floor(x / self.cell_size))
        cell_y = int(math.floor(y / self.cell_size))
        return (cell_x, cell_y)

    def add_entity(self, entity: Entity) -> None:
        self.entities[entity.entity_id] = entity
        cell_coords = self.get_cell_coords(entity.x, entity.y)
        self.entity_cell_map[entity.entity_id] = cell_coords

        if cell_coords not in self.cells:
            self.cells[cell_coords] = set()
        self.cells[cell_coords].add(entity.entity_id)

    def remove_entity(self, entity_id: int) -> Optional[Entity]:
        entity = self.entities.pop(entity_id, None)
        if not entity:
            return None

        cell_coords = self.entity_cell_map.pop(entity_id, None)
        if cell_coords and cell_coords in self.cells:
            self.cells[cell_coords].discard(entity_id)
            if not self.cells[cell_coords]:
                del self.cells[cell_coords]

        return entity

    def update_entity_position(self, entity: Entity, new_x: float, new_y: float) -> None:
        old_cell = self.entity_cell_map.get(entity.entity_id)
        new_cell = self.get_cell_coords(new_x, new_y)

        entity.x = new_x
        entity.y = new_y

        if old_cell != new_cell:
            if old_cell and old_cell in self.cells:
                self.cells[old_cell].discard(entity.entity_id)
                if not self.cells[old_cell]:
                    del self.cells[old_cell]

            if new_cell not in self.cells:
                self.cells[new_cell] = set()
            self.cells[new_cell].add(entity.entity_id)
            self.entity_cell_map[entity.entity_id] = new_cell

    def get_entities_in_aoi(self, center_x: float, center_y: float, radius_cells: int = 1) -> Set[int]:
        """
        Retrieves all entity IDs within the center cell and its adjacent neighboring cells.
        radius_cells = 1 queries a 3x3 block of cells (9 cells total).
        """
        center_cell_x, center_cell_y = self.get_cell_coords(center_x, center_y)
        visible_ids: Set[int] = set()

        for dx in range(-radius_cells, radius_cells + 1):
            for dy in range(-radius_cells, radius_cells + 1):
                neighbor_coords = (center_cell_x + dx, center_cell_y + dy)
                if neighbor_coords in self.cells:
                    visible_ids.update(self.cells[neighbor_coords])

        return visible_ids