from __future__ import annotations

import heapq
import logging
import functools
from typing import TYPE_CHECKING, Literal, Any
import rtree

import numpy
import shapely

from inire.geometry.components import Bend90, SBend, Straight, SEARCH_GRID_SNAP_UM, snap_search_grid
from inire.geometry.primitives import Port
from inire.router.config import RouterConfig
from inire.router.visibility import VisibilityManager

if TYPE_CHECKING:
    from inire.geometry.components import ComponentResult
    from inire.router.cost import CostEvaluator

logger = logging.getLogger(__name__)


class AStarNode:
    """
    A node in the A* search tree.
    """
    __slots__ = ('port', 'g_cost', 'h_cost', 'f_cost', 'parent', 'component_result')

    def __init__(
            self,
            port: Port,
            g_cost: float,
            h_cost: float,
            parent: AStarNode | None = None,
            component_result: ComponentResult | None = None,
            ) -> None:
        self.port = port
        self.g_cost = g_cost
        self.h_cost = h_cost
        self.f_cost = g_cost + h_cost
        self.parent = parent
        self.component_result = component_result

    def __lt__(self, other: AStarNode) -> bool:
        if self.f_cost < other.f_cost - 1e-6:
            return True
        if self.f_cost > other.f_cost + 1e-6:
            return False
        return self.h_cost < other.h_cost


