"""
Extract road graph from OSM PBF into a .yoromaps SQLite database.

Uses pyosmium to parse the PBF file in a single pass, extracting
highway ways and their nodes into a graph suitable for routing.
"""

from __future__ import annotations

import math
import sqlite3
from pathlib import Path

try:
    import osmium
except ImportError:
    osmium = None  # type: ignore[assignment]

# Road types we care about, with estimated speeds (km/h)
HIGHWAY_SPEEDS: dict[str, float] = {
    "motorway": 110, "motorway_link": 60,
    "trunk": 90, "trunk_link": 50,
    "primary": 70, "primary_link": 40,
    "secondary": 60, "secondary_link": 35,
    "tertiary": 50, "tertiary_link": 30,
    "residential": 30,
    "unclassified": 40,
    "living_street": 20,
    "service": 20,
    "track": 15,
    "path": 5,
}


def haversine(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    """Distance in meters between two GPS points."""
    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))


def _check_osmium():
    if osmium is None:
        raise ImportError("Install extract extra: pip install yoro-maps[extract]")


def extract_graph(pbf_path: str | Path, conn: sqlite3.Connection, progress=None) -> dict:
    """Extract road graph from a PBF file into the database.

    Args:
        pbf_path: Path to the .osm.pbf file.
        conn: Open SQLite connection to a .yoromaps database.
        progress: Optional callable(message, current, total) for progress reporting.

    Returns:
        Dict with keys: ``nodes``, ``edges``.
    """
    _check_osmium()

    class _NodeCollector(osmium.SimpleHandler):
        def __init__(self):
            super().__init__()
            self.highway_node_ids: set[int] = set()
            self.ways_data: list[tuple] = []

        def way(self, w):
            tags = dict(w.tags)
            hw = tags.get("highway")
            if hw not in HIGHWAY_SPEEDS:
                return
            node_ids = [n.ref for n in w.nodes]
            if len(node_ids) < 2:
                return
            oneway = tags.get("oneway", "no") in ("yes", "true", "1")
            name = tags.get("name", "")
            self.ways_data.append((node_ids, hw, oneway, name))
            self.highway_node_ids.update(node_ids)

    class _NodeLocator(osmium.SimpleHandler):
        def __init__(self, wanted_ids: set[int]):
            super().__init__()
            self.wanted = wanted_ids
            self.coords: dict[int, tuple[float, float]] = {}

        def node(self, n):
            if n.id in self.wanted:
                self.coords[n.id] = (n.location.lat, n.location.lon)

    pbf_path = str(pbf_path)

    # Pass 1: collect ways and their node references
    if progress:
        progress("Scanning ways...", 0, 2)
    collector = _NodeCollector()
    collector.apply_file(pbf_path)

    if progress:
        progress(f"Found {len(collector.ways_data)} roads, {len(collector.highway_node_ids)} nodes", 1, 2)

    # Pass 2: collect coordinates for relevant nodes
    locator = _NodeLocator(collector.highway_node_ids)
    locator.apply_file(pbf_path, locations=True)

    if progress:
        progress("Building graph...", 2, 2)

    # Insert nodes
    node_rows = [(nid, lat, lon) for nid, (lat, lon) in locator.coords.items()]
    conn.executemany("INSERT OR IGNORE INTO nodes (id, lat, lon) VALUES (?, ?, ?)", node_rows)

    # Insert edges
    edge_count = 0
    edge_rows = []
    for node_ids, hw_type, oneway, name in collector.ways_data:
        speed_kmh = HIGHWAY_SPEEDS[hw_type]
        speed_ms = speed_kmh / 3.6

        for i in range(len(node_ids) - 1):
            n1, n2 = node_ids[i], node_ids[i + 1]
            c1 = locator.coords.get(n1)
            c2 = locator.coords.get(n2)
            if not c1 or not c2:
                continue

            dist = haversine(c1[0], c1[1], c2[0], c2[1])
            duration = dist / speed_ms if speed_ms > 0 else 0

            # Forward edge
            edge_rows.append((n1, n2, dist, duration, hw_type, 1 if oneway else 0, name))

            # Reverse edge (if not oneway)
            if not oneway:
                edge_rows.append((n2, n1, dist, duration, hw_type, 0, name))

            edge_count += 1

    conn.executemany(
        "INSERT INTO edges (from_node, to_node, distance_m, duration_s, road_type, oneway, name) "
        "VALUES (?, ?, ?, ?, ?, ?, ?)",
        edge_rows,
    )
    conn.commit()

    return {"nodes": len(node_rows), "edges": edge_count}
