import pytest

from inire.geometry.collision import CollisionEngine
from inire.geometry.components import Bend90, Straight
from inire.geometry.primitives import Port
from inire.router.astar import AStarContext
from inire.router.cost import CostEvaluator
from inire.router.danger_map import DangerMap
from inire.router.outcomes import RoutingOutcome
from inire.router.pathfinder import PathFinder, RoutingResult
from inire.router.session import (
    create_routing_session_state,
    prepare_routing_session_state,
    run_routing_iteration,
)


@pytest.fixture
def basic_evaluator() -> CostEvaluator:
    engine = CollisionEngine(clearance=2.0)
    danger_map = DangerMap(bounds=(0, 0, 100, 100))
    danger_map.precompute([])
    return CostEvaluator(engine, danger_map)


def _build_manual_path(start: Port, width: float, clearance: float, steps: list[tuple[str, float | str]]) -> list:
    path = []
    curr = start
    dilation = clearance / 2.0
    for kind, value in steps:
        if kind == "B":
            comp = Bend90.generate(curr, 5.0, width, value, dilation=dilation)
        else:
            comp = Straight.generate(curr, value, width, dilation=dilation)
        path.append(comp)
        curr = comp.end_port
    return path


def test_pathfinder_parallel(basic_evaluator: CostEvaluator) -> None:
    context = AStarContext(basic_evaluator)
    pf = PathFinder(context)

    netlist = {
        "net1": (Port(0, 0, 0), Port(50, 0, 0)),
        "net2": (Port(0, 10, 0), Port(50, 10, 0)),
    }
    net_widths = {"net1": 2.0, "net2": 2.0}

    results = pf.route_all(netlist, net_widths)

    assert len(results) == 2
    assert results["net1"].is_valid
    assert results["net2"].is_valid
    assert results["net1"].collisions == 0
    assert results["net2"].collisions == 0


def test_pathfinder_crossing_detection(basic_evaluator: CostEvaluator) -> None:
    context = AStarContext(basic_evaluator)
    # Force a crossing by setting low iterations and low penalty
    pf = PathFinder(context, max_iterations=1, base_congestion_penalty=1.0, warm_start=None)

    # Net 1: (0, 25) -> (100, 25) Horizontal
    # Net 2: (50, 0) -> (50, 50) Vertical
    netlist = {
        "net1": (Port(0, 25, 0), Port(100, 25, 0)),
        "net2": (Port(50, 0, 90), Port(50, 50, 90)),
    }
    net_widths = {"net1": 2.0, "net2": 2.0}

    results = pf.route_all(netlist, net_widths)

    # Both should be invalid because they cross
    assert not results["net1"].is_valid
    assert not results["net2"].is_valid
    assert results["net1"].collisions > 0
    assert results["net2"].collisions > 0