class AStarRouter:
    """
    Waveguide router based on sparse A* search.
    """
    def __init__(self, cost_evaluator: CostEvaluator, node_limit: int | None = None, **kwargs) -> None:
        self.cost_evaluator = cost_evaluator
        self.config = RouterConfig()
        
        if node_limit is not None:
            self.config.node_limit = node_limit
            
        for k, v in kwargs.items():
            if hasattr(self.config, k):
                setattr(self.config, k, v)
        
        self.node_limit = self.config.node_limit
        
        # Visibility Manager for sparse jumps
        self.visibility_manager = VisibilityManager(self.cost_evaluator.collision_engine)
        
        self._hard_collision_set: set[tuple] = set()
        self._congestion_cache: dict[tuple, int] = {}
        self._static_safe_cache: set[tuple] = set()
        self._move_cache: dict[tuple, ComponentResult] = {}

        self.total_nodes_expanded = 0
        self.last_expanded_nodes: list[tuple[float, float, float]] = []
        
        self.metrics = {
            'nodes_expanded': 0,
            'moves_generated': 0,
            'moves_added': 0,
            'pruned_closed_set': 0,
            'pruned_hard_collision': 0,
            'pruned_cost': 0
        }

    def reset_metrics(self) -> None:
        """ Reset all performance counters. """
        for k in self.metrics:
            self.metrics[k] = 0
        self.cost_evaluator.collision_engine.reset_metrics()

    def get_metrics_summary(self) -> str:
        """ Return a human-readable summary of search performance. """
        m = self.metrics
        c = self.cost_evaluator.collision_engine.get_metrics_summary()
        return (f"Search Performance: \n"
                f"  Nodes Expanded: {m['nodes_expanded']}\n"
                f"  Moves: Generated={m['moves_generated']}, Added={m['moves_added']}\n"
                f"  Pruning: ClosedSet={m['pruned_closed_set']}, HardColl={m['pruned_hard_collision']}, Cost={m['pruned_cost']}\n"
                f"  {c}")

    @property
    def _self_dilation(self) -> float:
        return self.cost_evaluator.collision_engine.clearance / 2.0

    def route(
            self,
            start: Port,
            target: Port,
            net_width: float,
            net_id: str = 'default',
            bend_collision_type: Literal['arc', 'bbox', 'clipped_bbox'] | None = None,
            return_partial: bool = False,
            store_expanded: bool = False,
            skip_congestion: bool = False,
            ) -> list[ComponentResult] | None:
        """
        Route a single net using A*.
        """
        self._congestion_cache.clear()
        if store_expanded:
            self.last_expanded_nodes = []

        if bend_collision_type is not None:
            self.config.bend_collision_type = bend_collision_type
            
        open_set: list[AStarNode] = []
        snap = self.config.snap_size
        
        # (x_grid, y_grid, orientation_grid) -> min_g_cost
        closed_set: dict[tuple[int, int, int], float] = {}

        start_node = AStarNode(start, 0.0, self.cost_evaluator.h_manhattan(start, target))
        heapq.heappush(open_set, start_node)

        best_node = start_node
        nodes_expanded = 0
        
        node_limit = self.node_limit

        while open_set:
            if nodes_expanded >= node_limit:
                return self._reconstruct_path(best_node) if return_partial else None

            current = heapq.heappop(open_set)
            
            if current.h_cost < best_node.h_cost:
                best_node = current

            state = (int(round(current.port.x / snap)), int(round(current.port.y / snap)), int(round(current.port.orientation / 1.0)))
            if state in closed_set and closed_set[state] <= current.g_cost + 1e-6:
                continue
            closed_set[state] = current.g_cost

            if store_expanded:
                self.last_expanded_nodes.append((current.port.x, current.port.y, current.port.orientation))

            nodes_expanded += 1
            self.total_nodes_expanded += 1
            self.metrics['nodes_expanded'] += 1

            # Check if we reached the target exactly
            if (abs(current.port.x - target.x) < 1e-6 and
                    abs(current.port.y - target.y) < 1e-6 and
                    abs(current.port.orientation - target.orientation) < 0.1):
                return self._reconstruct_path(current)

            # Expansion
            self._expand_moves(current, target, net_width, net_id, open_set, closed_set, snap, nodes_expanded, skip_congestion=skip_congestion)

        return self._reconstruct_path(best_node) if return_partial else None

    def _expand_moves(
            self,
            current: AStarNode,
            target: Port,
            net_width: float,
            net_id: str,
            open_set: list[AStarNode],
            closed_set: dict[tuple[int, int, int], float],
            snap: float = 1.0,
            nodes_expanded: int = 0,
            skip_congestion: bool = False,
            ) -> None:
        cp = current.port
        base_ori = round(cp.orientation, 2)
        dx_t = target.x - cp.x
        dy_t = target.y - cp.y
        dist_sq = dx_t*dx_t + dy_t*dy_t
        
        rad = numpy.radians(base_ori)
        cos_v, sin_v = numpy.cos(rad), numpy.sin(rad)
        # 1. DIRECT JUMP TO TARGET (Priority 1)
        proj_t = dx_t * cos_v + dy_t * sin_v
        perp_t = -dx_t * sin_v + dy_t * cos_v

        # A. Straight Jump
        if proj_t > 0 and abs(perp_t) < 1e-3 and abs(cp.orientation - target.orientation) < 0.1:
             max_reach = self.cost_evaluator.collision_engine.ray_cast(cp, base_ori, proj_t + 1.0)
             if max_reach >= proj_t - 0.01:
                  self._process_move(current, target, net_width, net_id, open_set, closed_set, snap, f'S{proj_t}', 'S', (proj_t,), skip_congestion, skip_static=True, snap_to_grid=False)

        # 2. VISIBILITY JUMPS & MAX REACH (Priority 2)
        max_reach = self.cost_evaluator.collision_engine.ray_cast(cp, base_ori, self.config.max_straight_length)
        
        straight_lengths = set()
        if max_reach > self.config.min_straight_length:
             # milestone 1: exactly at max_reach (touching)
             straight_lengths.add(snap_search_grid(max_reach, snap))
             # milestone 2: space to turn before collision
             for radius in self.config.bend_radii:
                  if max_reach > radius + self.config.min_straight_length:
                       straight_lengths.add(snap_search_grid(max_reach - radius, snap))
             
             # milestone 3: small buffer for tight maneuvering
             if max_reach > self.config.min_straight_length + 5.0:
                  straight_lengths.add(snap_search_grid(max_reach - 5.0, snap))

        visible_corners = self.visibility_manager.get_visible_corners(cp, max_dist=max_reach)
        for cx, cy, dist in visible_corners:
            proj = (cx - cp.x) * cos_v + (cy - cp.y) * sin_v
            if proj > self.config.min_straight_length:
                 straight_lengths.add(snap_search_grid(proj, snap))
        
        # ALWAYS include the min length for maneuvering
        straight_lengths.add(self.config.min_straight_length)
        
        # If the jump is long, add an intermediate point to allow more flexible turning
        if max_reach > self.config.min_straight_length * 4:
             straight_lengths.add(snap_search_grid(max_reach / 2.0, snap))

        # Target alignment logic (for turning towards target) - Keep this as it's high value
        if abs(base_ori % 180) < 0.1: # Horizontal
             target_dist = abs(target.x - cp.x)
             if target_dist <= max_reach and target_dist > self.config.min_straight_length:
                  straight_lengths.add(snap_search_grid(target_dist, snap))
             
             # Space for turning: target_dist - R and target_dist - 2R
             for radius in self.config.bend_radii:
                  l1 = target_dist - radius
                  if l1 > self.config.min_straight_length:
                       s_l1 = snap_search_grid(l1, snap)
                       if s_l1 <= max_reach and s_l1 > 0.1:
                            straight_lengths.add(s_l1)
                  
                  l2 = target_dist - 2 * radius
                  if l2 > self.config.min_straight_length:
                       s_l2 = snap_search_grid(l2, snap)
                       if s_l2 <= max_reach and s_l2 > 0.1:
                            straight_lengths.add(s_l2)
        else: # Vertical
             target_dist = abs(target.y - cp.y)
             if target_dist <= max_reach and target_dist > self.config.min_straight_length:
                  straight_lengths.add(snap_search_grid(target_dist, snap))
             
             # Space for turning: target_dist - R and target_dist - 2R
             for radius in self.config.bend_radii:
                  l1 = target_dist - radius
                  if l1 > self.config.min_straight_length:
                       s_l1 = snap_search_grid(l1, snap)
                       if s_l1 <= max_reach and s_l1 > 0.1:
                            straight_lengths.add(s_l1)
                  
                  l2 = target_dist - 2 * radius
                  if l2 > self.config.min_straight_length:
                       s_l2 = snap_search_grid(l2, snap)
                       if s_l2 <= max_reach and s_l2 > 0.1:
                            straight_lengths.add(s_l2)

        # NO standard samples here! Only milestones.

        for length in sorted(straight_lengths, reverse=True):
            # Trust ray_cast: these lengths are <= max_reach
            self._process_move(current, target, net_width, net_id, open_set, closed_set, snap, f'S{length}', 'S', (length,), skip_congestion, skip_static=True)

        # 3. BENDS & SBENDS (Priority 3)
        angle_to_target = numpy.degrees(numpy.arctan2(target.y - cp.y, target.x - cp.x))
        allow_backwards = (dist_sq < 150*150)

        for radius in self.config.bend_radii:
            for direction in ['CW', 'CCW']:
                if not allow_backwards:
                    turn = 90 if direction == 'CCW' else -90
                    new_ori = (cp.orientation + turn) % 360
                    new_diff = (angle_to_target - new_ori + 180) % 360 - 180
                    if abs(new_diff) > 135:
                        continue
                self._process_move(current, target, net_width, net_id, open_set, closed_set, snap, f'B{radius}{direction}', 'B', (radius, direction), skip_congestion)

        # 4. SBENDS
        max_sbend_r = max(self.config.sbend_radii) if self.config.sbend_radii else 0
        if max_sbend_r > 0:
            user_offsets = self.config.sbend_offsets
            offsets: set[float] = set(user_offsets) if user_offsets is not None else set()
            
            dx_local = (target.x - cp.x) * cos_v + (target.y - cp.y) * sin_v
            dy_local = -(target.x - cp.x) * sin_v + (target.y - cp.y) * cos_v
            
            # Always try aligning with target if it's forward and within reach
            if dx_local > 0 and abs(dy_local) < 2 * max_sbend_r:
                 # Check if we have enough distance for the SBend
                 # Min distance D = sqrt(4RO - O^2). Smallest R is O/2.
                 min_d = numpy.sqrt(max(0, 4 * (abs(dy_local)/2.0) * abs(dy_local) - dy_local**2))
                 if dx_local >= min_d:
                      offsets.add(dy_local)
            
            # If no offsets provided by user (None), the router "chooses" offsets 
            # by trying grid-aligned steps up to the reach of the largest radius.
            if user_offsets is None:
                 # Try a selection of grid-aligned offsets. 
                 # Fibonacci-ish steps are useful to cover different scales efficiently.
                 for sign in [-1, 1]:
                      for i in [1, 2, 3, 5, 8, 13, 21, 34, 55, 89]:
                           o = sign * i * snap
                           if abs(o) < 2 * max_sbend_r:
                                offsets.add(o)

            for offset in sorted(offsets):
                for radius in self.config.sbend_radii:
                    if abs(offset) >= 2 * radius: continue
                    self._process_move(current, target, net_width, net_id, open_set, closed_set, snap, f'SB{offset}R{radius}', 'SB', (offset, radius), skip_congestion)

    def _process_move(
            self,
            parent: AStarNode,
            target: Port,
            net_width: float,
            net_id: str,
            open_set: list[AStarNode],
            closed_set: dict[tuple[int, int, int], float],
            snap: float,
            move_type: str,
            move_class: Literal['S', 'B', 'SB'],
            params: tuple,
            skip_congestion: bool,
            skip_static: bool = False,
            snap_to_grid: bool = True,
            ) -> None:
        cp = parent.port
        base_ori = round(cp.orientation, 2)
        state_key = (int(round(cp.x / snap)), int(round(cp.y / snap)), int(round(base_ori / 1.0)))
        
        abs_key = (state_key, move_class, params, net_width, self.config.bend_collision_type, snap_to_grid)
        if abs_key in self._move_cache:
            res = self._move_cache[abs_key]
            if move_class == 'B': move_radius = params[0]
            elif move_class == 'SB': move_radius = params[1]
            else: move_radius = None
            self._add_node(parent, res, target, net_width, net_id, open_set, closed_set, move_type, move_radius=move_radius, snap=snap, skip_congestion=skip_congestion)
            return

        rel_key = (base_ori, move_class, params, net_width, self.config.bend_collision_type, self._self_dilation, snap_to_grid)
        
        cache_key = (state_key[0], state_key[1], base_ori, move_type, net_width, snap_to_grid)
        if cache_key in self._hard_collision_set:
            return
            
        if rel_key in self._move_cache:
            res_rel = self._move_cache[rel_key]
            ex = res_rel.end_port.x + cp.x
            ey = res_rel.end_port.y + cp.y
            end_state = (int(round(ex / snap)), int(round(ey / snap)), int(round(res_rel.end_port.orientation / 1.0)))
            if end_state in closed_set and closed_set[end_state] <= parent.g_cost + 1e-6:
                return
            res = res_rel.translate(cp.x, cp.y)
        else:
            try:
                if move_class == 'S':
                    res_rel = Straight.generate(Port(0, 0, base_ori), params[0], net_width, dilation=self._self_dilation, snap_to_grid=snap_to_grid, snap_size=self.config.snap_size)
                elif move_class == 'B':
                    res_rel = Bend90.generate(Port(0, 0, base_ori), params[0], net_width, params[1], collision_type=self.config.bend_collision_type, clip_margin=self.config.bend_clip_margin, dilation=self._self_dilation, snap_to_grid=snap_to_grid, snap_size=self.config.snap_size)
                elif move_class == 'SB':
                    res_rel = SBend.generate(Port(0, 0, base_ori), params[0], params[1], net_width, collision_type=self.config.bend_collision_type, clip_margin=self.config.bend_clip_margin, dilation=self._self_dilation, snap_to_grid=snap_to_grid, snap_size=self.config.snap_size)
                else:
                    return
                self._move_cache[rel_key] = res_rel
                res = res_rel.translate(cp.x, cp.y)
            except (ValueError, ZeroDivisionError):
                return

        self._move_cache[abs_key] = res
        if move_class == 'B': move_radius = params[0]
        elif move_class == 'SB': move_radius = params[1]
        else: move_radius = None
            
        self._add_node(parent, res, target, net_width, net_id, open_set, closed_set, move_type, move_radius=move_radius, snap=snap, skip_congestion=skip_congestion)

    def _add_node(
            self,
            parent: AStarNode,
            result: ComponentResult,
            target: Port,
            net_width: float,
            net_id: str,
            open_set: list[AStarNode],
            closed_set: dict[tuple[int, int, int], float],
            move_type: str,
            move_radius: float | None = None,
            snap: float = 1.0,
            skip_congestion: bool = False,
            ) -> None:
        self.metrics['moves_generated'] += 1
        end_p = result.end_port
        state = (int(round(end_p.x / snap)), int(round(end_p.y / snap)), int(round(end_p.orientation / 1.0)))
        
        if state in closed_set and closed_set[state] <= parent.g_cost + 1e-6:
            self.metrics['pruned_closed_set'] += 1
            return

        parent_p = parent.port
        cache_key = (int(round(parent_p.x / snap)), int(round(parent_p.y / snap)), int(round(parent_p.orientation / 1.0)), move_type, net_width)
        
        if cache_key in self._hard_collision_set:
            self.metrics['pruned_hard_collision'] += 1
            return

        is_static_safe = (cache_key in self._static_safe_cache)
        if not is_static_safe:
            collision_engine = self.cost_evaluator.collision_engine
            # Fast check for straights
            if 'S' in move_type and 'SB' not in move_type:
                 if collision_engine.check_move_straight_static(parent_p, result.length):
                      self._hard_collision_set.add(cache_key)
                      self.metrics['pruned_hard_collision'] += 1
                      return
                 is_static_safe = True
            
            if not is_static_safe:
                if collision_engine.check_move_static(result, start_port=parent_p, end_port=end_p):
                    self._hard_collision_set.add(cache_key)
                    self.metrics['pruned_hard_collision'] += 1
                    return
                else:
                    self._static_safe_cache.add(cache_key)

        total_overlaps = 0
        if not skip_congestion:
            if cache_key in self._congestion_cache:
                total_overlaps = self._congestion_cache[cache_key]
            else:
                total_overlaps = self.cost_evaluator.collision_engine.check_move_congestion(result, net_id)
                self._congestion_cache[cache_key] = total_overlaps

        penalty = 0.0
        if 'SB' in move_type: penalty = self.config.sbend_penalty
        elif 'B' in move_type: penalty = self.config.bend_penalty

        # Scale penalty by radius (larger radius = smoother = lower penalty)
        if move_radius is not None and move_radius > 1e-6:
            penalty *= (10.0 / move_radius)**0.5

        move_cost = self.cost_evaluator.evaluate_move(
            result.geometry, result.end_port, net_width, net_id,
            start_port=parent_p, length=result.length,
            dilated_geometry=result.dilated_geometry, penalty=penalty,
            skip_static=True, skip_congestion=True
        )
        move_cost += total_overlaps * self.cost_evaluator.congestion_penalty

        if move_cost > 1e12:
            self.metrics['pruned_cost'] += 1
            return

        g_cost = parent.g_cost + move_cost
        if state in closed_set and closed_set[state] <= g_cost + 1e-6:
            self.metrics['pruned_closed_set'] += 1
            return
            
        h_cost = self.cost_evaluator.h_manhattan(result.end_port, target)
        new_node = AStarNode(result.end_port, g_cost, h_cost, parent, result)
        heapq.heappush(open_set, new_node)
        self.metrics['moves_added'] += 1

    def _reconstruct_path(self, end_node: AStarNode) -> list[ComponentResult]:
        path = []
        curr: AStarNode | None = end_node
        while curr and curr.component_result:
            path.append(curr.component_result)
            curr = curr.parent
        return path[::-1]
