Trainpulse/trainpulse/prepare.py

"""Fahrplan-Update: holt für jede Quelle (Schweiz, Deutschland) den GTFS-Fahrplan, filtert auf Bahn
und rechnet mit pfaedle die echten Streckenverläufe (shapes) aus OpenStreetMap. Läuft stündlich
und macht für eine Quelle nichts, solange dort kein neuer Fahrplan veröffentlicht wurde.

    python -m trainpulse.prepare                       # alle Quellen
    python -m trainpulse.prepare --source de --force   # eine Quelle neu aufbereiten
    python -m trainpulse.prepare --source ch --zip X   # lokale Datei statt Download (zum Testen)
"""
import argparse
import csv
import datetime as dt
import io
import json
import logging
import shutil
import subprocess
import time
import zipfile
from array import array
from pathlib import Path

import httpx

from . import border
from . import config as C
from .feed import route_mode

log = logging.getLogger("prepare")


# Aufbereitungs-Version: ändert sich, was aus dem Fahrplan übernommen wird, rechnet der Updater
# alle Fahrpläne einmal neu, auch ohne neuen Fahrplan (2: mit Trams und U-Bahnen)
FORMAT = 2


def timetable_year(today: dt.date) -> int:
    # Fahrplanwechsel ist am 2. Sonntag im Dezember
    if today.month == 12:
        first = dt.date(today.year, 12, 1)
        if today >= first + dt.timedelta(days=(6 - first.weekday()) % 7 + 7):
            return today.year + 1
    return today.year


def gtfs_url(src: str) -> str:
    url = C.SOURCES[src]["gtfs_url"]
    if src == "ch" and not url:
        # shortcut: in der Nacht des Fahrplanwechsels fehlen bis zum nächsten Stundenlauf die Züge
        url = f"https://data.opentransportdata.swiss/dataset/timetable-{timetable_year(dt.date.today())}-gtfs2020/permalink"
    return url


def download(url: str, dest: Path) -> None:
    tmp = dest.with_suffix(dest.suffix + ".part")
    with httpx.stream("GET", url, follow_redirects=True, timeout=120,
                      headers={"User-Agent": C.USER_AGENT}) as r:
        r.raise_for_status()
        total, done, t0, last = int(r.headers.get("content-length", 0)), 0, time.time(), time.time()
        with open(tmp, "wb") as f:
            for chunk in r.iter_bytes(1 << 20):
                f.write(chunk)
                done += len(chunk)
                if time.time() - last > 30:  # Fortschritt, damit ein langsamer Download nicht wie ein Hänger aussieht
                    last = time.time()
                    log.info("  … %d von %d MB (%.1f MB/s)", done >> 20, total >> 20, done / (last - t0) / 1e6)
    tmp.rename(dest)


def remote_id(src: str, url: str) -> tuple[str, str | None]:
    """Kennung der aktuell veröffentlichten Datei, ohne sie herunterzuladen: (Name, Datum JJJJMMTT oder None)."""
    if src == "ch":
        return latest_gtfs_name(url), None
    r = httpx.head(url, follow_redirects=True, timeout=60, headers={"User-Agent": C.USER_AGENT})
    r.raise_for_status()
    lm = r.headers.get("last-modified")
    if not lm:
        raise RuntimeError(f"{url}: kein Last-Modified-Header")
    date = dt.datetime.strptime(lm, "%a, %d %b %Y %H:%M:%S GMT").strftime("%Y%m%d")
    return f"{src}-{lm}", date


def latest_gtfs_name(url: str) -> str:
    # Der Permalink leitet auf .../gtfs_fp2026_YYYYMMDD.zip weiter. HEAD geht beim Speicher nicht,
    # darum nur die erste Weiterleitung lesen – so sehen wir ohne 300-MB-Download, ob es Neues gibt.
    r = httpx.get(url, follow_redirects=False, timeout=60, headers={"User-Agent": C.USER_AGENT})
    if r.status_code not in (301, 302, 303, 307, 308):
        r.raise_for_status()
        raise RuntimeError(f"Keine Weiterleitung von {url} (Status {r.status_code})")
    return r.headers["location"].rsplit("/", 1)[-1].split("?")[0]


def rows(z: zipfile.ZipFile, name: str):
    reader = csv.reader(io.TextIOWrapper(z.open(name), encoding="utf-8-sig", newline=""))
    header = next(reader)
    return header, reader


def write_csv(path: Path, header, data) -> int:
    n = 0
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(header)
        for row in data:
            w.writerow(row)
            n += 1
    return n


