"""
A* routing engine on the .yoromaps road graph.

Loads the graph adjacency list into memory for fast routing.
Typical memory usage: ~100-200 MB for a country like Mali or Togo.
"""

from __future__ import annotations

import heapq
import math
import sqlite3
from dataclasses import dataclass


def _haversine(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    R = 6_371_000
    dlat = math.radians(lat2 - lat1)
    dlon = math.radians(lon2 - lon1)
    a = (math.sin(dlat / 2) ** 2 +
         math.cos(math.radians(lat1)) * math.cos(math.radians(lat2)) *
         math.sin(dlon / 2) ** 2)
    return R * 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))


@dataclass
class RouteResult:
    """Result of a routing query."""

    distance_km: float
    duration_min: float
    nodes: list[int]
    geometry: dict  # GeoJSON LineString
    steps: list[dict]
    found: bool = True


class Graph:
    """In-memory road graph loaded from a .yoromaps database.

    Load once, route many times::

        graph = Graph.from_db(conn)
        r1 = graph.route(6.13, 1.22, 9.55, 1.18)
        r2 = graph.route(6.13, 1.22, 8.98, 1.13)
    """

    def __init__(self):
        self.coords: dict[int, tuple[float, float]] = {}  # node_id → (lat, lon)
        self.adj: dict[int, list[tuple[int, float, float, str]]] = {}  # node_id → [(to, dist_m, dur_s, name)]

    @classmethod
    def from_db(cls, conn: sqlite3.Connection) -> "Graph":
        """Load the entire graph into memory from a .yoromaps database."""
        g = cls()

        for row in conn.execute("SELECT id, lat, lon FROM nodes"):
            g.coords[row[0]] = (row[1], row[2])

        for row in conn.execute("SELECT from_node, to_node, distance_m, duration_s, name FROM edges"):
            from_n, to_n, dist, dur, name = row
            if from_n not in g.adj:
                g.adj[from_n] = []
            g.adj[from_n].append((to_n, dist, dur, name or ""))

        return g

    def _nearest(self, lat: float, lon: float) -> int | None:
        """Find the nearest node (brute-force on a subset for speed)."""
        best_id = None
        best_dist = float("inf")

        # First try a tight box
        for delta in (0.01, 0.05, 0.1, 0.5):
            for nid, (nlat, nlon) in self.coords.items():
                if abs(nlat - lat) > delta or abs(nlon - lon) > delta:
                    continue
                d = (nlat - lat) ** 2 + (nlon - lon) ** 2
                if d < best_dist:
                    best_dist = d
                    best_id = nid
            if best_id is not None:
                return best_id

        return None

    def route(
        self,
        start_lat: float,
        start_lon: float,
        end_lat: float,
        end_lon: float,
    ) -> RouteResult:
        """Find the shortest route between two GPS coordinates using A*."""

        start_node = self._nearest(start_lat, start_lon)
        end_node = self._nearest(end_lat, end_lon)

        empty = RouteResult(
            distance_km=0, duration_min=0, nodes=[],
            geometry={"type": "LineString", "coordinates": []},
            steps=[], found=False,
        )

        if not start_node or not end_node:
            return empty

        if start_node == end_node:
            c = self.coords[start_node]
            return RouteResult(
                distance_km=0, duration_min=0, nodes=[start_node],
                geometry={"type": "LineString", "coordinates": [[c[1], c[0]]]},
                steps=[], found=True,
            )

        end_lat_n, end_lon_n = self.coords[end_node]

        # A* algorithm — in-memory, fast
        open_set: list[tuple[float, int]] = [(0.0, start_node)]
        came_from: dict[int, int] = {}
        g_score: dict[int, float] = {start_node: 0.0}
        duration: dict[int, float] = {start_node: 0.0}
        visited: set[int] = set()

        while open_set:
            _, current = heapq.heappop(open_set)

            if current == end_node:
                path = [current]
                while current in came_from:
                    current = came_from[current]
                    path.append(current)
                path.reverse()

                return RouteResult(
                    distance_km=round(g_score[end_node] / 1000, 1),
                    duration_min=round(duration[end_node] / 60, 1),
                    nodes=path,
                    geometry=self._build_geometry(path),
                    steps=self._build_steps(path),
                    found=True,
                )

            if current in visited:
                continue
            visited.add(current)

            for to_node, dist_m, dur_s, _ in self.adj.get(current, []):
                if to_node in visited:
                    continue

                tentative_g = g_score[current] + dist_m
                if tentative_g < g_score.get(to_node, float("inf")):
                    came_from[to_node] = current
                    g_score[to_node] = tentative_g
                    duration[to_node] = duration[current] + dur_s

                    nlat, nlon = self.coords.get(to_node, (0, 0))
                    h = _haversine(nlat, nlon, end_lat_n, end_lon_n)
                    heapq.heappush(open_set, (tentative_g + h, to_node))

        return empty

    def _build_geometry(self, node_ids: list[int]) -> dict:
        coords = []
        for nid in node_ids:
            c = self.coords.get(nid)
            if c:
                coords.append([c[1], c[0]])  # [lon, lat]
        return {"type": "LineString", "coordinates": coords}

    def _build_steps(self, node_ids: list[int]) -> list[dict]:
        steps: list[dict] = []
        prev_name = ""
        seg_distance = 0.0

        for i in range(len(node_ids) - 1):
            n_from = node_ids[i]
            n_to = node_ids[i + 1]

            # Find edge
            name = ""
            dist = 0.0
            for to, d, _, nm in self.adj.get(n_from, []):
                if to == n_to:
                    name = nm or "road"
                    dist = d
                    break

            if name != prev_name and prev_name:
                steps.append({
                    "instruction": f"Continue on {prev_name}",
                    "distance_m": round(seg_distance),
                    "name": prev_name,
                })
                seg_distance = 0.0

            seg_distance += dist
            prev_name = name

        if prev_name:
            steps.append({
                "instruction": f"Arrive via {prev_name}",
                "distance_m": round(seg_distance),
                "name": prev_name,
            })

        return steps


