#!/usr/bin/env python3
"""Download transit route lines + stops from Overpass, per mode.

For each mode in config/modes.yaml, runs an Overpass query (with mirror
fallback + exponential backoff) and writes:
    data/raw/osm/{mode}.gpkg          (layers: "lines", "stops", EPSG:4326)
    data/raw/osm/{mode}.overpassql    (the exact query text)

Lines: one MultiLineString per route relation (concatenation of member way
geometries), carrying ref/name/network/colour/operator/route/mode tags.
Stops: one Point per platform node member, carrying ref/name/mode/route_refs
(route refs of routes using this stop, for disambiguation).
"""
import sys
import time
from pathlib import Path

import geopandas as gpd
import requests
import yaml
from shapely.geometry import LineString, MultiLineString, Point

from _common import CONFIG, OSM_RAW, PROCESSED, UA

MIRRORS = [
    "https://overpass-api.de/api/interpreter",
    "https://overpass.kumi.systems/api/interpreter",
    "https://overpass.private.coffee/api/interpreter",
]

PLATFORM_ROLES = {"platform", "platform_entry", "platform_exit"}
STOP_ROLES = {"stop", "stop_entry", "stop_exit", "stop_entry_only", "stop_exit_only"}
# bus stops: use platform nodes (where passengers wait)
# rail stops: use stop nodes (platforms are often areas, stop nodes are reliable points)
BUS_STOP_ROLES = PLATFORM_ROLES
RAIL_STOP_ROLES = STOP_ROLES | PLATFORM_ROLES


def build_query(mode_cfg, bbox):
    s, n, w, e = bbox
    filters = []
    for k, v in mode_cfg["osm"].items():
        if k.endswith("_neq"):
            filters.append(f'["{k[:-4]}"!="{v}"]')
        else:
            filters.append(f'["{k}"="{v}"]')
    filt = "".join(filters)
    return f"""[out:json][timeout:180];
relation{filt}({s},{w},{n},{e});
out geom;
>;
out body qt;
"""


def overpass(ql):
    last_err = None
    for mi, mirror in enumerate(MIRRORS):
        for attempt in range(4):
            try:
                r = requests.post(
                    mirror,
                    data={"data": ql},
                    headers={"User-Agent": UA},
                    timeout=300,
                )
                if r.status_code == 429 or r.status_code >= 500:
                    raise RuntimeError(f"HTTP {r.status_code}")
                r.raise_for_status()
                return r.json()
            except Exception as e:
                last_err = e
                wait = 2 ** (attempt + mi)
                print(f"  [{mirror}] attempt {attempt+1} failed: {e}; retry in {wait}s", flush=True)
                time.sleep(wait)
    raise RuntimeError(f"all mirrors failed: {last_err}")


def parse(data, mode):
    is_bus = mode == "bus"
    stop_roles = BUS_STOP_ROLES if is_bus else RAIL_STOP_ROLES

    nodes = {}
    for e in data["elements"]:
        if e["type"] == "node":
            nodes[e["id"]] = e

    lines = []
    # collect route_refs per stop node: {node_id: set(route_refs)}
    node_routes = {}
    for e in data["elements"]:
        if e["type"] != "relation" or "tags" not in e:
            continue
        tags = e["tags"]
        if tags.get("route") is None:
            continue
        route_ref = tags.get("ref")
        # collect line geometries
        segs = []
        for m in e["members"]:
            if m["type"] == "way" and "geometry" in m and m.get("role") not in PLATFORM_ROLES:
                coords = [(g["lon"], g["lat"]) for g in m["geometry"]]
                if len(coords) >= 2:
                    segs.append(LineString(coords))
        geom = MultiLineString(segs) if len(segs) > 1 else (segs[0] if segs else None)
        if geom is None:
            continue
        lines.append({
            "mode": mode,
            "ref": route_ref,
            "name": tags.get("name"),
            "network": tags.get("network"),
            "colour": tags.get("colour"),
            "operator": tags.get("operator"),
            "route": tags.get("route"),
            "geometry": geom,
        })
        # record route_ref on each stop member node
        for m in e["members"]:
            if m["type"] != "node" or m.get("role") not in stop_roles:
                continue
            nid = m["ref"]
            if nid not in nodes:
                continue
            node_routes.setdefault(nid, set())
            if route_ref:
                node_routes[nid].add(route_ref)

    # build stops from collected nodes
    stops = []
    stop_seen = set()
    for nid, routes in node_routes.items():
        if nid in stop_seen:
            continue
        nd = nodes[nid]
        if "lat" not in nd or "lon" not in nd:
            continue
        stop_seen.add(nid)
        ntags = nd.get("tags", {})
        stops.append({
            "mode": mode,
            "name": ntags.get("name") or ntags.get("public_transport") or "stop",
            "ref": ntags.get("ref"),
            "public_transport": ntags.get("public_transport"),
            "railway": ntags.get("railway"),
            "route_refs": ";".join(sorted(routes)),
            "geometry": Point(nd["lon"], nd["lat"]),
        })
    return lines, stops


def stations_query(bbox):
    s, n, w, e = bbox
    return f"""[out:json][timeout:180];
(
  node["railway"="station"]({s},{w},{n},{e});
  node["public_transport"="station"]({s},{w},{n},{e});
  way["railway"="station"]({s},{w},{n},{e});
  way["public_transport"="station"]({s},{w},{n},{e});
);
out center;
"""