def filter_by_col(z: zipfile.ZipFile, name: str, col: str, keys: set, out: Path, collect: str | None = None):
    """Filter für die Riesendateien (stop_times.txt hat mehrere GB): behält Zeilen, deren Spalte `col`
    in `keys` ist, und schreibt sie unverändert weg. Steht `col` vorne, wird nur das erste Feld als
    Rohtext verglichen (schnell). Gibt (Anzahl, Werte der Spalte `collect`) zurück."""
    n, seen = 0, set()
    with z.open(name) as raw, open(out, "w", newline="", encoding="utf-8") as f:
        lines = io.TextIOWrapper(raw, encoding="utf-8-sig", newline="")
        header = next(csv.reader([lines.readline()]))
        f.write(",".join(header) + "\n")
        ki, ci = header.index(col), header.index(collect) if collect else None
        for line in lines:
            if ki == 0:
                if line.split(",", 1)[0].strip('"') not in keys:
                    continue
                row = next(csv.reader([line])) if ci is not None else None
            else:
                row = next(csv.reader([line]))
                if row[ki] not in keys:
                    continue
            f.write(line)
            n += 1
            if ci is not None:
                seen.add(row[ci])
    return n, seen


def filter_rail(zip_path: Path, out: Path, version: str | None = None) -> dict:
    out.mkdir(parents=True)
    z = zipfile.ZipFile(zip_path)

    h, rd = rows(z, "routes.txt")
    routes = [r for r in rd if route_mode(r[h.index("route_type")])]
    write_csv(out / "routes.txt", h, routes)
    route_ids = {r[h.index("route_id")] for r in routes}

    h, rd = rows(z, "trips.txt")
    trips = [r for r in rd if r[h.index("route_id")] in route_ids]
    if "shape_id" in h:  # shapes.txt kopieren wir nicht (pfaedle rechnet neu) -> Verweise leeren (z. B. OVapi)
        for r in trips:
            r[h.index("shape_id")] = ""
    write_csv(out / "trips.txt", h, trips)
    trip_ids = {r[h.index("trip_id")] for r in trips}
    service_ids = {r[h.index("service_id")] for r in trips}

    if "frequencies.txt" in z.namelist():
        fh, frd = rows(z, "frequencies.txt")
        n = sum(1 for r in frd if r[fh.index("trip_id")] in trip_ids)
        if n:
            log.warning("%d Bahnfahrten nur über frequencies.txt – werden ignoriert", n)

    n_st, used_stops = filter_by_col(z, "stop_times.txt", "trip_id", trip_ids, out / "stop_times.txt", "stop_id")

    h, rd = rows(z, "stops.txt")
    sid, par = h.index("stop_id"), h.index("parent_station")
    stops = list(rd)
    parents = {r[par] for r in stops if r[sid] in used_stops}
    write_csv(out / "stops.txt", h, (r for r in stops if r[sid] in used_stops or r[sid] in parents))

    for name in ("calendar.txt", "calendar_dates.txt"):
        if name in z.namelist():
            filter_by_col(z, name, "service_id", service_ids, out / name)
    for name in ("agency.txt", "feed_info.txt"):
        if name in z.namelist():
            h, rd = rows(z, name)
            write_csv(out / name, h, rd)

    h, rd = rows(z, "feed_info.txt") if "feed_info.txt" in z.namelist() else ([], iter([]))
    info = dict(zip(h, next(rd, [])))
    fv = info.get("feed_version", "")
    return {"feed_version": fv if fv.isdigit() else (version or zip_path.stem),  # gtfs.de: "latest-free"
            "routes": len(routes), "trips": len(trips), "stop_times": n_st}