def test_prepare_routing_session_state_builds_warm_start_and_sorts_nets(
    basic_evaluator: CostEvaluator,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    context = AStarContext(basic_evaluator)
    pf = PathFinder(context)
    calls: list[tuple[str, list[str]]] = []
    cleared: list[bool] = []

    def fake_build(
        netlist: dict[str, tuple[Port, Port]],
        net_widths: dict[str, float],
        order: str,
    ) -> dict[str, list]:
        calls.append((order, list(netlist.keys())))
        return {"warm": []}

    monkeypatch.setattr(PathFinder, "_build_greedy_warm_start_paths", lambda self, netlist, net_widths, order: fake_build(netlist, net_widths, order))
    monkeypatch.setattr(AStarContext, "clear_static_caches", lambda self: cleared.append(True))

    netlist = {
        "short": (Port(0, 0, 0), Port(10, 0, 0)),
        "long": (Port(0, 0, 0), Port(40, 10, 0)),
        "mid": (Port(0, 0, 0), Port(20, 0, 0)),
    }
    state = create_routing_session_state(
        pf,
        netlist,
        {net_id: 2.0 for net_id in netlist},
        store_expanded=False,
        iteration_callback=None,
        shuffle_nets=False,
        sort_nets="longest",
        initial_paths=None,
        seed=None,
    )

    prepare_routing_session_state(pf, state)

    assert calls == [("longest", ["short", "long", "mid"])]
    assert cleared == [True]
    assert state.initial_paths == {"warm": []}
    assert state.all_net_ids == ["long", "mid", "short"]


def test_run_routing_iteration_updates_results_and_invokes_callback(
    basic_evaluator: CostEvaluator,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    context = AStarContext(basic_evaluator)
    pf = PathFinder(context)
    callback_results: list[dict[str, RoutingResult]] = []

    def fake_route_once(
        net_id: str,
        start: Port,
        target: Port,
        width: float,
        iteration: int,
        initial_paths: dict[str, list] | None,
        store_expanded: bool,
        needs_self_collision_check: set[str],
    ) -> tuple[RoutingResult, RoutingOutcome]:
        _ = (start, target, width, iteration, initial_paths, store_expanded, needs_self_collision_check)
        result = RoutingResult(
            net_id,
            [],
            net_id == "net1",
            int(net_id == "net2"),
            reached_target=True,
            outcome="completed" if net_id == "net1" else "colliding",
        )
        return result, result.outcome

    monkeypatch.setattr(
        PathFinder,
        "_route_net_once",
        lambda self, net_id, start, target, width, iteration, initial_paths, store_expanded, needs_self_collision_check: fake_route_once(
            net_id,
            start,
            target,
            width,
            iteration,
            initial_paths,
            store_expanded,
            needs_self_collision_check,
        ),
    )
    state = create_routing_session_state(
        pf,
        {"net1": (Port(0, 0, 0), Port(10, 0, 0)), "net2": (Port(0, 10, 0), Port(10, 10, 0))},
        {"net1": 2.0, "net2": 2.0},
        store_expanded=True,
        iteration_callback=lambda iteration, results: callback_results.append(dict(results)),
        shuffle_nets=False,
        sort_nets=None,
        initial_paths={"seeded": []},
        seed=None,
    )

    outcomes = run_routing_iteration(pf, state, iteration=0)

    assert outcomes == {"net1": "completed", "net2": "colliding"}
    assert set(state.results) == {"net1", "net2"}
    assert callback_results and set(callback_results[0]) == {"net1", "net2"}
    assert state.results["net1"].is_valid
    assert not state.results["net2"].is_valid
    assert state.results["net2"].outcome == "colliding"


def test_run_routing_iteration_timeout_finalizes_tree(
    basic_evaluator: CostEvaluator,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    context = AStarContext(basic_evaluator)
    pf = PathFinder(context)
    finalized: list[bool] = []
    monkeypatch.setattr(type(pf.path_state), "finalize_dynamic_tree", lambda self: finalized.append(True))

    state = create_routing_session_state(
        pf,
        {"net1": (Port(0, 0, 0), Port(10, 0, 0))},
        {"net1": 2.0},
        store_expanded=False,
        iteration_callback=None,
        shuffle_nets=False,
        sort_nets=None,
        initial_paths={},
        seed=None,
    )
    state.start_time = 0.0
    state.session_timeout = 0.0

    result = run_routing_iteration(pf, state, iteration=0)

    assert result is None
    assert finalized == [True]


def test_route_all_retries_partial_paths_across_iterations(
    basic_evaluator: CostEvaluator,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    context = AStarContext(basic_evaluator)
    pf = PathFinder(context, max_iterations=3, warm_start=None, refine_paths=False)
    calls: list[int] = []

    class FakeComponent:
        def __init__(self, start_port: Port, end_port: Port) -> None:
            self.start_port = start_port
            self.end_port = end_port

    def fake_route_astar(
        start: Port,
        target: Port,
        width: float,
        *,
        context: AStarContext,
        metrics: object,
        net_id: str,
        bend_collision_type: str,
        return_partial: bool,
        store_expanded: bool,
        skip_congestion: bool,
        self_collision_check: bool,
        node_limit: int,
    ) -> list[FakeComponent]:
        _ = (
            width,
            context,
            metrics,
            net_id,
            bend_collision_type,
            return_partial,
            store_expanded,
            skip_congestion,
            self_collision_check,
            node_limit,
        )
        calls.append(len(calls))
        if len(calls) == 1:
            return [FakeComponent(start, Port(5, 0, 0))]
        return [FakeComponent(start, target)]

    monkeypatch.setattr("inire.router.pathfinder.route_astar", fake_route_astar)
    monkeypatch.setattr(type(pf.path_state), "install_path", lambda self, net_id, path: None)
    monkeypatch.setattr(type(pf.path_state), "remove_path", lambda self, net_id: None)
    monkeypatch.setattr(
        type(pf.path_state),
        "verify_path_report",
        lambda self, net_id, path: basic_evaluator.collision_engine.verify_path_report(net_id, []),
    )
    monkeypatch.setattr(type(pf.path_state), "finalize_dynamic_tree", lambda self: None)

    results = pf.route_all({"net": (Port(0, 0, 0), Port(10, 0, 0))}, {"net": 2.0})

    assert calls == [0, 1]
    assert results["net"].reached_target
    assert results["net"].is_valid
    assert results["net"].outcome == "completed"


def test_pathfinder_refine_paths_reduces_locked_detour_bends() -> None:
    bounds = (0, -50, 100, 50)

    def build_pathfinder(*, refine_paths: bool) -> tuple[CollisionEngine, PathFinder]:
        engine = CollisionEngine(clearance=2.0)
        danger_map = DangerMap(bounds=bounds)
        danger_map.precompute([])
        evaluator = CostEvaluator(engine, danger_map, bend_penalty=250.0, sbend_penalty=500.0)
        context = AStarContext(evaluator, bend_radii=[10.0])
        return engine, PathFinder(context, refine_paths=refine_paths)

    base_engine, base_pf = build_pathfinder(refine_paths=False)
    base_pf.route_all({"netA": (Port(10, 0, 0), Port(90, 0, 0))}, {"netA": 2.0})
    base_engine.lock_net("netA")
    base_result = base_pf.route_all({"netB": (Port(50, -20, 90), Port(50, 20, 90))}, {"netB": 2.0})["netB"]

    refined_engine, refined_pf = build_pathfinder(refine_paths=True)
    refined_pf.route_all({"netA": (Port(10, 0, 0), Port(90, 0, 0))}, {"netA": 2.0})
    refined_engine.lock_net("netA")
    refined_result = refined_pf.route_all({"netB": (Port(50, -20, 90), Port(50, 20, 90))}, {"netB": 2.0})["netB"]

    base_bends = sum(1 for comp in base_result.path if comp.move_type == "Bend90")
    refined_bends = sum(1 for comp in refined_result.path if comp.move_type == "Bend90")

    assert base_result.is_valid
    assert refined_result.is_valid
    assert refined_bends < base_bends
    assert refined_pf._path_cost(refined_result.path) < base_pf._path_cost(base_result.path)


def test_pathfinder_refine_paths_simplifies_triple_crossing_detours() -> None:
    bounds = (0, 0, 100, 100)
    netlist = {
        "horizontal": (Port(10, 50, 0), Port(90, 50, 0)),
        "vertical_up": (Port(45, 10, 90), Port(45, 90, 90)),
        "vertical_down": (Port(55, 90, 270), Port(55, 10, 270)),
    }
    net_widths = {net_id: 2.0 for net_id in netlist}

    def build_pathfinder(*, refine_paths: bool) -> PathFinder:
        engine = CollisionEngine(clearance=2.0)
        danger_map = DangerMap(bounds=bounds)
        danger_map.precompute([])
        evaluator = CostEvaluator(engine, danger_map, greedy_h_weight=1.5, bend_penalty=250.0, sbend_penalty=500.0)
        context = AStarContext(evaluator, bend_radii=[10.0], sbend_radii=[10.0])
        return PathFinder(context, base_congestion_penalty=1000.0, refine_paths=refine_paths)

    base_results = build_pathfinder(refine_paths=False).route_all(netlist, net_widths)
    refined_results = build_pathfinder(refine_paths=True).route_all(netlist, net_widths)

    for net_id in ("vertical_up", "vertical_down"):
        base_result = base_results[net_id]
        refined_result = refined_results[net_id]
        base_bends = sum(1 for comp in base_result.path if comp.move_type == "Bend90")
        refined_bends = sum(1 for comp in refined_result.path if comp.move_type == "Bend90")

        assert base_result.is_valid
        assert refined_result.is_valid
        assert refined_bends < base_bends


def test_refine_path_handles_same_orientation_lateral_offset() -> None:
    engine = CollisionEngine(clearance=2.0)
    danger_map = DangerMap(bounds=(-20, -20, 120, 120))
    danger_map.precompute([])
    evaluator = CostEvaluator(engine, danger_map, bend_penalty=250.0, sbend_penalty=500.0)
    context = AStarContext(evaluator, bend_radii=[5.0, 10.0])
    pf = PathFinder(context, refine_paths=True)

    start = Port(0, 0, 0)
    width = 2.0
    path = _build_manual_path(
        start,
        width,
        engine.clearance,
        [
            ("B", "CCW"),
            ("S", 10.0),
            ("B", "CW"),
            ("S", 20.0),
            ("B", "CW"),
            ("S", 10.0),
            ("B", "CCW"),
            ("S", 10.0),
            ("B", "CCW"),
            ("S", 5.0),
            ("B", "CW"),
        ],
    )
    target = path[-1].end_port

    refined = pf._refine_path("net", start, target, width, path)

    assert target == Port(60, 15, 0)
    assert sum(1 for comp in path if comp.move_type == "Bend90") == 6
    assert sum(1 for comp in refined if comp.move_type == "Bend90") == 4
    assert refined[-1].end_port == target
    assert pf._path_cost(refined) < pf._path_cost(path)


def test_refine_path_can_simplify_subpath_with_different_global_orientation() -> None:
    engine = CollisionEngine(clearance=2.0)
    danger_map = DangerMap(bounds=(-20, -20, 120, 120))
    danger_map.precompute([])
    evaluator = CostEvaluator(engine, danger_map, bend_penalty=250.0, sbend_penalty=500.0)
    context = AStarContext(evaluator, bend_radii=[5.0, 10.0])
    pf = PathFinder(context, refine_paths=True)

    start = Port(0, 0, 0)
    width = 2.0
    path = _build_manual_path(
        start,
        width,
        engine.clearance,
        [
            ("B", "CCW"),
            ("S", 10.0),
            ("B", "CW"),
            ("S", 20.0),
            ("B", "CW"),
            ("S", 10.0),
            ("B", "CCW"),
            ("S", 10.0),
            ("B", "CCW"),
            ("S", 5.0),
            ("B", "CW"),
            ("B", "CCW"),
            ("S", 10.0),
        ],
    )
    target = path[-1].end_port

    refined = pf._refine_path("net", start, target, width, path)

    assert target == Port(65, 30, 90)
    assert sum(1 for comp in path if comp.move_type == "Bend90") == 7
    assert sum(1 for comp in refined if comp.move_type == "Bend90") == 5
    assert refined[-1].end_port == target
    assert pf._path_cost(refined) < pf._path_cost(path)