def parse_stations(data):
    """Stations (nodes or area centroids) with proper human-readable names."""
    rows = []
    seen = set()
    for e in data["elements"]:
        if e["type"] not in ("node", "way"):
            continue
        tags = e.get("tags", {})
        name = tags.get("name")
        if not name:
            continue
        if e["type"] == "node":
            geom = Point(e["lon"], e["lat"])
        else:
            c = e.get("center")
            if not c:
                continue
            geom = Point(c["lon"], c["lat"])
        key = (name, round(geom.x, 5), round(geom.y, 5))
        if key in seen:
            continue
        seen.add(key)
        rows.append({
            "name": name,
            "railway": tags.get("railway"),
            "public_transport": tags.get("public_transport"),
            "station": tags.get("station"),
            "subway": tags.get("subway"),
            "light_rail": tags.get("light_rail"),
            "train": tags.get("train"),
            "uic_ref": tags.get("uic_ref"),
            "geometry": geom,
        })
    return rows


def orphan_bus_stops_query(bbox):
    s, n, w, e = bbox
    return f"""[out:json][timeout:180];
(
  node["highway"="bus_stop"]["public_transport"="platform"]({s},{w},{n},{e});
);
out body;
"""


def parse_orphan_bus_stops(data, known_node_ids):
    """Parse standalone bus platform nodes not already captured by routes."""
    stops = []
    for e in data["elements"]:
        if e["type"] != "node" or e["id"] in known_node_ids:
            continue
        tags = e.get("tags", {})
        if tags.get("public_transport") != "platform":
            continue
        name = tags.get("name")
        if not name:
            continue
        stops.append({
            "mode": "bus",
            "name": name,
            "ref": tags.get("ref"),
            "public_transport": tags.get("public_transport"),
            "railway": tags.get("railway"),
            "route_refs": "",  # orphan — no route membership
            "geometry": Point(e["lon"], e["lat"]),
        })
    return stops


def main():
    modes = yaml.safe_load((CONFIG / "modes.yaml").read_text())["modes"]
    area = gpd.read_file(PROCESSED / "area.gpkg")
    w, s, e, n = area.total_bounds  # (minx, miny, maxx, maxy) = (west, south, east, north)
    bbox = (s, n, w, e)

    OSM_RAW.mkdir(parents=True, exist_ok=True)
    bus_stop_node_ids = set()
    for mode, cfg in modes.items():
        ql = build_query(cfg, bbox)
        (OSM_RAW / f"{mode}.overpassql").write_text(ql)
        print(f"[{mode}] querying Overpass bbox=({s:.4f},{w:.4f},{n:.4f},{e:.4f})...", flush=True)
        data = overpass(ql)
        lines, stops = parse(data, mode)
        # track bus stop node ids to avoid dupes with orphans
        if mode == "bus":
            for e in data["elements"]:
                if e["type"] == "node" and e.get("tags", {}).get("public_transport") == "platform":
                    bus_stop_node_ids.add(e["id"])
        print(f"  -> {len(lines)} lines, {len(stops)} stops", flush=True)
        if not lines:
            continue
        lgdf = gpd.GeoDataFrame(lines, crs="EPSG:4326")
        lgdf.to_file(OSM_RAW / f"{mode}.gpkg", driver="GPKG", layer="lines")
        if stops:
            sgdf = gpd.GeoDataFrame(stops, crs="EPSG:4326")
            sgdf.to_file(OSM_RAW / f"{mode}.gpkg", driver="GPKG", layer="stops")

    # orphan bus stops (platform nodes not in any route relation)
    oql = orphan_bus_stops_query(bbox)
    (OSM_RAW / "orphan_bus_stops.overpassql").write_text(oql)
    print("[orphan_bus_stops] querying Overpass...", flush=True)
    odata = overpass(oql)
    orphans = parse_orphan_bus_stops(odata, bus_stop_node_ids)
    print(f"  -> {len(orphans)} orphan bus stops (not in route relations)", flush=True)
    if orphans:
        # append to bus.gpkg stops layer
        existing = gpd.read_file(OSM_RAW / "bus.gpkg", layer="stops")
        combined = gpd.GeoDataFrame(
            __import__("pandas").concat([existing, gpd.GeoDataFrame(orphans, crs="EPSG:4326")],
                                        ignore_index=True),
            crs="EPSG:4326"
        )
        # rewrite stops layer
        import tempfile, shutil
        tmp = OSM_RAW / "bus_tmp.gpkg"
        lines_gdf = gpd.read_file(OSM_RAW / "bus.gpkg", layer="lines")
        lines_gdf.to_file(tmp, driver="GPKG", layer="lines")
        combined.to_file(tmp, driver="GPKG", layer="stops")
        shutil.move(str(tmp), str(OSM_RAW / "bus.gpkg"))
        print(f"  -> bus stops layer now: {len(combined)} total", flush=True)

    # stations (named station nodes/areas) for labelling
    sql = stations_query(bbox)
    (OSM_RAW / "stations.overpassql").write_text(sql)
    print("[stations] querying Overpass...", flush=True)
    sdata = overpass(sql)
    st = parse_stations(sdata)
    print(f"  -> {len(st)} named stations", flush=True)
    if st:
        gpd.GeoDataFrame(st, crs="EPSG:4326").to_file(
            OSM_RAW / "stations.gpkg", driver="GPKG", layer="stations"
        )
    print("done.")


if __name__ == "__main__":
    main()