#!/usr/bin/env python3
"""Split/Second lobby server: sessions, leaderboards, and the UDP helper for NAT traversal.

One file, Python 3.8 or later, standard library only. Run it with

    python3 splitsecond_lobby.py [--host 0.0.0.0] [--http-port 31860] [--udp-port 31861] [--db lobby.sqlite]
                                 [--stats-dir DIR] [--verbose]

HTTP (JSON, plain, curl-able) - what a copy of the game with ONLINE set to INTERNET talks to:

    GET    /v1/whoami                      the caller's public address
    POST   /v1/sessions                    a host creates a session            -> {session_id, ticket}
    GET    /v1/sessions?title=<id>         search (open public slots first)    -> {sessions: [...]}
    GET    /v1/sessions/<id>               details
    POST   /v1/sessions/<id>/join          a player joins                      -> {session, ticket}
    POST   /v1/sessions/<id>/leave         a player leaves
    POST   /v1/sessions/<id>/heartbeat     the host is alive (every 20 s; a session unheard for 90 s goes)
    POST   /v1/sessions/<id>/migrate       another player becomes the host
    POST   /v1/sessions/<id>/modify        slots, flags, properties
    DELETE /v1/sessions/<id>               the host deletes the session
    GET    /v1/sessions/<id>/peers         every member's ticket and last seen UDP address
    POST   /v1/stats                       write a leaderboard row
    GET    /v1/leaderboard/<view>?title=   read a leaderboard
    GET    /stats                          usage statistics, counts only (below)

UDP (one port, for the copies' own game sockets): the reflector and the relay - see handle_udp.

The session and wire design follows what the Xenia emulator's netplay fork (BSD-3-Clause) and
Xenia-WebServices (MIT) do for the console's XSession calls; the code here is written for this
title and shares none of theirs.
"""
import argparse
import asyncio
import hashlib
import json
import os
import re
import secrets
import signal
import sqlite3
import struct
import sys
import time
from urllib.parse import parse_qs, urlsplit

SESSION_TTL = 90.0      # seconds without a heartbeat
ENDPOINT_TTL = 120.0    # seconds without a UDP hello
MAX_BODY = 64 * 1024
MAX_SESSIONS = 5000
MAX_CONNS_PER_IP = 64   # open HTTP connections from one address (loopback is exempt)
WRITE_TIMEOUT = 10.0    # a client that does not read its answer for this long is dropped
VERSION = "1.1"
START_TIME = time.time()

sessions = {}           # session id -> dict
tickets = {}            # ticket (hex, 16 chars) -> {"session": id, "xuid": int, "endpoint": (ip, port) | None, "seen": t}
udp_transport = None
db = None
args = None


def now():
    return time.time()


def log(*parts):
    print(time.strftime("%H:%M:%S"), *parts, flush=True)


def verbose():
    return bool(args and args.verbose)


def who(*parts):
    """Gamertags and addresses go into the log only with --verbose."""
    return " ".join(str(p) for p in parts) if verbose() else "(hidden)"


# --------------------------------------------------------------------------- sessions

def public_view(s, with_endpoints=False):
    members = [{"xuid": m["xuid"], "gamertag": m.get("gamertag", ""), "private": m.get("private", False)}
               for m in s["members"]]
    out = {
        "session_id": s["id"],
        "title_id": s["title_id"],
        "xnkid": s["xnkid"],
        "xnaddr": s["xnaddr"],          # the host's XNADDR as the game made it (hex, 36 bytes)
        "key": s["key"],                # the session key (hex, 16 bytes)
        "flags": s["flags"],
        "public_slots": s["public_slots"],
        "private_slots": s["private_slots"],
        "open_public": max(0, s["public_slots"] - sum(1 for m in s["members"] if not m.get("private"))),
        "open_private": max(0, s["private_slots"] - sum(1 for m in s["members"] if m.get("private"))),
        "host": {
            "xuid": s["host_xuid"],
            "gamertag": s["host_gamertag"],
            "public_ip": s["host_public_ip"],
            "lan_ip": s["host_lan_ip"],
            "port": s["host_port"],
            "ticket": s["host_ticket"],
        },
        "state": s["state"],
        "nonce": s["nonce"],
        "properties": s["properties"],
        "members": members,
        "created": int(s["created"]),
    }
    if with_endpoints:
        out["peers"] = peers_of(s)
    return out


