from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable

if TYPE_CHECKING:
    from inire.geometry.components import ComponentResult
    from inire.geometry.primitives import Port
    from inire.router.astar import AStarRouter
    from inire.router.cost import CostEvaluator

logger = logging.getLogger(__name__)


@dataclass
class RoutingResult:
    """
    Result of a single net routing operation.
    """
    net_id: str
    """ Identifier for the net """

    path: list[ComponentResult]
    """ List of moves forming the path """

    is_valid: bool
    """ Whether the path is collision-free and reached the target """

    collisions: int
    """ Number of detected collisions/overlaps """

    reached_target: bool = False
    """ Whether the final port matches the target port """


class PathFinder:
    """
    Multi-net router using Negotiated Congestion.
    """
    __slots__ = ('router', 'cost_evaluator', 'max_iterations', 'base_congestion_penalty', 'use_tiered_strategy')

    router: AStarRouter
    """ The A* search engine """

    cost_evaluator: CostEvaluator
    """ The evaluator for path costs """

    max_iterations: int
    """ Maximum number of rip-up and reroute iterations """

    base_congestion_penalty: float
    """ Starting penalty for overlaps """

    use_tiered_strategy: bool
    """ If True, use simpler collision models in early iterations for speed """

    def __init__(
            self,
            router: AStarRouter,
            cost_evaluator: CostEvaluator,
            max_iterations: int = 10,
            base_congestion_penalty: float = 100.0,
            use_tiered_strategy: bool = True,
            ) -> None:
        """
        Initialize the PathFinder.

        Args:
            router: The A* search engine.
            cost_evaluator: The evaluator for path costs.
            max_iterations: Maximum number of rip-up and reroute iterations.
            base_congestion_penalty: Starting penalty for overlaps.
            use_tiered_strategy: Whether to use simplified collision models in early iterations.
        """
        self.router = router
        self.cost_evaluator = cost_evaluator
        self.max_iterations = max_iterations
        self.base_congestion_penalty = base_congestion_penalty
        self.use_tiered_strategy = use_tiered_strategy

    def route_all(
            self,
            netlist: dict[str, tuple[Port, Port]],
            net_widths: dict[str, float],
            store_expanded: bool = False,
            iteration_callback: Callable[[int, dict[str, RoutingResult]], None] | None = None,
            ) -> dict[str, RoutingResult]:
        """
        Route all nets in the netlist using Negotiated Congestion.

        Args:
            netlist: Mapping of net_id to (start_port, target_port).
            net_widths: Mapping of net_id to waveguide width.
            store_expanded: Whether to store expanded nodes for the last iteration.
            iteration_callback: Optional callback(iteration_idx, current_results).

        Returns:
            Mapping of net_id to RoutingResult.
        """
        results: dict[str, RoutingResult] = {}
        self.cost_evaluator.congestion_penalty = self.base_congestion_penalty

        start_time = time.monotonic()
        num_nets = len(netlist)
        session_timeout = max(60.0, 10.0 * num_nets * self.max_iterations)

        for iteration in range(self.max_iterations):
            any_congestion = False
            logger.info(f'PathFinder Iteration {iteration}...')

            # Sequence through nets
            for net_id, (start, target) in netlist.items():
                # Timeout check
                elapsed = time.monotonic() - start_time
                if elapsed > session_timeout:
                    logger.warning(f'PathFinder TIMEOUT after {elapsed:.2f}s')
                    return self._finalize_results(results, netlist)

                width = net_widths.get(net_id, 2.0)

                # 1. Rip-up existing path
                self.cost_evaluator.collision_engine.remove_path(net_id)

                # 2. Reroute with current congestion info
                target_coll_model = self.router.config.bend_collision_type
                coll_model = target_coll_model
                skip_cong = False
                if self.use_tiered_strategy and iteration == 0:
                    skip_cong = True
                    if target_coll_model == "arc":
                        coll_model = "clipped_bbox"
                
                # Dynamic node limit: increase if it failed previously
                base_node_limit = self.router.config.node_limit
                current_node_limit = base_node_limit
                if net_id in results and not results[net_id].reached_target:
                     current_node_limit = base_node_limit * (iteration + 1)
                
                net_start = time.monotonic()
                # Store expanded only in the last potential iteration or if specifically requested
                do_store = store_expanded and (iteration == self.max_iterations - 1)
                
                # Temporarily override node_limit
                original_limit = self.router.node_limit
                self.router.node_limit = current_node_limit
                
                path = self.router.route(start, target, width, net_id=net_id, bend_collision_type=coll_model, return_partial=True, store_expanded=do_store, skip_congestion=skip_cong)
                
                # Restore
                self.router.node_limit = original_limit
                
                logger.debug(f'  Net {net_id} routed in {time.monotonic() - net_start:.4f}s using {coll_model}')

                if path:
                    # Check if reached exactly
                    last_p = path[-1].end_port
                    reached = (abs(last_p.x - target.x) < 1e-6 and
                               abs(last_p.y - target.y) < 1e-6 and
                               abs(last_p.orientation - target.orientation) < 0.1)
                    
                    all_geoms = []
                    all_dilated = []
                    # 3. Add to index ONLY if it reached the target
                    # (Prevents failed paths from blocking others forever)
                    if reached:
                        for res in path:
                            all_geoms.extend(res.geometry)
                            if res.dilated_geometry:
                                all_dilated.extend(res.dilated_geometry)
                            else:
                                # Fallback dilation
                                dilation = self.cost_evaluator.collision_engine.clearance / 2.0
                                all_dilated.extend([p.buffer(dilation) for p in res.geometry])
                                
                        self.cost_evaluator.collision_engine.add_path(net_id, all_geoms, dilated_geometry=all_dilated)

                    # Check if this new path has any congestion
                    collision_count = 0
                    # Always check for congestion to decide if more iterations are needed
                    if reached:
                        for i, poly in enumerate(all_geoms):
                            overlaps = self.cost_evaluator.collision_engine.check_congestion(
                                poly, net_id, dilated_geometry=all_dilated[i]
                            )
                            if overlaps > 0:
                                collision_count += overlaps

                    if collision_count > 0:
                        any_congestion = True

                    results[net_id] = RoutingResult(net_id, path, (reached and collision_count == 0), collision_count, reached_target=reached)
                else:
                    results[net_id] = RoutingResult(net_id, [], False, 0, reached_target=False)
                    any_congestion = True

            if iteration_callback:
                iteration_callback(iteration, results)

            if not any_congestion:
                # Check if all reached target
                all_reached = all(r.reached_target for r in results.values())
                if all_reached:
                    break

            # 4. Inflate congestion penalty
            self.cost_evaluator.congestion_penalty *= 1.5

        return self._finalize_results(results, netlist)

    def _finalize_results(
            self,
            results: dict[str, RoutingResult],
            netlist: dict[str, tuple[Port, Port]],
            ) -> dict[str, RoutingResult]:
        """
        Final check: re-verify all nets against the final static paths.

        Args:
            results: Results from the routing loop.
            netlist: The original netlist.

        Returns:
            Refined results with final collision counts.
        """
        logger.debug(f'Finalizing results for nets: {list(results.keys())}')
        final_results = {}
        for net_id in netlist:
            res = results.get(net_id)
            if not res or not res.path:
                final_results[net_id] = RoutingResult(net_id, [], False, 0)
                continue

            collision_count = 0
            for comp in res.path:
                for i, poly in enumerate(comp.geometry):
                    dil_poly = comp.dilated_geometry[i] if comp.dilated_geometry else None
                    overlaps = self.cost_evaluator.collision_engine.check_collision(
                        poly, net_id, buffer_mode='congestion', dilated_geometry=dil_poly
                    )
                    if isinstance(overlaps, int):
                        collision_count += overlaps

            reached = False
            if res.path:
                target_p = netlist[net_id][1]
                last_p = res.path[-1].end_port
                reached = (abs(last_p.x - target_p.x) < 1e-6 and
                           abs(last_p.y - target_p.y) < 1e-6 and
                           abs(last_p.orientation - target_p.orientation) < 0.1)

            final_results[net_id] = RoutingResult(net_id, res.path, (collision_count == 0 and reached), collision_count, reached_target=reached)

        return final_results
