Trainpulse/trainpulse/server.py

"""Webserver: liefert die Karte aus und berechnet laufend die Zugpositionen.

    uvicorn trainpulse.server:app --host 0.0.0.0 --port 8080
"""
import datetime as dt
import gzip
import json
import logging
import threading
import time
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.staticfiles import StaticFiles

from . import config as C
from . import border, realtime
from .feed import TZ, Feed

log = logging.getLogger("trainpulse")
STEP, STEPS = 2, 5  # Positionen alle 2 s für die nächsten 10 s, damit der Browser Kurven nachfährt
WEB = Path(__file__).resolve().parent.parent / "web"


class Source:
    """Ein Fahrplan (Schweiz, Deutschland, …) mit seinem Echtzeit-Feed."""

    def __init__(self, name: str):
        self.name, self.cfg = name, C.SOURCES[name]
        self.feed: Feed | None = None
        self.rt: dict = {}
        self.rt_version: str | None = None
        self.rt_at: float | None = None
        self.rt_error: str | None = None
        self.rt_tried_at: float | None = None
        self.on_disk = False  # ein aufbereiteter Fahrplan liegt bereit (wird gleich geladen)

    @property
    def has_rt(self) -> bool:
        return bool(self.cfg["rt_url"]) and (bool(self.cfg["rt_key"]) or not self.cfg["needs_key"])


class State:
    sources = {name: Source(name) for name in C.ENABLED}
    vehicles: bytes = b'{"type":"FeatureCollection","features":[]}'
    vehicles_gz: bytes = gzip.compress(vehicles)
    count = 0
    computed_at: float | None = None


S = State()


def loop(every: float, fn, *args):
    def run():
        while True:
            t = time.time()
            try:
                fn(*args)
            except Exception:
                log.exception("Fehler in %s %s", fn.__name__, args)
            time.sleep(max(1.0, every - (time.time() - t)))
    threading.Thread(target=run, daemon=True, name=fn.__name__).start()


def watch_feeds():
    """Lädt je Quelle den passenden Fahrplan: den, auf den sich die Echtzeitdaten beziehen (Schweiz:
    die schalten erst um 15 Uhr auf einen neuen Fahrplan um), sonst den neusten."""
    today = dt.datetime.now(TZ).date()
    found = {}
    for src in S.sources.values():  # erst alle markieren, damit die Karte nicht mit einem Land allein startet
        found[src.name] = sorted(m.parent.name for m in (C.FEEDS_DIR / src.name).glob("*/meta.json"))
        src.on_disk = bool(found[src.name])
    for src in S.sources.values():
        ready = found[src.name]
        if not ready:
            continue
        want = src.rt_version if src.rt_version in ready else ready[-1]
        meta = json.loads((C.FEEDS_DIR / src.name / want / "meta.json").read_text())
        if (src.feed and src.feed.version == want and src.feed.covers(today + dt.timedelta(days=1))
                and src.feed.meta.get("prepared_at") == meta.get("prepared_at")):  # gleiche Version neu aufbereitet?
            continue
        log.info("[%s] Lade Fahrplan %s …", src.name, want)
        t = time.time()
        src.feed = Feed(C.FEEDS_DIR / src.name / want)
        log.info("[%s] Fahrplan %s geladen in %.0f s (%d Fahrten)", src.name, want, time.time() - t,
                 len(src.feed.trip_ids))
        try:  # Echtzeit sofort holen statt auf den nächsten Abruf zu warten
            poll_realtime(src)
        except Exception:
            log.exception("[%s] Erster Echtzeit-Abruf fehlgeschlagen", src.name)
        compute_positions()


def poll_realtime(src: Source):
    if not src.has_rt or not src.feed:
        return
    try:
        src.rt_version, src.rt = realtime.parse(realtime.fetch(src.cfg["rt_url"], src.cfg["rt_key"]),
                                                src.feed.trip_index)
        src.rt_at, src.rt_error = time.time(), None
    except Exception as e:  # Daten älter als 10 min sind wertlos -> verwerfen statt falsch anzeigen
        src.rt_error = str(e)
        if src.rt_at and time.time() - src.rt_at > 600:
            src.rt = {}
        raise
    finally:
        src.rt_tried_at = time.time()