# ---------- Convenience functions ----------

def _nearest_node(conn: sqlite3.Connection, lat: float, lon: float) -> int | None:
    """Find the nearest graph node to a GPS coordinate (SQL-based, for tests)."""
    for delta in (0.01, 0.05, 0.1, 0.5, 1.0):
        row = conn.execute("""
            SELECT id, lat, lon FROM nodes
            WHERE lat BETWEEN ? AND ? AND lon BETWEEN ? AND ?
            ORDER BY ((lat - ?) * (lat - ?) + (lon - ?) * (lon - ?))
            LIMIT 1
        """, (lat - delta, lat + delta, lon - delta, lon + delta,
              lat, lat, lon, lon)).fetchone()
        if row:
            return row["id"]
    return None


def route(
    conn: sqlite3.Connection,
    start_lat: float,
    start_lon: float,
    end_lat: float,
    end_lon: float,
) -> RouteResult:
    """Find the shortest route between two GPS coordinates.

    Loads the graph into memory on each call. For multiple routes,
    use ``Graph.from_db()`` directly.
    """
    graph = Graph.from_db(conn)
    return graph.route(start_lat, start_lon, end_lat, end_lon)


def route_from_codes(conn: sqlite3.Connection, codes: list[str]) -> list[RouteResult]:
    """Route through multiple Yoro codes in order.

    Returns a list of RouteResult, one per leg.
    """
    import yoro

    graph = Graph.from_db(conn)
    points = [(yoro.decode(c)["lat"], yoro.decode(c)["lon"]) for c in codes]
    return [
        graph.route(points[i][0], points[i][1], points[i + 1][0], points[i + 1][1])
        for i in range(len(points) - 1)
    ]