def ensure_osm() -> Path:
    """Schienennetz aus OpenStreetMap für alle Länder: herunterladen, mit osmium auf Bahn-Objekte
    filtern (Deutschland: ~4 GB -> wenige 100 MB) und zusammenführen. Monatlich erneuert."""
    d = C.DATA_DIR / "osm"
    rail, urls_file = d / "rail.osm.pbf", d / "urls.txt"
    wanted = "\n".join(C.OSM_URLS)
    if rail.exists() and time.time() - rail.stat().st_mtime < 30 * 86400 \
            and urls_file.exists() and urls_file.read_text() == wanted:
        return rail
    d.mkdir(parents=True, exist_ok=True)
    parts = []
    for i, url in enumerate(C.OSM_URLS):
        raw, part = d / f"download-{i}.osm.pbf", d / f"rail-{i}.osm.pbf"
        log.info("Lade OSM-Daten %s", url)
        download(url, raw)
        log.info("Filtere Schienennetz …")
        subprocess.run([C.OSMIUM, "tags-filter", str(raw), "nwr/railway", "nwr/public_transport", "r/route",
                        "-o", str(part), "--overwrite"], check=True)
        raw.unlink()  # spart Platz
        parts.append(str(part))
    merged, tmp = d / "rail.merged.osm.pbf", d / "rail.tmp.osm.pbf"
    subprocess.run([C.OSMIUM, "merge", *parts, "-o", str(merged), "--overwrite"], check=True)
    # Die Länder-Dateien sind unterschiedlich alt: Objekte an den Grenzen stecken dann in zwei
    # Versionen drin, woran osmium extract scheitert. time-filter behält je Objekt die neueste.
    subprocess.run([C.OSMIUM, "time-filter", str(merged), "-o", str(tmp), "--overwrite"], check=True)
    # IDs durchnummerieren: osmium extract braucht sonst ~4 GB RAM für die weit verstreuten OSM-IDs
    subprocess.run([C.OSMIUM, "renumber", str(tmp), "-o", str(merged), "--overwrite"], check=True)
    tmp.unlink()
    merged.rename(rail)
    urls_file.write_text(wanted)
    for p in parts:
        Path(p).unlink()
    for old in d.glob("region.osm.pbf"):  # Datei der alten Version (nur Schweiz)
        old.unlink()
    return rail


def run_pfaedle(work: Path, rail: Path, source: str) -> Path | None:
    if not shutil.which(C.PFAEDLE):
        log.warning("pfaedle nicht gefunden – Züge fahren Luftlinie zwischen den Haltestellen")
        return None
    # Netz aufs Land (+ Rand) zuschneiden: die Fahrpläne enthalten Züge bis Paris oder Hamburg,
    # pfaedle würde sonst das Netz aller Länder laden. Ausserhalb zeigt ohnehin der Nachbar-Fahrplan.
    osm = work / "region.osm.pbf"
    subprocess.run([C.OSMIUM, "extract", "-b", border.bbox(source), str(ensure_osm()), "-o", str(osm),
                    "--overwrite"], check=True)
    log.info("[%s] pfaedle rechnet Streckenverläufe …", source)
    t = time.time()
    res = subprocess.run([C.PFAEDLE, "-D", "-m", "rail,tram,subway", "-x", str(osm), str(rail)], cwd=work)
    osm.unlink()
    out = work / "gtfs-out"
    if res.returncode != 0 or not (out / "shapes.txt").exists():
        log.error("[%s] pfaedle fehlgeschlagen (Code %s) – nutze Fahrplan ohne Streckenverläufe", source, res.returncode)
        return None
    log.info("[%s] pfaedle fertig in %.0f s", source, time.time() - t)
    return out


def pack_shapes(d: Path) -> None:
    """shapes.txt ist für die Schweiz ~2.6 GB (56 Mio. Punkte) – viel zu gross für den Speicher.
    Umwandeln in shapes.bin (float32: lon, lat, Meter ab Start) + shapes.json (id -> [Start, Anzahl]),
    der Server liest daraus nur die Strecken, die er gerade braucht."""
    src = d / "shapes.txt"
    if not src.exists():
        return
    index, cur, pts, n = {}, None, [], 0
    with open(d / "shapes.bin", "wb") as out:
        def flush():
            nonlocal n
            if cur is not None:
                pts.sort()
                array("f", [v for p in pts for v in p[1:]]).tofile(out)
                index[cur] = [n, len(pts)]
                n += len(pts)
        with open(src, encoding="utf-8-sig", newline="") as f:
            r = csv.reader(f)
            h = {k: i for i, k in enumerate(next(r))}
            sid, lat, lon, seq = h["shape_id"], h["shape_pt_lat"], h["shape_pt_lon"], h["shape_pt_sequence"]
            dist = h["shape_dist_traveled"]
            for row in r:
                if row[sid] != cur:
                    flush()
                    assert row[sid] not in index, "shapes.txt nicht nach shape_id gruppiert"
                    cur, pts = row[sid], []
                pts.append((int(row[seq]), float(row[lon]), float(row[lat]), float(row[dist])))
            flush()
    (d / "shapes.json").write_text(json.dumps(index, separators=(",", ":")))
    src.unlink()


