Trainpulse/trainpulse/prepare.py

"""Fahrplan-Update: holt den Schweizer GTFS-Fahrplan, filtert auf Bahn und rechnet mit pfaedle
die echten Streckenverläufe (shapes) aus OpenStreetMap. Läuft stündlich per systemd-Timer und
macht nichts, solange opentransportdata.swiss keinen neuen Fahrplan veröffentlicht hat.

    python -m trainpulse.prepare            # normaler Lauf
    python -m trainpulse.prepare --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 pathlib import Path

import httpx

from . import config as C

log = logging.getLogger("prepare")


def is_rail(route_type: str) -> bool:
    t = int(route_type)
    return t == 2 or 100 <= t <= 199  # GTFS + "extended route types" (Schweiz nutzt 100–117)


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() -> str:
    # shortcut: in der Nacht des Fahrplanwechsels fehlen bis zum nächsten Stundenlauf die Züge
    return C.GTFS_URL or (
        f"https://data.opentransportdata.swiss/dataset/timetable-{timetable_year(dt.date.today())}-gtfs2020/permalink")


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()
        with open(tmp, "wb") as f:
            for chunk in r.iter_bytes(1 << 20):
                f.write(chunk)
    tmp.rename(dest)


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_first_col(z: zipfile.ZipFile, name: str, keys: set, out: Path, collect: str | None = None):
    """Schnellfilter für die Riesendateien (stop_times.txt hat ~3.7 GB): vergleicht nur das erste Feld
    als Rohtext, schreibt Treffer unverändert weg. 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()]))
        assert header[0] in ("trip_id", "service_id"), header
        f.write(",".join(header) + "\n")
        ci = header.index(collect) if collect else None
        for line in lines:
            if line.split(",", 1)[0].strip('"') in keys:
                f.write(line)
                n += 1
                if ci is not None:
                    seen.add(next(csv.reader([line]))[ci])
    return n, seen


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

    h, rd = rows(z, "routes.txt")
    routes = [r for r in rd if is_rail(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]
    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[0] in trip_ids)
        if n:
            log.warning("%d Bahnfahrten nur über frequencies.txt – werden ignoriert", n)

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

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

    for name in ("calendar.txt", "calendar_dates.txt"):
        if name in z.namelist():
            filter_by_first_col(z, name, 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, [])))
    return {"feed_version": info.get("feed_version") or zip_path.stem,
            "routes": len(routes), "trips": len(trips), "stop_times": n_st}


def ensure_osm() -> Path:
    osm = C.DATA_DIR / "osm" / "region.osm.pbf"
    if not osm.exists() or time.time() - osm.stat().st_mtime > 30 * 86400:
        osm.parent.mkdir(parents=True, exist_ok=True)
        log.info("Lade OSM-Daten %s", C.OSM_URL)
        download(C.OSM_URL, osm)
    return osm


def run_pfaedle(work: Path, rail: Path) -> Path | None:
    if not shutil.which(C.PFAEDLE):
        log.warning("pfaedle nicht gefunden – Züge fahren Luftlinie zwischen den Haltestellen")
        return None
    osm = ensure_osm()
    log.info("pfaedle rechnet Streckenverläufe …")
    t = time.time()
    res = subprocess.run([C.PFAEDLE, "-D", "-m", "rail", "-x", str(osm), str(rail)], cwd=work)
    out = work / "gtfs-out"
    if res.returncode != 0 or not (out / "shapes.txt").exists():
        log.error("pfaedle fehlgeschlagen (Code %s) – nutze Fahrplan ohne Streckenverläufe", res.returncode)
        return None
    log.info("pfaedle fertig in %.0f s", time.time() - t)
    return out


def publish(src: Path, meta: dict) -> None:
    (src / "meta.json").write_text(json.dumps(meta))
    dest = C.FEEDS_DIR / 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 C.FEEDS_DIR.iterdir() if p.is_dir())[:-3]:
        shutil.rmtree(old)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--zip", type=Path, help="lokale GTFS-Zip statt Download")
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    C.FEEDS_DIR.mkdir(parents=True, exist_ok=True)
    state_file = C.DATA_DIR / "state.json"
    state = json.loads(state_file.read_text()) if state_file.exists() else {}
    work = C.DATA_DIR / "work"
    shutil.rmtree(work, ignore_errors=True)
    work.mkdir()

    if args.zip:
        zip_path, name = args.zip, args.zip.name
    else:
        url = gtfs_url()
        name = latest_gtfs_name(url)
        if name == state.get("last_file") and not args.force:
            log.info("Kein neuer Fahrplan (%s)", name)
            return
        zip_path = work / name
        log.info("Lade %s", name)
        download(url, zip_path)

    t = time.time()
    meta = filter_rail(zip_path, work / "rail")
    log.info("Bahn-Filter: %s in %.0f s", meta, time.time() - t)
    if not args.zip:
        zip_path.unlink()

    shaped = run_pfaedle(work, work / "rail")
    meta["shapes"] = shaped is not None
    meta["source_file"] = name
    meta["prepared_at"] = int(time.time())
    publish(shaped or work / "rail", meta)

    state["last_file"] = name
    state_file.write_text(json.dumps(state))
    shutil.rmtree(work, ignore_errors=True)
    log.info("Fahrplan %s bereit", meta["feed_version"])


if __name__ == "__main__":
    main()