def peers_of(s):
    out = []
    for m in s["members"]:
        t = tickets.get(m.get("ticket", ""))
        ep = t["endpoint"] if t and t["endpoint"] and now() - t["seen"] < ENDPOINT_TTL else None
        out.append({"xuid": m["xuid"], "gamertag": m.get("gamertag", ""), "ticket": m.get("ticket", ""),
                    "udp": {"ip": ep[0], "port": ep[1]} if ep else None})
    return out


def bind_ticket(requested, session_id, xuid):
    """A copy of the game makes its own ticket (8 random bytes, hex) and registers it by UDP; the
    session calls bind it to a session. A ticket that is taken by another player is refused."""
    t = str(requested or "").lower()
    if len(t) != 16 or any(c not in "0123456789abcdef" for c in t):
        t = secrets.token_hex(8)
    entry = tickets.get(t)
    if entry and entry["xuid"] not in (0, xuid):
        t = secrets.token_hex(8)
        entry = None
    if entry is None:
        entry = tickets[t] = {"session": None, "xuid": xuid, "endpoint": None, "seen": now()}
    entry["session"] = session_id
    entry["xuid"] = xuid
    return t


def drop_session(sid):
    s = sessions.pop(sid, None)
    if not s:
        return
    st.session_ended(s)
    for t in [t for t, v in tickets.items() if v["session"] == sid]:
        tickets[t]["session"] = None


def expire():
    t = now()
    for sid in [sid for sid, s in sessions.items() if t - s["seen"] > SESSION_TTL]:
        log("session", sid, "expired")
        drop_session(sid)
    for k in [k for k, v in tickets.items() if v["session"] is None and t - v["seen"] > 600]:
        del tickets[k]


def need(body, *keys):
    for k in keys:
        if k not in body:
            raise ApiError(400, "missing field %s" % k)


class ApiError(Exception):
    def __init__(self, status, message):
        super().__init__(message)
        self.status = status
        self.message = message


def hexfield(body, key, nbytes):
    v = str(body.get(key, ""))
    if len(v) != nbytes * 2:
        raise ApiError(400, "%s must be %d hex chars" % (key, nbytes * 2))
    try:
        bytes.fromhex(v)
    except ValueError:
        raise ApiError(400, "%s is not hex" % key)
    return v.lower()


def create_session(body, peer_ip):
    need(body, "title_id", "host_xuid")
    expire()
    if len(sessions) >= MAX_SESSIONS:
        st.refused()
        raise ApiError(503, "too many sessions")
    sid = secrets.token_hex(4)
    xuid = int(body["host_xuid"])
    ticket = bind_ticket(body.get("ticket"), sid, xuid)
    s = {
        "id": sid,
        "title_id": int(body["title_id"]),
        "xnkid": hexfield(body, "xnkid", 8),
        "xnaddr": hexfield(body, "xnaddr", 36),
        "key": hexfield(body, "key", 16),
        "nonce": str(body.get("nonce", "")),
        "flags": int(body.get("flags", 0)),
        "public_slots": int(body.get("public_slots", 0)),
        "private_slots": int(body.get("private_slots", 0)),
        "host_xuid": xuid,
        "host_gamertag": str(body.get("host_gamertag", ""))[:32],
        "host_public_ip": peer_ip,
        "host_lan_ip": str(body.get("lan_ip", "")),
        "host_port": int(body.get("port", 0)),
        "host_ticket": ticket,
        "state": int(body.get("state", 0)),
        "properties": body.get("properties", []),
        "members": [{"xuid": xuid, "gamertag": str(body.get("host_gamertag", ""))[:32], "ticket": ticket}],
        "created": now(),
        "seen": now(),
        "mode": mode_of(body.get("properties", [])),
        "peak_members": 1,
        "relay_first": None,    # first and last relayed packet of the session (the relay's share, below)
        "relay_last": None,
    }
    sessions[sid] = s
    log("session", sid, "created by", who(s["host_gamertag"], peer_ip), "title %08X" % s["title_id"],
        "mode", s["mode"])
    if verbose():
        log("session", sid, "properties", json.dumps(s["properties"])[:400])
    st.entered(s["host_gamertag"] or str(xuid))
    st.created_game(s["mode"])
    st.sample()
    return {"session_id": sid, "ticket": ticket, "session": public_view(s)}


