#!/usr/bin/env python3
"""dijkstra_optimized.py

A practical, vendor-neutral reference implementation of Dijkstra's algorithm
with common performance optimizations:

- adjacency list graph representation
- min-heap priority queue
- stale-entry skipping
- optional early exit for a single target node
- basic validation for non-negative edge weights

This script is intended as a reusable starting point for production code.
It does not include any environment-specific dependencies or secrets.
"""

from __future__ import annotations

import argparse
import json
import math
import sys
from collections import defaultdict
from dataclasses import dataclass
from heapq import heappop, heappush
from typing import DefaultDict, Dict, Iterable, List, Optional, Sequence, Tuple

Node = str
Weight = float
Edge = Tuple[Node, Weight]
Graph = Dict[Node, List[Edge]]


@dataclass
class DijkstraStats:
    relaxations: int = 0
    queue_pushes: int = 0
    queue_pops: int = 0
    stale_skips: int = 0


def build_graph(edges: Sequence[Sequence[object]], bidirectional: bool = False) -> Graph:
    """Build an adjacency list from a sequence of [u, v, weight] records."""
    graph: DefaultDict[Node, List[Edge]] = defaultdict(list)

    for idx, edge in enumerate(edges, start=1):
        if len(edge) != 3:
            raise ValueError(f"Edge #{idx} must have exactly 3 items: [from, to, weight]")

        u, v, weight = edge
        if not isinstance(u, str) or not isinstance(v, str):
            raise TypeError(f"Edge #{idx} nodes must be strings")
        if not isinstance(weight, (int, float)):
            raise TypeError(f"Edge #{idx} weight must be numeric")
        if weight < 0:
            raise ValueError("Dijkstra's algorithm requires non-negative edge weights")

        graph[u].append((v, float(weight)))
        graph.setdefault(v, [])
        if bidirectional:
            graph[v].append((u, float(weight)))

    return dict(graph)


def dijkstra(graph: Graph, source: Node, target: Optional[Node] = None) -> Tuple[Dict[Node, float], Dict[Node, Optional[Node]], DijkstraStats]:
    """Compute shortest paths from source using an optimized Dijkstra implementation."""
    if source not in graph:
        raise ValueError(f"Source node '{source}' is not present in the graph")
    if target is not None and target not in graph:
        raise ValueError(f"Target node '{target}' is not present in the graph")

    dist: Dict[Node, float] = {node: math.inf for node in graph}
    prev: Dict[Node, Optional[Node]] = {node: None for node in graph}
    stats = DijkstraStats()

    dist[source] = 0.0
    pq: List[Tuple[float, Node]] = [(0.0, source)]
    stats.queue_pushes += 1

    while pq:
        current_dist, u = heappop(pq)
        stats.queue_pops += 1

        if current_dist != dist[u]:
            stats.stale_skips += 1
            continue

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

        for v, weight in graph[u]:
            stats.relaxations += 1
            candidate = current_dist + weight
            if candidate < dist[v]:
                dist[v] = candidate
                prev[v] = u
                heappush(pq, (candidate, v))
                stats.queue_pushes += 1

    return dist, prev, stats


def reconstruct_path(prev: Dict[Node, Optional[Node]], source: Node, target: Node) -> List[Node]:
    """Reconstruct a shortest path from source to target."""
    if source == target:
        return [source]

    path: List[Node] = []
    cur: Optional[Node] = target
    while cur is not None:
        path.append(cur)
        if cur == source:
            break
        cur = prev.get(cur)

    if not path or path[-1] != source:
        return []

    path.reverse()
    return path


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Optimized Dijkstra shortest-path reference script"
    )
    parser.add_argument(
        "--input",
        required=True,
        help="Path to a JSON file containing graph data",
    )
    parser.add_argument(
        "--source",
        required=True,
        help="Source node name",
    )
    parser.add_argument(
        "--target",
        default=None,
        help="Optional target node name for early exit and path reconstruction",
    )
    parser.add_argument(
        "--bidirectional",
        action="store_true",
        help="Treat each edge as undirected by adding reverse edges",
    )
    parser.add_argument(
        "--show-stats",
        action="store_true",
        help="Print internal performance counters",
    )
    return parser.parse_args(argv)


def load_graph_from_json(path: str, bidirectional: bool = False) -> Graph:
    """Load a graph from JSON.

    Expected format:
    {
      "edges": [
        ["A", "B", 3],
        ["B", "C", 5]
      ]
    }
    """
    try:
        with open(path, "r", encoding="utf-8") as f:
            payload = json.load(f)
    except FileNotFoundError as exc:
        raise FileNotFoundError(f"Input file not found: {path}") from exc
    except json.JSONDecodeError as exc:
        raise ValueError(f"Invalid JSON in input file: {exc}") from exc

    if not isinstance(payload, dict):
        raise ValueError("Input JSON must be an object with an 'edges' field")

    edges = payload.get("edges")
    if not isinstance(edges, list):
        raise ValueError("Input JSON must contain an 'edges' array")

    return build_graph(edges, bidirectional=bidirectional)


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)
    graph = load_graph_from_json(args.input, bidirectional=args.bidirectional)

    dist, prev, stats = dijkstra(graph, args.source, target=args.target)

    if args.target is None:
        output = {
            "source": args.source,
            "distances": dist,
        }
    else:
        path = reconstruct_path(prev, args.source, args.target)
        output = {
            "source": args.source,
            "target": args.target,
            "distance": dist.get(args.target, math.inf),
            "path": path,
        }

    print(json.dumps(output, indent=2, sort_keys=True))

    if args.show_stats:
        print(
            json.dumps(
                {
                    "relaxations": stats.relaxations,
                    "queue_pushes": stats.queue_pushes,
                    "queue_pops": stats.queue_pops,
                    "stale_skips": stats.stale_skips,
                },
                indent=2,
                sort_keys=True,
            ),
            file=sys.stderr,
        )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())