from __future__ import annotations

import numpy
from typing import TYPE_CHECKING
import rtree
from shapely.geometry import Point, LineString

if TYPE_CHECKING:
    from inire.geometry.collision import CollisionEngine
    from inire.geometry.primitives import Port


from inire.geometry.primitives import Port

class VisibilityManager:
    """
    Manages corners of static obstacles for sparse A* / Visibility Graph jumps.
    """
    __slots__ = ('collision_engine', 'corners', 'corner_index', '_corner_graph', '_static_visibility_cache')

    def __init__(self, collision_engine: CollisionEngine) -> None:
        self.collision_engine = collision_engine
        self.corners: list[tuple[float, float]] = []
        self.corner_index = rtree.index.Index()
        self._corner_graph: dict[int, list[tuple[float, float, float]]] = {}
        self._static_visibility_cache: dict[tuple[int, int], list[tuple[float, float, float]]] = {}
        self._build()

    def _build(self) -> None:
        """
        Extract corners and pre-compute corner-to-corner visibility.
        """
        raw_corners = []
        for obj_id, poly in self.collision_engine.static_dilated.items():
            coords = list(poly.exterior.coords)
            if coords[0] == coords[-1]:
                coords = coords[:-1]
            raw_corners.extend(coords)
            for ring in poly.interiors:
                coords = list(ring.coords)
                if coords[0] == coords[-1]:
                    coords = coords[:-1]
                raw_corners.extend(coords)

        if not raw_corners:
            return

        # Deduplicate and snap to 1nm
        seen = set()
        for x, y in raw_corners:
            sx, sy = round(x, 3), round(y, 3)
            if (sx, sy) not in seen:
                seen.add((sx, sy))
                self.corners.append((sx, sy))

        # Build spatial index for corners
        for i, (x, y) in enumerate(self.corners):
            self.corner_index.insert(i, (x, y, x, y))

        # Pre-compute visibility graph between corners
        num_corners = len(self.corners)
        if num_corners > 200:
             # Limit pre-computation if too many corners
             return

        for i in range(num_corners):
            self._corner_graph[i] = []
            p1 = Port(self.corners[i][0], self.corners[i][1], 0)
            for j in range(num_corners):
                if i == j: continue
                cx, cy = self.corners[j]
                dx, dy = cx - p1.x, cy - p1.y
                dist = numpy.sqrt(dx**2 + dy**2)
                angle = numpy.degrees(numpy.arctan2(dy, dx))
                reach = self.collision_engine.ray_cast(p1, angle, max_dist=dist + 0.05)
                if reach >= dist - 0.01:
                    self._corner_graph[i].append((cx, cy, dist))

    def get_visible_corners(self, origin: Port, max_dist: float = 1000.0) -> list[tuple[float, float, float]]:
        """
        Find all corners visible from the origin.
        Returns list of (x, y, distance).
        """
        if max_dist < 0:
            return []
            
        ox, oy = round(origin.x, 3), round(origin.y, 3)
        
        # 1. Exact corner check
        # Use spatial index to find if origin is AT a corner
        nearby = list(self.corner_index.intersection((ox - 0.001, oy - 0.001, ox + 0.001, oy + 0.001)))
        for idx in nearby:
            cx, cy = self.corners[idx]
            if abs(cx - ox) < 1e-4 and abs(cy - oy) < 1e-4:
                # We are at a corner! Return pre-computed graph (filtered by max_dist)
                if idx in self._corner_graph:
                    return [c for c in self._corner_graph[idx] if c[2] <= max_dist]

        # 2. Cache check for arbitrary points
        # Grid-based caching for arbitrary points is tricky, 
        # but since static obstacles don't change, we can cache exact coordinates.
        cache_key = (int(ox * 1000), int(oy * 1000))
        if cache_key in self._static_visibility_cache:
             return self._static_visibility_cache[cache_key]

        # 3. Full visibility check
        bounds = (origin.x - max_dist, origin.y - max_dist, origin.x + max_dist, origin.y + max_dist)
        candidates = list(self.corner_index.intersection(bounds))
        
        visible = []
        for i in candidates:
            cx, cy = self.corners[i]
            dx, dy = cx - origin.x, cy - origin.y
            dist = numpy.sqrt(dx**2 + dy**2)
            
            if dist > max_dist or dist < 1e-3:
                continue
                
            angle = numpy.degrees(numpy.arctan2(dy, dx))
            reach = self.collision_engine.ray_cast(origin, angle, max_dist=dist + 0.05)
            if reach >= dist - 0.01:
                visible.append((cx, cy, dist))
        
        self._static_visibility_cache[cache_key] = visible
        return visible