def search_sessions(query):
    expire()
    title = int(query.get("title", ["0"])[0] or 0)
    limit = min(100, int(query.get("limit", ["50"])[0]))
    found = []
    for s in sessions.values():
        if title and s["title_id"] != title:
            continue
        v = public_view(s)
        found.append(v)
    found.sort(key=lambda v: (v["open_public"] == 0, -v["created"]))
    return {"sessions": found[:limit]}


def get_session(sid):
    s = sessions.get(sid)
    if not s:
        raise ApiError(404, "no such session")
    return s


def join_session(sid, body, peer_ip):
    s = get_session(sid)
    need(body, "xuid")
    xuid = int(body["xuid"])
    private = bool(body.get("private", False))
    for m in s["members"]:
        if m["xuid"] == xuid:
            return {"session": public_view(s, True), "ticket": m["ticket"]}
    v = public_view(s)
    if (private and v["open_private"] == 0) or (not private and v["open_public"] == 0):
        raise ApiError(409, "session full")
    ticket = bind_ticket(body.get("ticket"), sid, xuid)
    s["members"].append({"xuid": xuid, "gamertag": str(body.get("gamertag", ""))[:32], "private": private,
                         "ticket": ticket, "public_ip": peer_ip})
    log("session", sid, "joined by", who(body.get("gamertag", ""), peer_ip))
    s["peak_members"] = max(s["peak_members"], len(s["members"]))
    st.entered(str(body.get("gamertag", "")) or str(xuid))
    st.sample()
    return {"session": public_view(s, True), "ticket": ticket}


def leave_session(sid, body):
    s = get_session(sid)
    need(body, "xuid")
    xuid = int(body["xuid"])
    left = [m for m in s["members"] if m["xuid"] == xuid]
    s["members"] = [m for m in s["members"] if m["xuid"] != xuid]
    for m in left:
        if m.get("ticket") in tickets:
            tickets[m["ticket"]]["session"] = None
    if xuid == s["host_xuid"] and s["members"]:
        # the host left without migrating: the oldest member is offered the session
        s["state"] = s["state"]
    if not s["members"]:
        drop_session(sid)
    return {"ok": True}


def migrate_session(sid, body, peer_ip):
    s = get_session(sid)
    need(body, "xuid")
    xuid = int(body["xuid"])
    member = next((m for m in s["members"] if m["xuid"] == xuid), None)
    if not member:
        raise ApiError(403, "not a member")
    s["host_xuid"] = xuid
    s["host_gamertag"] = member.get("gamertag", "")
    s["host_ticket"] = member["ticket"]
    s["host_public_ip"] = peer_ip
    s["host_lan_ip"] = str(body.get("lan_ip", s["host_lan_ip"]))
    s["host_port"] = int(body.get("port", s["host_port"]))
    if body.get("xnaddr"):
        s["xnaddr"] = hexfield(body, "xnaddr", 36)
    if body.get("xnkid"):
        s["xnkid"] = hexfield(body, "xnkid", 8)
    if body.get("key"):
        s["key"] = hexfield(body, "key", 16)
    s["seen"] = now()
    log("session", sid, "migrated to", who(s["host_gamertag"]))
    return {"session": public_view(s, True)}


def modify_session(sid, body):
    s = get_session(sid)
    for k in ("flags", "public_slots", "private_slots"):
        if k in body:
            s[k] = int(body[k])
    if "state" in body:
        new = int(body["state"])
        if new == STATE_INGAME and s["state"] != STATE_INGAME:
            st.started_game(s)
        s["state"] = new
    if "properties" in body:
        s["properties"] = body["properties"]
        s["mode"] = mode_of(s["properties"])
    s["seen"] = now()
    return {"session": public_view(s)}


# --------------------------------------------------------------------------- leaderboards

def stats_init(path):
    global db
    db = sqlite3.connect(path, check_same_thread=False)
    db.execute("""CREATE TABLE IF NOT EXISTS stats (
        title INTEGER, view INTEGER, xuid INTEGER, gamertag TEXT, rating INTEGER, columns TEXT, updated REAL,
        PRIMARY KEY (title, view, xuid))""")
    db.commit()


def stats_write(body):
    need(body, "title_id", "view", "xuid")
    rating = int(body.get("rating", 0))
    db.execute("""INSERT INTO stats VALUES (?,?,?,?,?,?,?)
                  ON CONFLICT(title, view, xuid) DO UPDATE SET gamertag=excluded.gamertag,
                  rating=excluded.rating, columns=excluded.columns, updated=excluded.updated""",
               (int(body["title_id"]), int(body["view"]), int(body["xuid"]),
                str(body.get("gamertag", ""))[:32], rating, json.dumps(body.get("columns", [])), now()))
    db.commit()
    st.race_result()
    return {"ok": True}