def ensure_tiles() -> None:
    """Kartenkacheln (OpenStreetMap, von Protomaps) einmalig für die Region herunterladen
    (erneut, wenn sich der Ausschnitt geändert hat)."""
    dest = C.DATA_DIR / "tiles" / "region.pmtiles"
    info = C.DATA_DIR / "tiles" / "region.txt"
    wanted = f"{C.TILES_BBOX} z{C.TILES_MAXZOOM}"
    if dest.exists() and dest.stat().st_size and info.exists() and info.read_text() == wanted:
        return
    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp = dest.with_suffix(".tmp")
    for days in range(1, 8):  # Protomaps baut täglich, ältere Builds werden gelöscht
        d = (dt.date.today() - dt.timedelta(days=days)).strftime("%Y%m%d")
        log.info("Lade Kartenkacheln vom %s …", d)
        r = subprocess.run([C.PMTILES, "extract", f"https://build.protomaps.com/{d}.pmtiles", str(tmp),
                            f"--bbox={C.TILES_BBOX}", f"--maxzoom={C.TILES_MAXZOOM}"])
        if r.returncode == 0 and tmp.exists() and tmp.stat().st_size:
            tmp.rename(dest)
            info.write_text(wanted)
            return
    log.error("Kartenkacheln konnten nicht geladen werden – Karte bleibt ohne Hintergrund")


def publish(source: str, src: Path, meta: dict) -> None:
    (src / "meta.json").write_text(json.dumps(meta))
    base = C.FEEDS_DIR / source
    base.mkdir(parents=True, exist_ok=True)
    dest = base / meta["feed_version"]
    if dest.exists():
        shutil.rmtree(dest)
    src.rename(dest)  # gleiche Partition -> atomar; der Server sieht nur fertige Ordner
    for old in sorted(p for p in base.iterdir() if p.is_dir())[:-2]:
        shutil.rmtree(old)


def prepare(source: str, state: dict, zip_path: Path | None = None, force: bool = False) -> None:
    work = C.DATA_DIR / "work"
    shutil.rmtree(work, ignore_errors=True)
    work.mkdir()

    version = None
    if zip_path:
        name = zip_path.name
    else:
        url = gtfs_url(source)
        name, version = remote_id(source, url)
        if f"{name} #{FORMAT}" == state.get(source) and not force:
            log.info("[%s] Kein neuer Fahrplan (%s)", source, name)
            return
        zip_path = work / "gtfs.zip"
        log.info("[%s] Lade Fahrplan %s", source, name)
        download(url, zip_path)

    t = time.time()
    meta = filter_rail(zip_path, work / "rail", version)
    log.info("[%s] Bahn-Filter: %s in %.0f s", source, meta, time.time() - t)
    if zip_path.parent == work:
        zip_path.unlink()

    shaped = run_pfaedle(work, work / "rail", source)
    meta["shapes"] = shaped is not None
    if shaped:
        pack_shapes(shaped)
    meta["source_file"] = name
    meta["prepared_at"] = int(time.time())
    publish(source, shaped or work / "rail", meta)
    shutil.rmtree(work, ignore_errors=True)
    state[source] = f"{name} #{FORMAT}"
    log.info("[%s] Fahrplan %s bereit", source, meta["feed_version"])


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--source", choices=list(C.SOURCES), help="nur diese Quelle (Standard: alle aktiven)")
    ap.add_argument("--zip", type=Path, help="lokale GTFS-Zip statt Download (nur mit --source)")
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    if args.zip and not args.source:
        ap.error("--zip braucht --source")

    C.FEEDS_DIR.mkdir(parents=True, exist_ok=True)
    for p in C.FEEDS_DIR.iterdir():
        if not p.is_dir() or p.name in C.SOURCES:
            continue
        if p.name.isdigit() and (p / "meta.json").exists():  # alte Version: feeds/<version> war die Schweiz
            (C.FEEDS_DIR / "ch").mkdir(exist_ok=True)
            p.rename(C.FEEDS_DIR / "ch" / p.name)
        else:
            shutil.rmtree(p)
    if not args.zip:
        ensure_tiles()
    state_file = C.DATA_DIR / "state.json"
    state = json.loads(state_file.read_text()) if state_file.exists() else {}

    failed = False
    for source in [args.source] if args.source else C.ENABLED:
        try:
            prepare(source, state, args.zip, args.force)
        except Exception:  # eine kaputte Quelle soll die anderen nicht aufhalten
            log.exception("[%s] Aufbereitung fehlgeschlagen", source)
            failed = True
        state_file.write_text(json.dumps(state))
    if failed:
        raise SystemExit(1)


if __name__ == "__main__":
    main()