def positions(src: Source, t: float) -> list:
    feats = []
    for f in src.feed.vehicles(t, src.rt):
        if len(S.sources) > 1 and not border.shows(src.name, S.sources, *f["geometry"]["coordinates"]):
            continue  # gehört zum Fahrplan des Nachbarlandes
        f["properties"]["id"] = f"{src.name}/{f['properties']['id']}"
        feats.append(f)
    return feats


def compute_positions():
    now = time.time()
    feats, future = [], [{} for _ in range(STEPS)]
    for src in list(S.sources.values()):
        if not src.feed:
            continue
        feats += positions(src, now)
        for k in range(STEPS):
            future[k].update({f"{src.name}/{f['properties']['id']}": f["geometry"]["coordinates"] + [f["properties"]["heading"]]
                              for f in src.feed.vehicles(now + STEP * (k + 1), src.rt)})
    for f in feats:
        path = []
        for pos in future:
            p = pos.get(f["properties"]["id"])
            if p is None:
                break  # Fahrt endet
            path.append(p)
        f["path"] = path
    S.vehicles = json.dumps({"type": "FeatureCollection", "ts": now, "step": STEP, "features": feats},
                            separators=(",", ":")).encode()
    S.vehicles_gz = gzip.compress(S.vehicles, 5)
    S.count, S.computed_at = len(feats), time.time()


@asynccontextmanager
async def lifespan(app):
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    C.FEEDS_DIR.mkdir(parents=True, exist_ok=True)
    loop(60, watch_feeds)
    for src in S.sources.values():
        loop(src.cfg["rt_interval"], poll_realtime, src)
    loop(3, compute_positions)
    yield


app = FastAPI(lifespan=lifespan, docs_url=None, redoc_url=None)


@app.middleware("http")
async def no_stale_frontend(request: Request, call_next):
    # Nach einem Update sonst alte app.js aus dem Browser-Cache; mit no-cache fragt der Browser
    # jedes Mal kurz nach (304, wenn unverändert). Kacheln dürfen normal gecacht werden.
    resp = await call_next(request)
    if not request.url.path.startswith(("/tiles/", "/api/")):
        resp.headers["Cache-Control"] = "no-cache"
    return resp


@app.get("/api/vehicles")
def vehicles(request: Request):
    headers = {"Cache-Control": "no-store", "X-Now": str(time.time())}  # für den Uhrenabgleich im Browser
    if "gzip" in request.headers.get("accept-encoding", ""):
        return Response(S.vehicles_gz, media_type="application/json", headers={**headers, "Content-Encoding": "gzip"})
    return Response(S.vehicles, media_type="application/json", headers=headers)


@app.get("/api/trip/{source}/{date}/{ti}")
def trip(source: str, date: str, ti: int):
    src = S.sources.get(source)
    f = src.feed if src else None
    if not f or not (0 <= ti < len(f.trip_ids)) or not (len(date) == 8 and date.isdigit()):
        raise HTTPException(404)
    return f.trip_detail(date, ti, time.time(), src.rt)


@app.get("/api/status")
def status():
    loaded = [s for s in S.sources.values() if s.feed]
    pending = [s for s in S.sources.values() if s.on_disk and not s.feed]  # nicht halbe Karte zeigen
    rt = [s for s in loaded if s.has_rt]
    ok = [s for s in rt if not s.rt_error and s.rt_at and time.time() - s.rt_at < 180]
    return {
        "loaded": bool(loaded),
        # erst "bereit", wenn Positionen mit (dem ersten Versuch) Echtzeit berechnet sind
        "ready": bool(loaded and not pending and S.computed_at and all(s.rt_tried_at and S.computed_at >= s.rt_tried_at for s in rt)),
        "vehicles": S.count,
        "rt_ok": bool(ok) and len(ok) == len(loaded),
        "rt_some": bool(ok),
        "sources": {s.name: {
            "feed_version": s.feed.version if s.feed else None,
            "shapes": s.feed.meta.get("shapes") if s.feed else None,
            "realtime": s.has_rt,
            "rt_version": s.rt_version,
            "rt_age": round(time.time() - s.rt_at) if s.rt_at else None,
            "rt_error": s.rt_error,
            "rt_trips": len(s.rt),
        } for s in S.sources.values()},
    }


app.mount("/tiles", StaticFiles(directory=C.DATA_DIR / "tiles", check_dir=False), name="tiles")  # PMTiles, Range-Requests
app.mount("/", StaticFiles(directory=WEB, html=True), name="web")