def stats_read(view, query):
    title = int(query.get("title", ["0"])[0])
    xuids = [int(x) for x in query.get("xuid", [])]
    start = int(query.get("start", ["1"])[0])
    count = min(100, int(query.get("count", ["10"])[0]))
    rows = db.execute("SELECT xuid, gamertag, rating, columns FROM stats WHERE title=? AND view=? "
                      "ORDER BY rating DESC, updated ASC", (title, int(view))).fetchall()
    ranked = [{"rank": i + 1, "xuid": r[0], "gamertag": r[1], "rating": r[2], "columns": json.loads(r[3])}
              for i, r in enumerate(rows)]
    if xuids:
        picked = [r for r in ranked if r["xuid"] in xuids]
    else:
        picked = ranked[start - 1:start - 1 + count]
    return {"view": int(view), "total": len(ranked), "rows": picked}


# --------------------------------------------------------------------------- usage statistics
#
# Counts only: no address and no gamertag is kept or written. "Unique players" is a set of hashes of the
# lower-cased gamertag under a random salt made at the start of each UTC day; the set lives in memory
# only, is never written, and is dropped at midnight. With --stats-dir the day's totals go to
# daily.jsonl (one JSON line per day, appended when the day ends) and current.json (today so far,
# rewritten every minute). A restart on the same day carries on from current.json (a player seen before
# and after the restart then counts twice).
#
# "Sign-in" here is a player entering a game (a host creating a session, a player joining one).
# "Race results" are leaderboard rows written by the title (one per player per race).
# The relay's share: a session of two or more players counts as "relayed" if its relay traffic went on for
# more than RELAY_STAY seconds (the title relays first and then hole-punches a direct path, so a short
# burst at the start is normal); otherwise it counts as "direct".

STATE_INGAME = 2        # XSESSION_STATE: 0 lobby, 1 registration, 2 in game, 3 reporting, 4 deleted
MODE_CONTEXT = 0x0000800B   # the console's game-mode context (the title sets it before it creates a game)
MODE_LABELS = {}        # "mode_<value>" -> a name, once the title's values are known; else "Mode <value>"
MAX_START_PLAYERS = 8
RELAY_STAY = 30.0
RELAY_LIVE = 10.0       # relay traffic this recently = "relaying now"
HISTORY_DAYS = 30
COUNTERS = ("connections", "refused_full", "signins", "games_created", "games_started", "race_results",
            "relay_packets", "relay_bytes", "sessions_relayed", "sessions_direct")
DAY_FIELDS = ("date", "peak_players", "peak_sessions", "signins", "unique_players", "games_started",
              "race_results")


def mode_of(properties):
    for p in properties or []:
        if isinstance(p, dict) and p.get("kind") == "context" and p.get("id") == MODE_CONTEXT:
            try:
                return "mode_%d" % int(p.get("value", 0))
            except (TypeError, ValueError):
                break
    return "other"


def mode_label(key):
    return MODE_LABELS.get(key) or ("Other" if key == "other" else key.replace("mode_", "Mode "))


def mode_key(label):
    for k in list(MODE_LABELS) + ["other"]:
        if mode_label(k) == label:
            return k
    return "other" if label == "Other" else label.replace("Mode ", "mode_")


def date_of(day):
    return time.strftime("%Y-%m-%d", time.gmtime(day * 86400))


