"""
Spatial index for fast token lookup.
"""

from .models import TokenLike


class TokenIndex:
    """
    Spatial index for tokens to enable fast nearby token lookup.

    Uses grid-based spatial hashing for O(1) average lookup instead of O(n).
    """

    def __init__(self, tokens: list[TokenLike], grid_size: float = 100.0):
        """
        Build spatial index from tokens.

        Args:
            tokens: List of tokens to index
            grid_size: Size of grid cells in pixels
        """
        self.tokens = tokens
        self.grid_size = grid_size
        self._grid: dict[tuple[int, int], list[TokenLike]] = {}
        self._token_centers: dict[int, tuple[float, float]] = {}
        self._token_text_lower: dict[int, str] = {}

        # Build index
        for i, token in enumerate(tokens):
            # Cache center coordinates
            center_x = (token.bbox[0] + token.bbox[2]) / 2
            center_y = (token.bbox[1] + token.bbox[3]) / 2
            self._token_centers[id(token)] = (center_x, center_y)

            # Cache lowercased text
            self._token_text_lower[id(token)] = token.text.lower()

            # Add to grid cell
            grid_x = int(center_x / grid_size)
            grid_y = int(center_y / grid_size)
            key = (grid_x, grid_y)
            if key not in self._grid:
                self._grid[key] = []
            self._grid[key].append(token)

    def get_center(self, token: TokenLike) -> tuple[float, float]:
        """Get cached center coordinates for token."""
        return self._token_centers.get(id(token), (
            (token.bbox[0] + token.bbox[2]) / 2,
            (token.bbox[1] + token.bbox[3]) / 2
        ))

    def get_text_lower(self, token: TokenLike) -> str:
        """Get cached lowercased text for token."""
        return self._token_text_lower.get(id(token), token.text.lower())

    def find_nearby(self, token: TokenLike, radius: float) -> list[TokenLike]:
        """
        Find all tokens within radius of the given token.

        Uses grid-based lookup for O(1) average case instead of O(n).
        """
        center = self.get_center(token)
        center_x, center_y = center

        # Determine which grid cells to search
        cells_to_check = int(radius / self.grid_size) + 1
        grid_x = int(center_x / self.grid_size)
        grid_y = int(center_y / self.grid_size)

        nearby = []
        radius_sq = radius * radius

        # Check all nearby grid cells
        for dx in range(-cells_to_check, cells_to_check + 1):
            for dy in range(-cells_to_check, cells_to_check + 1):
                key = (grid_x + dx, grid_y + dy)
                if key not in self._grid:
                    continue

                for other in self._grid[key]:
                    if other is token:
                        continue

                    other_center = self.get_center(other)
                    dist_sq = (center_x - other_center[0]) ** 2 + (center_y - other_center[1]) ** 2

                    if dist_sq <= radius_sq:
                        nearby.append(other)

        return nearby
