from __future__ import annotations

from collections.abc import Hashable, Iterable, Mapping
from heapq import heappop, heappush
from itertools import count
from typing import TypeVar

Node = TypeVar("Node", bound=Hashable)
Graph = Mapping[Node, Iterable[tuple[Node, float]]]


def dijkstra(
    graph: Graph[Node],
    source: Node,
    target: Node | None = None,
) -> tuple[dict[Node, float], dict[Node, Node | None]]:
    """Return shortest distances and predecessor links from source."""
    distances: dict[Node, float] = {node: float("inf") for node in graph}
    previous: dict[Node, Node | None] = {source: None}
    distances[source] = 0.0

    tie_breaker = count()
    queue: list[tuple[float, int, Node]] = [(0.0, next(tie_breaker), source)]

    while queue:
        current_distance, _, current = heappop(queue)

        # Skip stale heap entries created before a better path was found.
        if current_distance != distances[current]:
            continue

        if target is not None and current == target:
            break

        for neighbor, weight in graph.get(current, ()):
            if weight < 0:
                raise ValueError("Dijkstra's algorithm requires non-negative edge weights")

            new_distance = current_distance + weight
            if new_distance < distances.get(neighbor, float("inf")):
                distances[neighbor] = new_distance
                previous[neighbor] = current
                heappush(queue, (new_distance, next(tie_breaker), neighbor))

    return distances, previous


def shortest_path(graph: Graph[Node], source: Node, target: Node) -> list[Node]:
    """Reconstruct the shortest path from source to target."""
    distances, previous = dijkstra(graph, source, target)

    if distances.get(target, float("inf")) == float("inf"):
        return []

    path: list[Node] = []
    current: Node | None = target
    while current is not None:
        path.append(current)
        current = previous[current]

    return list(reversed(path))


def _test_shortest_path() -> None:
    graph: dict[str, list[tuple[str, float]]] = {
        "A": [("B", 4), ("C", 1)],
        "B": [("D", 1)],
        "C": [("B", 2), ("D", 5)],
        "D": [],
    }

    distances, _ = dijkstra(graph, "A")
    assert distances["D"] == 4
    assert shortest_path(graph, "A", "D") == ["A", "C", "B", "D"]


def _test_unreachable() -> None:
    graph: dict[str, list[tuple[str, float]]] = {
        "A": [("B", 1)],
        "B": [],
        "C": [],
    }

    distances, _ = dijkstra(graph, "A")
    assert distances["C"] == float("inf")
    assert shortest_path(graph, "A", "C") == []


if __name__ == "__main__":
    _test_shortest_path()
    _test_unreachable()

    example_graph: dict[str, list[tuple[str, float]]] = {
        "A": [("B", 4), ("C", 1)],
        "B": [("D", 1)],
        "C": [("B", 2), ("D", 5)],
        "D": [],
    }

    print(shortest_path(example_graph, "A", "D"))