class Stats:
    def __init__(self, directory=""):
        self.dir = directory
        self.written = 0.0
        self.cache = (0.0, None)
        self.history_cache = (None, [])
        self.reset(int(now() // 86400))
        if directory:
            os.makedirs(directory, exist_ok=True)
            self.restore()

    def reset(self, day):
        self.day = day
        self.salt = secrets.token_bytes(16)
        self.seen = set()
        self.unique_base = 0
        self.n = dict.fromkeys(COUNTERS, 0)
        self.peak_sessions = self.peak_players = 0
        self.created = {}
        self.started = {}
        self.start_players = {}

    # ---- counting
    def refused(self):
        self.n["refused_full"] += 1

    def connection(self):
        self.n["connections"] += 1

    def entered(self, gamertag):
        self.n["signins"] += 1
        if len(self.seen) < (1 << 20):
            digest = hashlib.blake2b(str(gamertag).lower().encode("utf-8", "replace"), key=self.salt,
                                     digest_size=8).digest()
            self.seen.add(digest)

    def created_game(self, mode):
        self.n["games_created"] += 1
        self.created[mode] = self.created.get(mode, 0) + 1

    def started_game(self, s):
        self.n["games_started"] += 1
        self.started[s["mode"]] = self.started.get(s["mode"], 0) + 1
        size = max(1, min(MAX_START_PLAYERS, len(s["members"])))
        self.start_players[size] = self.start_players.get(size, 0) + 1

    def race_result(self):
        self.n["race_results"] += 1

    def relayed(self, session, size):
        self.n["relay_packets"] += 1
        self.n["relay_bytes"] += size
        if session is not None:
            t = now()
            if session["relay_first"] is None:
                session["relay_first"] = t
            session["relay_last"] = t

    def session_ended(self, s):
        if s["peak_members"] < 2:
            return
        if s["relay_first"] is not None and s["relay_last"] - s["relay_first"] > RELAY_STAY:
            self.n["sessions_relayed"] += 1
        else:
            self.n["sessions_direct"] += 1

    def sample(self):
        """Peaks, from the live numbers."""
        live = self.live()
        self.peak_sessions = max(self.peak_sessions, live["sessions"])
        self.peak_players = max(self.peak_players, live["players"])

    # ---- the live view
    def live(self):
        t = now()
        in_games = set()
        relaying = 0
        racing = 0
        for s in sessions.values():
            for m in s["members"]:
                in_games.add(m["xuid"])
            if s["relay_last"] is not None and t - s["relay_last"] < RELAY_LIVE:
                relaying += 1
            if s["state"] == STATE_INGAME:
                racing += 1
        clients = sum(1 for v in tickets.values() if v["endpoint"] and t - v["seen"] < 30)
        return {"sessions": len(sessions), "players": max(len(in_games), clients),
                "players_in_games": len(in_games), "games_in_progress": racing, "sessions_relaying": relaying}

    def open_games(self):
        out = []
        for s in sessions.values():
            slots = s["public_slots"] + s["private_slots"]
            label = hashlib.blake2b(s["id"].encode(), key=self.salt, digest_size=2).hexdigest().upper()
            out.append({"name": "Game " + label, "mode": mode_label(s["mode"]), "players": len(s["members"]),
                        "max": max(slots, len(s["members"])),
                        "state": "racing" if s["state"] == STATE_INGAME else "waiting", "created": int(s["created"])})
        out.sort(key=lambda g: (g["state"] != "waiting", -g["created"]))
        return out[:50]

    # ---- the day's totals
    def today(self):
        d = {"date": date_of(self.day)}
        d["unique_players"] = self.unique_base + len(self.seen)
        d["peak_sessions"] = self.peak_sessions
        d["peak_players"] = self.peak_players
        d.update(self.n)
        d["created_by_mode"] = {mode_label(k): v for k, v in sorted(self.created.items())}
        d["started_by_mode"] = {mode_label(k): v for k, v in sorted(self.started.items())}
        d["start_players"] = {str(k): v for k, v in sorted(self.start_players.items())}
        return d

    def roll(self):
        """At the first call of a new UTC day: write the old day's line and start afresh."""
        day = int(now() // 86400)
        if day != self.day:
            self.append_daily(self.today())
            self.reset(day)
            self.write_current()

    # ---- files
    def path(self, name):
        return os.path.join(self.dir, name)

    def append_daily(self, today):
        if not self.dir:
            return
        p = self.path("daily.jsonl")
        try:
            last = ""
            if os.path.exists(p):
                with open(p, "rb") as f:
                    f.seek(max(0, os.path.getsize(p) - 8192))
                    lines = f.read().decode("utf-8", "replace").strip().splitlines()
                    last = lines[-1] if lines else ""
            if '"date":"%s"' % today["date"] in last:
                return              # that day is already there
            with open(p, "a") as f:
                f.write(json.dumps(today, separators=(",", ":")) + "\n")
        except OSError as e:
            log("stats: cannot write daily.jsonl:", e)

    def write_current(self):
        if not self.dir:
            return
        body = self.today()
        body.update(self.live())
        body["updated"] = int(now())
        # the unique count carries over a restart; the hashed set itself is never written
        tmp = self.path("current.json.tmp")
        try:
            with open(tmp, "w") as f:
                f.write(json.dumps(body, separators=(",", ":")))
            os.replace(tmp, self.path("current.json"))
            self.written = now()
        except OSError as e:
            log("stats: cannot write current.json:", e)

    def restore(self):
        try:
            with open(self.path("current.json")) as f:
                old = json.load(f)
        except (OSError, ValueError):
            return
        if not isinstance(old, dict) or "date" not in old:
            return
        if old["date"] != date_of(self.day):
            # the server was down at midnight: the old day's line has not been written yet
            live_keys = ("updated", "sessions", "players", "players_in_games", "games_in_progress",
                         "sessions_relaying")
            self.append_daily({k: v for k, v in old.items() if k not in live_keys})
            return
        try:
            for k in COUNTERS:
                self.n[k] = int(old.get(k, 0))
            self.unique_base = int(old.get("unique_players", 0))
            self.peak_sessions = int(old.get("peak_sessions", 0))
            self.peak_players = int(old.get("peak_players", 0))
            for name, field in (("created_by_mode", self.created), ("started_by_mode", self.started)):
                for label, v in (old.get(name) or {}).items():
                    field[mode_key(label)] = int(v)
            self.start_players = {int(k): int(v) for k, v in (old.get("start_players") or {}).items()}
        except (TypeError, ValueError, AttributeError):
            pass
        log("stats: carried on from current.json (%s)" % old["date"])

    def history(self):
        """The last HISTORY_DAYS finished days, from daily.jsonl (cached by file time)."""
        if not self.dir:
            return []
        p = self.path("daily.jsonl")
        try:
            stamp = os.stat(p).st_mtime
        except OSError:
            return []
        if self.history_cache[0] != stamp:
            rows = []
            try:
                with open(p) as f:
                    lines = f.read().splitlines()[-HISTORY_DAYS:]
                for line in lines:
                    try:
                        r = json.loads(line)
                        rows.append({k: r.get(k, 0) for k in DAY_FIELDS})
                    except (ValueError, AttributeError):
                        pass
            except OSError:
                pass
            self.history_cache = (stamp, rows)
        return self.history_cache[1]

    # ---- GET /stats
    def view(self):
        t = now()
        if self.cache[1] is not None and t - self.cache[0] < 2.0:
            return self.cache[1]
        expire()
        self.roll()
        self.sample()
        today = self.today()
        days = self.history()[-(HISTORY_DAYS - 1):] + [{k: today[k] for k in DAY_FIELDS}]
        out = {"server": "splitsecond-lobby", "version": VERSION, "uptime_seconds": int(t - START_TIME),
               "time": int(t), "now": self.live(), "today": today, "history": days, "open_games": self.open_games()}
        self.cache = (t, out)
        return out


st = Stats()


async def stats_loop():
    """Every 5 s: expire sessions, sample the peaks, roll the day; every minute: write current.json."""
    while True:
        await asyncio.sleep(5)
        try:
            expire()
            st.roll()
            st.sample()
            if st.dir and now() - st.written >= 60:
                st.write_current()
        except Exception as e:  # statistics must never stop the server
            log("stats error:", e)


# --------------------------------------------------------------------------- HTTP

def route(method, path, query, body, peer_ip):
    parts = [p for p in path.split("/") if p]
    if parts == ["stats"] and method == "GET":
        return st.view()
    if parts[:1] != ["v1"]:
        raise ApiError(404, "not found")
    parts = parts[1:]
    if parts == ["whoami"] and method == "GET":
        return {"ip": peer_ip}
    if parts == ["health"] and method == "GET":
        return {"ok": True, "sessions": len(sessions)}
    if parts == ["sessions"]:
        if method == "POST":
            return create_session(body, peer_ip)
        if method == "GET":
            return search_sessions(query)
    if len(parts) >= 2 and parts[0] == "sessions":
        sid = parts[1]
        if len(parts) == 2:
            if method == "GET":
                return {"session": public_view(get_session(sid), True)}
            if method == "DELETE":
                get_session(sid)
                drop_session(sid)
                return {"ok": True}
        elif len(parts) == 3 and method == "POST":
            action = parts[2]
            if action == "join":
                return join_session(sid, body, peer_ip)
            if action == "leave":
                return leave_session(sid, body)
            if action == "heartbeat":
                get_session(sid)["seen"] = now()
                return {"ok": True}
            if action == "migrate":
                return migrate_session(sid, body, peer_ip)
            if action == "modify":
                return modify_session(sid, body)
        elif len(parts) == 3 and parts[2] == "peers" and method == "GET":
            return {"peers": peers_of(get_session(sid))}
    if parts == ["stats"] and method == "POST":
        return stats_write(body)
    if len(parts) == 2 and parts[0] == "leaderboard" and method == "GET":
        return stats_read(parts[1], query)
    raise ApiError(404, "not found")


conns_by_ip = {}        # address -> open HTTP connections


def is_loopback(ip):
    return ip.startswith("127.") or ip in ("::1", "")


async def handle_http(reader, writer):
    """One connection: at most MAX_CONNS_PER_IP per address (loopback exempt), then serve it."""
    peer = writer.get_extra_info("peername")
    ip = peer[0] if peer else ""
    loopback = is_loopback(ip)
    if not loopback:
        if conns_by_ip.get(ip, 0) >= MAX_CONNS_PER_IP:
            st.refused()
            writer.transport.abort()
            return
        conns_by_ip[ip] = conns_by_ip.get(ip, 0) + 1
        st.connection()
    try:
        await serve_http(reader, writer, ip)
    finally:
        if not loopback:
            conns_by_ip[ip] -= 1
            if conns_by_ip[ip] <= 0:
                del conns_by_ip[ip]


async def serve_http(reader, writer, peer_ip):
    status, payload = 500, {"error": "internal"}
    try:
        head = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), 10)
        lines = head.decode("latin-1").split("\r\n")
        method, target, _ = lines[0].split(" ", 2)
        headers = {}
        for line in lines[1:]:
            if ":" in line:
                k, v = line.split(":", 1)
                headers[k.strip().lower()] = v.strip()
        # Behind a reverse proxy, the client address is the proxy's: trust X-Forwarded-For only if asked.
        if args.trust_forwarded and "x-forwarded-for" in headers:
            peer_ip = headers["x-forwarded-for"].split(",")[0].strip()
        length = int(headers.get("content-length", "0"))
        if length > MAX_BODY:
            raise ApiError(413, "body too large")
        body = {}
        if length:
            raw = await asyncio.wait_for(reader.readexactly(length), 10)
            try:
                body = json.loads(raw.decode("utf-8"))
            except ValueError:
                raise ApiError(400, "body is not JSON")
            if not isinstance(body, dict):
                raise ApiError(400, "body must be an object")
        url = urlsplit(target)
        try:
            payload = route(method, url.path, parse_qs(url.query), body, peer_ip)
            status = 200
        except ApiError as e:
            status, payload = e.status, {"error": e.message}
        except (ValueError, TypeError, KeyError) as e:
            status, payload = 400, {"error": "bad request: %s" % e}
    except ApiError as e:
        status, payload = e.status, {"error": e.message}
    except (asyncio.TimeoutError, asyncio.IncompleteReadError, ConnectionError, asyncio.LimitOverrunError):
        writer.close()
        return
    data = json.dumps(payload).encode()
    reason = {200: "OK", 400: "Bad Request", 403: "Forbidden", 404: "Not Found", 409: "Conflict",
              413: "Payload Too Large", 500: "Internal Server Error", 503: "Service Unavailable"}.get(status, "")
    writer.write(("HTTP/1.1 %d %s\r\nContent-Type: application/json\r\nContent-Length: %d\r\n"
                  "Cache-Control: no-store\r\nConnection: close\r\n\r\n" % (status, reason, len(data))).encode() + data)
    try:
        # a client that does not read its answer is dropped
        await asyncio.wait_for(writer.drain(), WRITE_TIMEOUT)
    except (ConnectionError, asyncio.TimeoutError):
        writer.transport.abort()
        return
    writer.close()


# --------------------------------------------------------------------------- UDP: reflector and relay
#
#   "SSR?" + nonce(8)                    -> "SSR!" + nonce(8) + ip(4) + port(2)       where the caller is seen from
#   "SSA"  + ticket(8)                   -> "SSA!" + ticket(8) + ip(4) + port(2)      registers the UDP address
#                                                                                      of a ticket (repeat every 15 s)
#   "SSQ"  + ticket(8)                   -> "SSQ!" + ticket(8) + ip(4) + port(2)      where a ticket was last heard from
#                                                                                      (zeros if unknown)
#   "SSD"  + dest_ticket(8) + payload    -> dest gets "SSD" + source_ticket(8) + payload   the relay: both tickets registered,
#                                                                                      at least one of them in a session
# A ticket is 8 random bytes (16 hex digits) a copy of the game makes for itself; it registers it
# with "SSA" and names it in the HTTP create / join calls, which bind it to a session. Peers learn
# each other's tickets from the session (members' "ticket", /peers).

class UdpHandler(asyncio.DatagramProtocol):
    def connection_made(self, transport):
        global udp_transport
        udp_transport = transport

    def datagram_received(self, data, addr):
        try:
            handle_udp(data, addr)
        except Exception as e:  # a bad packet must not stop the server
            log("udp error", who(addr), e)


def endpoint_bytes(addr):
    return bytes(int(x) for x in addr[0].split(".")) + struct.pack(">H", addr[1])


def handle_udp(data, addr):
    if len(data) >= 12 and data[:4] == b"SSR?":
        udp_transport.sendto(b"SSR!" + data[4:12] + endpoint_bytes(addr), addr)
    elif len(data) >= 11 and data[:3] == b"SSA":
        ticket = data[3:11].hex()
        t = tickets.get(ticket)
        if t is None:
            if len(tickets) > 50000:
                return
            t = tickets[ticket] = {"session": None, "xuid": 0, "endpoint": None, "seen": now()}
        t["endpoint"] = (addr[0], addr[1])
        t["seen"] = now()
        udp_transport.sendto(b"SSA!" + data[3:11] + endpoint_bytes(addr), addr)
    elif len(data) >= 11 and data[:3] == b"SSQ":
        t = tickets.get(data[3:11].hex())
        ep = t["endpoint"] if t and t["endpoint"] and now() - t["seen"] < ENDPOINT_TTL else None
        udp_transport.sendto(b"SSQ!" + data[3:11] + (endpoint_bytes(ep) if ep else bytes(6)), addr)
    elif len(data) > 11 and data[:3] == b"SSD":
        dest = tickets.get(data[3:11].hex())
        if not dest or not dest["endpoint"]:
            return
        # the source is whichever ticket this address registered
        src_ticket = next((k for k, v in tickets.items() if v["endpoint"] == (addr[0], addr[1])), None)
        if not src_ticket:
            return
        # A game's players talk before they are in the server's session (the game connects, then
        # joins the session): one of the two has to be in a session, so the relay is not an open
        # forwarder between strangers.
        if not dest["session"] and not tickets[src_ticket]["session"]:
            return
        udp_transport.sendto(b"SSD" + bytes.fromhex(src_ticket) + data[11:], dest["endpoint"])
        st.relayed(sessions.get(tickets[src_ticket]["session"]) or sessions.get(dest["session"]), len(data) - 11)


# --------------------------------------------------------------------------- main

async def main():
    global args
    p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--http-port", type=int, default=31860)
    p.add_argument("--udp-port", type=int, default=31861)
    p.add_argument("--db", default="lobby.sqlite", help="leaderboards (SQLite file; ':memory:' for none kept)")
    p.add_argument("--trust-forwarded", action="store_true",
                   help="take the client address from X-Forwarded-For (behind a reverse proxy only)")
    p.add_argument("--stats-dir", default="",
                   help="usage statistics (counts only): daily.jsonl and current.json go here; empty = not kept")
    p.add_argument("--verbose", action="store_true", help="log gamertags and addresses (off by default)")
    args = p.parse_args()
    global st
    st = Stats(args.stats_dir)
    stats_init(args.db)
    loop = asyncio.get_running_loop()
    await loop.create_datagram_endpoint(UdpHandler, local_addr=(args.host, args.udp_port))
    server = await asyncio.start_server(handle_http, args.host, args.http_port)
    asyncio.get_running_loop().create_task(stats_loop())
    log("lobby: HTTP %s:%d, UDP %s:%d, leaderboards in %s, usage statistics in %s" % (
        args.host, args.http_port, args.host, args.udp_port, args.db, args.stats_dir or "(not kept)"))
    # SIGTERM (systemd stop or restart): write the day's counts first, so a restart loses nothing
    stop = asyncio.Event()
    try:
        loop.add_signal_handler(signal.SIGTERM, stop.set)
    except (NotImplementedError, ValueError):
        pass
    async with server:
        await stop.wait()
    st.write_current()
    log("lobby: stopped")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass
