#!/usr/bin/env python3
"""A* Pathfinding Algorithm Reference Script

This script demonstrates a practical, production-friendly A* implementation for
finding an efficient route from a single start node to a single goal node.

Features:
- Python 3 standard-library only
- argparse-based CLI
- Input validation
- Support for weighted directed graphs
- Optional heuristic input via JSON file
- Path reconstruction
- Safe defaults and no external dependencies

Example usage:
  python astar_pathfinding.py --graph graph.json --start A --goal G
  python astar_pathfinding.py --use-demo --start A --goal G

Graph JSON format:
{
  "nodes": ["A", "B", "C"],
  "edges": [
    {"from": "A", "to": "B", "cost": 1.0},
    {"from": "B", "to": "C", "cost": 2.5}
  ],
  "heuristic": {
    "A": 3.0,
    "B": 1.5,
    "C": 0.0
  }
}

Notes:
- Heuristic values should be admissible for exact shortest-path behavior.
- Edge costs must be non-negative.
- The graph is treated as directed unless you add both directions explicitly.
"""

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 Any, Dict, Hashable, Iterable, List, Optional, Tuple


Node = Hashable
Graph = Dict[Node, List[Tuple[Node, float]]]


@dataclass
class SearchResult:
    path: List[Node]
    cost: float
    expanded_nodes: int


class GraphError(ValueError):
    pass


def validate_non_negative_cost(cost: Any) -> float:
    if not isinstance(cost, (int, float)):
        raise GraphError(f"Edge cost must be numeric, got {type(cost).__name__}")
    if math.isnan(cost) or math.isinf(cost):
        raise GraphError("Edge cost must be a finite number")
    if cost < 0:
        raise GraphError("A* requires non-negative edge costs")
    return float(cost)


def load_graph_from_json(path: str) -> Tuple[Graph, Dict[Node, float]]:
    try:
        with open(path, "r", encoding="utf-8") as f:
            payload = json.load(f)
    except OSError as exc:
        raise GraphError(f"Could not read graph file: {exc}") from exc
    except json.JSONDecodeError as exc:
        raise GraphError(f"Invalid JSON in graph file: {exc}") from exc

    if not isinstance(payload, dict):
        raise GraphError("Graph JSON must be an object")

    nodes = payload.get("nodes", [])
    edges = payload.get("edges", [])
    heuristic = payload.get("heuristic", {})

    if not isinstance(nodes, list):
        raise GraphError('"nodes" must be a list')
    if not isinstance(edges, list):
        raise GraphError('"edges" must be a list')
    if not isinstance(heuristic, dict):
        raise GraphError('"heuristic" must be an object')

    graph: Graph = defaultdict(list)
    for node in nodes:
        graph[node] = []

    for idx, edge in enumerate(edges, start=1):
        if not isinstance(edge, dict):
            raise GraphError(f"Edge #{idx} must be an object")
        try:
            src = edge["from"]
            dst = edge["to"]
            cost = validate_non_negative_cost(edge["cost"])
        except KeyError as exc:
            raise GraphError(f"Edge #{idx} is missing required field {exc}") from exc
        graph[src].append((dst, cost))
        if dst not in graph:
            graph[dst] = []

    heuristic_map: Dict[Node, float] = {}
    for node, value in heuristic.items():
        heuristic_map[node] = validate_non_negative_cost(value)

    return dict(graph), heuristic_map


def build_demo_graph() -> Tuple[Graph, Dict[Node, float]]:
    # A small, readable example graph.
    graph: Graph = {
        "A": [("B", 1.0), ("C", 4.0)],
        "B": [("D", 2.0), ("E", 5.0)],
        "C": [("E", 1.0)],
        "D": [("G", 3.0)],
        "E": [("G", 1.0)],
        "G": [],
    }

    # Example heuristic values; in production these should be tied to your domain.
    heuristic = {
        "A": 5.0,
        "B": 4.0,
        "C": 2.0,
        "D": 2.0,
        "E": 1.0,
        "G": 0.0,
    }
    return graph, heuristic


def heuristic_value(heuristic: Dict[Node, float], node: Node) -> float:
    return float(heuristic.get(node, 0.0))


def reconstruct_path(came_from: Dict[Node, Node], current: Node) -> List[Node]:
    path = [current]
    while current in came_from:
        current = came_from[current]
        path.append(current)
    path.reverse()
    return path


def astar_search(graph: Graph, start: Node, goal: Node, heuristic: Dict[Node, float]) -> SearchResult:
    if start not in graph:
        raise GraphError(f"Start node {start!r} not found in graph")
    if goal not in graph:
        raise GraphError(f"Goal node {goal!r} not found in graph")

    open_heap: List[Tuple[float, float, Node]] = []
    heappush(open_heap, (heuristic_value(heuristic, start), 0.0, start))

    came_from: Dict[Node, Node] = {}
    g_score: Dict[Node, float] = {start: 0.0}
    closed_set = set()
    expanded_nodes = 0

    while open_heap:
        current_f, current_g, current = heappop(open_heap)

        if current in closed_set:
            continue

        expanded_nodes += 1
        if current == goal:
            return SearchResult(
                path=reconstruct_path(came_from, current),
                cost=current_g,
                expanded_nodes=expanded_nodes,
            )

        closed_set.add(current)

        for neighbor, edge_cost in graph.get(current, []):
            if edge_cost < 0:
                raise GraphError("A* requires non-negative edge costs")

            tentative_g = current_g + edge_cost
            if tentative_g < g_score.get(neighbor, math.inf):
                came_from[neighbor] = current
                g_score[neighbor] = tentative_g
                f_score = tentative_g + heuristic_value(heuristic, neighbor)
                heappush(open_heap, (f_score, tentative_g, neighbor))

    raise GraphError(f"No path found from {start!r} to {goal!r}")


def parse_args(argv: Optional[Iterable[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Demonstrate A* pathfinding on a weighted graph.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--graph", help="Path to a JSON file containing nodes, edges, and optional heuristic")
    parser.add_argument("--start", required=True, help="Start node")
    parser.add_argument("--goal", required=True, help="Goal node")
    parser.add_argument(
        "--use-demo",
        action="store_true",
        help="Use a built-in demo graph instead of loading a JSON file",
    )
    return parser.parse_args(argv)


def main(argv: Optional[Iterable[str]] = None) -> int:
    args = parse_args(argv)

    if args.use_demo:
        graph, heuristic = build_demo_graph()
    else:
        if not args.graph:
            print("Error: provide --graph PATH or use --use-demo", file=sys.stderr)
            return 2
        graph, heuristic = load_graph_from_json(args.graph)

    try:
        result = astar_search(graph, args.start, args.goal, heuristic)
    except GraphError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1

    print("A* search completed successfully")
    print(f"Start: {args.start}")
    print(f"Goal: {args.goal}")
    print(f"Expanded nodes: {result.expanded_nodes}")
    print(f"Total cost: {result.cost}")
    print("Path: " + " -> ".join(str(node) for node in result.path))
    return 0


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