#!/usr/bin/env python
"""btconsole.py -- minimal Tesla/WinTesla CONSOLE emulator (the operator station).

The pod-side engine (MUNGA_L4/L4NET.CPP) boots `-net <port>` into ConsoleOnly state:
it listens on TCP <port> and waits for a console to connect and stream the mission
egg.  The real console software is absent from every archive, so this tool speaks
the console's wire protocol:

  packet := NetworkPacketHeader + NetworkManager::ReceiveEggFileMessage
    off  0  int32  clientID      = 0    (NetworkClient::NetworkManagerClientID)
    off  4  int32  gameID        = 0    (NETWORK.cpp:121 -- always 0)
    off  8  int32  fromHost      = 1    (FirstLegalHostID == the console)
    off 12  int32  timeStamp     = 0    (Time; unused by the egg handler)
    off 16  uint32 messageLength = 1024 (sizeof(ReceiveEggFileMessage))
    off 20  int32  messageID     = 3    (NetworkManager::ReceiveEggFileMessageID)
    off 24  uint32 messageFlags  = 1    (Receiver::Message::ReliableFlag)
    off 28  int32  sequenceNumber      (0..n-1; -1 is the solo-local path only)
    off 32  int32  notationFileLength  (total egg bytes)
    off 36  int32  thisMessageLength   (bytes used in this chunk, <= 1000)
    off 40  char   notationData[1000]
    total 1040 bytes per chunk
  (constants verified from the port build via `btl4.exe` BT_NET_PROBE=1)

The pod reassembles chunks sequentially (L4NET.CPP:773-816: seq 0 allocates,
append until notationFileLength reached), sets NormalState, CreateMission ->
StartConnecting (the egg's [pilots] list forms the pod mesh), and sends an
AcknowledgeEggFile back on this socket.  The console must STAY CONNECTED (a
disconnect trips the pod's console-loss path, which also closes the game
listener -- an engine bug).

Usage (legacy dial-out console -- the pod-authentic direction):
  python btconsole.py <egg-file> <host:port> [<host:port> ...]
Example (two instances on one box, -net 1501 / -net 1601):
  python btconsole.py MP.EGG 127.0.0.1:1501 127.0.0.1:1601

RELAY MODE (D1 -- internet play; pods dial OUT to this process):
  python btconsole.py --relay <consolePort> <egg-file> [--bind ADDR] [--udp-drop PCT]
Example:
  python btconsole.py --relay 1500 MP_RELAY.EGG
The relay listens on <consolePort>/tcp (console protocol: streams the egg to each
pod that connects, sends the LAUNCH pair once ALL pods have the egg), on
<consolePort>+1/tcp (the GAME relay: envelope-framed frame router, see
RelayEnvelope below), and on <consolePort>+1/udp (the unreliable channel
forwarder).  Pods run with BT_RELAY=<host>:<consolePort> and a unique
BT_SELF=<their [pilots] entry>.

  RelayEnvelope (game TCP): { int32 route; uint32 length; } + length payload bytes
    route >= 2 : client->relay: unicast to that hostID
                 relay->client: the sender's hostID (frames also self-identify
                 via NetworkPacketHeader.fromHost)
    route == -1: broadcast to every registered pod except the sender
    route == -2: HELLO   (payload: int32 magic 'BTR1', int32 myHostID)
    route == -3: PEER_UP   (payload: int32 hostID)  relay->client only
    route == -4: PEER_DOWN (payload: int32 hostID)  relay->client only
  UDP envelope: { int32 route; int32 fromHost; uint32 seq; } + frame
    route as above; -2 = HELLO/keepalive (no payload) -> relay answers -5
    (HELLO-ACK).  The relay learns each pod's public endpoint from every
    inbound datagram and forwards verbatim; unknown target endpoint falls
    back to wrapping the frame onto the target's game TCP connection.
"""
import random
import selectors
import socket
import struct
import sys
import threading
import time

CHUNK = 1000
PKT_FMT_HDR = "<iiii"          # clientID, gameID, fromHost, timeStamp
PKT_FMT_MSG = "<IiIiii"        # messageLength, messageID, messageFlags, seq, fileLen, thisLen
MESSAGE_LENGTH = 1024          # sizeof(ReceiveEggFileMessage)
EGG_MESSAGE_ID = 3             # NetworkManager::ReceiveEggFileMessageID
RELIABLE_FLAG = 1
CONSOLE_HOST_ID = 1            # FirstLegalHostID

# The LAUNCH: pods load to WaitingForLaunch and sit until the console sends
# Application::RunMissionMessage (the operator's launch button).  The handler
# ladder needs it twice: WaitingForLaunch->LaunchingMission (dispatches the
# player MissionStarting/translocation), then LaunchingMission->RunningMission.
# Sending it early (state==CreatingMission) hits the handler's Fail() -- hence
# the settle delay.  Constants from BT_NET_PROBE=1.
APPLICATION_CLIENT_ID = 4      # NetworkClient::ApplicationClientID
RUN_MISSION_MESSAGE_ID = 5     # Application::RunMissionMessageID
STOP_MISSION_MESSAGE_ID = 6    # Application::StopMissionMessageID (APP.h:383)
LAUNCH_SETTLE_SECONDS = 20.0   # egg -> first launch (pods must reach WaitingForLaunch)
LAUNCH_STEP_SECONDS = 4.0      # first launch -> second (Launching -> Running)
STOP_GRACE_SECONDS = 2.0       # past the pods' own zero so the ranking shows


def run_mission_packet():
    pkt = struct.pack(PKT_FMT_HDR, APPLICATION_CLIENT_ID, 0, CONSOLE_HOST_ID, 0)
    pkt += struct.pack("<IiI", 12, RUN_MISSION_MESSAGE_ID, RELIABLE_FLAG)
    assert len(pkt) == 16 + 12
    return pkt


def stop_mission_packet():
    # Application::StopMissionMessage {exitCode=NullExitCodeID} (id 6,
    # APPMSG.cpp:21).  The 1995 console ENDED the timed mission -- the pods
    # only display the countdown (secondsRemainingInGame, APP.cpp:652) and
    # pop the ranking window for the final 30s (DIRECTOR.cpp:113); the stop
    # itself was always console-side, which is us.
    pkt = struct.pack(PKT_FMT_HDR, APPLICATION_CLIENT_ID, 0, CONSOLE_HOST_ID, 0)
    pkt += struct.pack("<IiIi", 16, STOP_MISSION_MESSAGE_ID, RELIABLE_FLAG, 0)
    assert len(pkt) == 16 + 16
    return pkt


def parse_egg_mission_length(egg_path):
    """[mission] length= in seconds (0 = untimed)."""
    section = None
    try:
        for raw in open(egg_path, "r", encoding="latin-1"):
            line = raw.strip()
            if line.startswith("[") and line.endswith("]"):
                section = line[1:-1].lower()
            elif section == "mission" and "=" in line:
                key, _, value = line.partition("=")
                if key.strip().lower() == "length":
                    return float(value.strip())
    except (OSError, ValueError):
        pass
    return 0.0


def egg_wire_bytes(egg_text_bytes):
    """Convert raw egg TEXT to the wire form NotationFile::ReadText expects:
    NUL-separated lines (NOTATION.cpp:1043 walks `strchr(buffer,'\\0')+1` per
    line).  The real 1996 console pre-processed the egg the same way; sending
    the raw file text makes the whole egg parse as one garbage line ("ERROR:
    no map in egg!")."""
    lines = egg_text_bytes.replace(b"\r\n", b"\n").split(b"\n")
    return b"".join(line + b"\0" for line in lines)


def egg_packets(egg_bytes):
    total = len(egg_bytes)
    seq = 0
    off = 0
    while off < total:
        chunk = egg_bytes[off:off + CHUNK]
        pkt = struct.pack(PKT_FMT_HDR, 0, 0, CONSOLE_HOST_ID, 0)
        pkt += struct.pack(PKT_FMT_MSG, MESSAGE_LENGTH, EGG_MESSAGE_ID, RELIABLE_FLAG,
                           seq, total, len(chunk))
        pkt += chunk.ljust(CHUNK, b"\0")
        assert len(pkt) == 16 + MESSAGE_LENGTH
        yield pkt
        off += len(chunk)
        seq += 1


def serve_pod(hostport, egg_bytes):
    host, port = hostport.rsplit(":", 1)
    port = int(port)
    name = f"{host}:{port}"
    # The pod may not be listening yet -- retry like the pods do to each other.
    while True:
        try:
            s = socket.create_connection((host, port), timeout=5)
            break
        except OSError as e:
            print(f"[{name}] waiting for pod ({e})")
            time.sleep(1)
    print(f"[{name}] connected; streaming egg ({len(egg_bytes)} bytes)")
    n = 0
    for pkt in egg_packets(egg_bytes):
        s.sendall(pkt)
        n += 1
    print(f"[{name}] egg sent in {n} chunk(s); launch in {LAUNCH_SETTLE_SECONDS}s")
    s.settimeout(2.0)
    launch_at = time.time() + LAUNCH_SETTLE_SECONDS
    launches_sent = 0
    while True:
        if launches_sent < 2 and time.time() >= launch_at:
            s.sendall(run_mission_packet())
            launches_sent += 1
            launch_at = time.time() + LAUNCH_STEP_SECONDS
            print(f"[{name}] RunMission #{launches_sent} sent")
        try:
            data = s.recv(4096)
            if not data:
                print(f"[{name}] pod closed the console socket", flush=True)
                return
            # Expect the AcknowledgeEggFile (and possibly other console traffic).
            if len(data) >= 24:
                (clid, gid, fh, ts) = struct.unpack_from("<iiii", data, 0)
                (mlen, mid) = struct.unpack_from("<Ii", data, 16)
                print(f"[{name}] pod->console packet: clientID={clid} fromHost={fh} "
                      f"msgID={mid} len={mlen} ({len(data)} bytes)")
            else:
                print(f"[{name}] pod->console {len(data)} bytes")
        except socket.timeout:
            continue
        except OSError as e:
            # Previously UNHANDLED: a reset/abort here killed the thread silently
            # (buffered stdout lost) and the console process exited -- leaving the
            # pods holding dead console sockets.  Log it and keep the thread alive
            # (idle) so the process + diagnosis survive.
            print(f"[{name}] console socket error: {e!r} -- idling", flush=True)
            while True:
                time.sleep(60)


#############################################################################
# RELAY MODE (D1)
#############################################################################

ENV_FMT_TCP = "<iI"            # route, length
ENV_FMT_UDP = "<iiI"           # route, fromHost, seq
ENV_TCP_SIZE = 8
ENV_UDP_SIZE = 12
ROUTE_BCAST = -1
ROUTE_HELLO = -2
ROUTE_PEER_UP = -3
ROUTE_PEER_DOWN = -4
ROUTE_UDP_ACK = -5
ROUTE_SEAT_REQUEST = -6   # client->relay: assign me a free roster seat
ROUTE_SEAT_ASSIGN = -7    # relay->client: int32 hostID + NUL-terminated tag
ROUTE_SEAT_FULL = -8      # relay->client: no free seats
SEAT_RESERVE_SECONDS = 60.0
HELLO_MAGIC = 0x31525442       # 'BTR1' little-endian
FRAME_CAP = 1600               # NETWORKMANAGER_BUFFER_SIZE (NETWORK.h:89)
FIRST_GAME_HOST_ID = 2         # FirstLegalHostID(1) == the console; pods 2..N+1
STATS_PERIOD = 10.0

# LAN auto-discovery (BT_RELAY=auto): pods broadcast DISC_PROBE on this UDP
# port; the relay answers DISC_REPLY + <H consolePort>.  The pod combines the
# replier's source IP with the advertised port.  Broadcast never crosses
# routers, so this is LAN-only by construction; internet players keep using
# the explicit host:port.
DISCOVERY_PORT = 15999
DISC_PROBE = b"BTR1DISC"
DISC_REPLY = b"BTR1HERE"


def parse_egg_roster(egg_path):
    """Count the [pilots] pilot= entries -> expected hostIDs 2..N+1 in egg order.
    The relay's ONLY game knowledge; it never parses game frames."""
    entries = []
    in_pilots = False
    for raw in open(egg_path, "rb").read().replace(b"\r\n", b"\n").split(b"\n"):
        line = raw.strip()
        if line.startswith(b"["):
            in_pilots = (line.lower() == b"[pilots]")
            continue
        if in_pilots and line.lower().startswith(b"pilot="):
            entries.append(line.split(b"=", 1)[1].decode("ascii", "replace"))
    return entries


class RelayGameConn:
    """One accepted game-TCP connection: ACCEPTED -> (HELLO) REGISTERED -> GONE."""
    def __init__(self, sock, addr):
        self.sock = sock
        self.addr = addr
        self.buf = b""
        self.host_id = None            # None until HELLO registers it

    def name(self):
        hid = self.host_id if self.host_id is not None else "?"
        return f"game[{self.addr[0]}:{self.addr[1]} host={hid}]"


class RelayConsoleConn:
    """One accepted console connection (legacy raw protocol, relay-terminated)."""
    def __init__(self, sock, addr):
        self.sock = sock
        self.addr = addr
        self.egg_sent = False
        self.acked = False      # pod sent AcknowledgeEggFile -- a REAL pod.
        # Internet hardening: the console port is exposed, and random scanners
        # connect and would otherwise count as pods.  Only ACKED connections
        # count toward the launch gate; scanners never speak the protocol.


class Relay:
    def __init__(self, console_port, egg_path, bind_addr, udp_drop_pct,
                 manual_launch=False, reserved_tags=None):
        self.console_port = console_port
        self.game_port = console_port + 1
        self.bind_addr = bind_addr
        self.udp_drop_pct = udp_drop_pct
        self.manual_launch = manual_launch
        self._reserved_tags = set(reserved_tags or ())
        self.launch_requested = False            # set by the operator (stdin)
        self._launch_blocked_warned = False      # warned once about empty seats
        self.stop_requested = False              # operator 'stop' (End Mission)
        self.mission_length = parse_egg_mission_length(egg_path)
        self.mission_started_at = None           # set when RunMission #2 fires
        self.stop_sent = False
        self.eggs_done_at = None
        self.egg_bytes = egg_wire_bytes(open(egg_path, "rb").read())
        self.roster = parse_egg_roster(egg_path)
        self.expected_ids = set(range(FIRST_GAME_HOST_ID,
                                      FIRST_GAME_HOST_ID + len(self.roster)))
        # OPERATOR-RESERVED SEATS: roster seats the operator flies LOCALLY (with
        # an explicit BT_SELF).  Those instances take ~15s to boot, so a remote
        # join.bat (no BT_SELF -> SEAT_REQUEST) would otherwise be assigned the
        # operator's seat first, and the operator's later HELLO would collide
        # ("already registered" -> dropped -> game never launches).  Seat
        # ASSIGNMENT skips these host_ids; the explicit HELLO still claims them.
        self.reserved_host_ids = set()
        for i, tag in enumerate(self.roster):
            if tag in self._reserved_tags:
                self.reserved_host_ids.add(FIRST_GAME_HOST_ID + i)
        if self.reserved_host_ids:
            print(f"[relay] operator-reserved seats (no auto-assign): "
                  f"{sorted(self.reserved_host_ids)}", flush=True)
        self.sel = selectors.DefaultSelector()
        self.console_conns = []
        self.game_conns = []                     # RelayGameConn (incl. unregistered)
        self.by_host = {}                        # hostID -> RelayGameConn
        self.seat_reservations = {}              # hostID -> expiry (auto-seats)
        self.udp_endpoint = {}                   # hostID -> (addr, port)
        self.udp_sock = None
        self.launch_at = None
        self.launches_sent = 0
        self.stats = {"tcp_rx": 0, "tcp_tx": 0, "udp_rx": 0, "udp_tx": 0,
                      "udp_dropped": 0, "udp_tcp_fallback": 0}
        self.stats_at = time.time() + STATS_PERIOD

    # ---------------- lifecycle ----------------

    def run(self):
        print(f"[relay] roster: {len(self.roster)} pilot(s) -> hostIDs "
              f"{sorted(self.expected_ids)}: {self.roster}", flush=True)
        con_l = self._listener(self.console_port)
        game_l = self._listener(self.game_port)
        self.udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.udp_sock.bind((self.bind_addr, self.game_port))
        self.udp_sock.setblocking(False)
        self.sel.register(con_l, selectors.EVENT_READ, ("con_listen", None))
        self.sel.register(game_l, selectors.EVENT_READ, ("game_listen", None))
        self.sel.register(self.udp_sock, selectors.EVENT_READ, ("udp", None))
        # LAN auto-discovery responder (best-effort: another relay may own it).
        self.disc_sock = None
        try:
            self.disc_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.disc_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.disc_sock.bind((self.bind_addr, DISCOVERY_PORT))
            self.disc_sock.setblocking(False)
            self.sel.register(self.disc_sock, selectors.EVENT_READ,
                              ("disc", None))
            print(f"[relay] LAN discovery answering on udp/{DISCOVERY_PORT} "
                  f"(pods may use BT_RELAY=auto)", flush=True)
        except OSError as e:
            print(f"[relay] LAN discovery unavailable ({e!r}) -- "
                  f"pods must use the explicit address", flush=True)
            self.disc_sock = None
        print(f"[relay] console {self.bind_addr}:{self.console_port} | game tcp/udp "
              f"{self.bind_addr}:{self.game_port} | udp-drop {self.udp_drop_pct}%",
              flush=True)
        while True:
            for key, _ in self.sel.select(timeout=0.5):
                kind, obj = key.data
                if kind == "con_listen":
                    self._accept_console(key.fileobj)
                elif kind == "game_listen":
                    self._accept_game(key.fileobj)
                elif kind == "console":
                    self._console_read(obj)
                elif kind == "game":
                    self._game_read(obj)
                elif kind == "udp":
                    self._udp_read()
                elif kind == "disc":
                    self._disc_read()
            self._tick_launch()
            self._tick_stop()
            self._tick_stats()

    def _listener(self, port):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind((self.bind_addr, port))
        s.listen(16)
        s.setblocking(False)
        return s

    # ---------------- console side (legacy protocol, relay-terminated) ----------------

    def _accept_console(self, listener):
        sock, addr = listener.accept()
        sock.setblocking(False)
        conn = RelayConsoleConn(sock, addr)
        self.console_conns.append(conn)
        self.sel.register(sock, selectors.EVENT_READ, ("console", conn))
        # Stream the egg immediately (blocking sendall is fine: the pod's console
        # pad drains fast and the egg is a few KB).
        sock.setblocking(True)
        try:
            n = 0
            for pkt in egg_packets(self.egg_bytes):
                sock.sendall(pkt)
                n += 1
            conn.egg_sent = True
            print(f"[relay] console conn {addr[0]}:{addr[1]}: egg sent "
                  f"({n} chunks); awaiting pod ACK", flush=True)
        except OSError as e:
            print(f"[relay] console conn {addr[0]}:{addr[1]}: egg send failed {e!r}",
                  flush=True)
            self._drop_console(conn)
            return
        finally:
            try:
                sock.setblocking(False)
            except OSError:
                pass

    def _pods_ready(self):
        """REAL pods = console connections that ACKED the egg (scanners never
        speak the protocol, so they can't hold a roster slot)."""
        return sum(1 for c in self.console_conns if c.acked)

    def _check_launch_gate(self):
        # BACK-TO-BACK MISSIONS (2026-07-18): a previous mission finished
        # (launches_sent==2) and now all seats have RE-ACKED -- i.e. every pod
        # exited, re-ran join.bat, and reconnected fresh.  Reset the launch
        # state so the relay (and the operator's LAUNCH button) re-arm for a
        # new round -- no Stop/Start needed.  The binary needs nothing: a
        # rejoined pod is a brand-new process at WaitingForLaunch.  (Only ever
        # reached on a pod ACK, which never happens mid-mission.)
        if self.launches_sent >= 2 and self._pods_ready() >= len(self.roster):
            print("[relay] all seats rejoined after the last mission -- "
                  "re-arming for a NEW mission", flush=True)
            self.launches_sent = 0
            self.stop_sent = False
            self.stop_requested = False
            self.mission_started_at = None
            self.launch_at = None
            self.launch_requested = False
            self._launch_blocked_warned = False
            self.eggs_done_at = None
            # fall through -> the normal gate re-prints WAITING FOR OPERATOR

        # GLOBAL launch: arm the timer when the LAST pod has ACKED its egg
        # (auto mode), or wait for the operator's 'launch' command (manual
        # mode -- the operator app's Launch button writes it to our stdin).
        if self._pods_ready() >= len(self.roster) and self.launches_sent == 0 \
                and self.eggs_done_at is None:
            self.eggs_done_at = time.time()
            if self.manual_launch:
                print("[relay] all pods ACKED the egg; WAITING FOR OPERATOR "
                      "LAUNCH", flush=True)
            else:
                self.launch_at = self.eggs_done_at + LAUNCH_SETTLE_SECONDS
                print(f"[relay] all pods ACKED the egg; LAUNCH pair in "
                      f"{LAUNCH_SETTLE_SECONDS}s", flush=True)

    def _console_read(self, conn):
        try:
            data = conn.sock.recv(4096)
        except OSError:
            data = b""
        if not data:
            print(f"[relay] console conn {conn.addr[0]}:{conn.addr[1]} closed",
                  flush=True)
            self._drop_console(conn)
            return
        if len(data) >= 24:
            (clid, _gid, fh, _ts) = struct.unpack_from("<iiii", data, 0)
            (mlen, mid) = struct.unpack_from("<Ii", data, 16)
            print(f"[relay] pod->console clientID={clid} fromHost={fh} msgID={mid} "
                  f"len={mlen}", flush=True)
            # AcknowledgeEggFile (clientID=0 NetworkManager, msgID=4): this
            # connection is a REAL pod -- count it toward the launch gate.
            if clid == 0 and mid == 4 and not conn.acked:
                conn.acked = True
                print(f"[relay] pod ACK from {conn.addr[0]}:{conn.addr[1]} "
                      f"({self._pods_ready()}/{len(self.roster)} ready)",
                      flush=True)
                self._check_launch_gate()

    def _drop_console(self, conn):
        try:
            self.sel.unregister(conn.sock)
        except (KeyError, ValueError):
            pass
        try:
            conn.sock.close()
        except OSError:
            pass
        if conn in self.console_conns:
            self.console_conns.remove(conn)

    def _log_launch_readiness(self):
        # PRE-LAUNCH READINESS (2026-07-18): make the "launching short" footgun
        # VISIBLE.  A roster seat with no registered pod means the mission will
        # STALL -- every pod waits on that empty seat's connection and the
        # All-connections-completed gate never fires.  Log filled vs empty
        # seats so the operator sees it BEFORE the confusion.
        filled, empty = [], []
        for i, tag in enumerate(self.roster):
            host_id = FIRST_GAME_HOST_ID + i
            (filled if host_id in self.by_host else empty).append(
                f"seat{i + 1}({tag})")
        print(f"[relay] LAUNCH readiness: {len(filled)}/{len(self.roster)} "
              f"seats filled -- {', '.join(filled) or 'NONE'}", flush=True)
        if empty:
            print(f"[relay] *** WARNING: {len(empty)} EMPTY seat(s): "
                  f"{', '.join(empty)} -- the mission will STALL (every pod "
                  f"waits on the empty seat).  Stop, reduce the roster to the "
                  f"players present, and re-launch.", flush=True)

    def _tick_launch(self):
        # Manual mode: the operator's 'launch' arms the timer, still honouring
        # the settle window (pods must reach WaitingForLaunch or the handler
        # Fail()s -- the same reason the auto timer waits).
        if (self.manual_launch and self.launch_requested
                and self.launch_at is None and self.launches_sent == 0):
            if self.eggs_done_at is not None:
                self.launch_at = max(time.time(),
                                     self.eggs_done_at + LAUNCH_SETTLE_SECONDS)
                wait = max(0.0, self.launch_at - time.time())
                print(f"[relay] operator LAUNCH received; firing in "
                      f"{wait:.0f}s", flush=True)
                self._log_launch_readiness()
            elif not self._launch_blocked_warned:
                # Operator pressed LAUNCH but the gate is NOT ready (not every
                # roster seat has ACKed its egg) -- previously this did NOTHING
                # visible ("I pressed launch and nothing happened").  Say WHY.
                self._launch_blocked_warned = True
                print("[relay] LAUNCH pressed but NOT all seats are filled:",
                      flush=True)
                self._log_launch_readiness()
        if self.launch_at is None or self.launches_sent >= 2:
            return
        if time.time() < self.launch_at:
            return
        pkt = run_mission_packet()
        for conn in list(self.console_conns):
            try:
                conn.sock.setblocking(True)
                conn.sock.sendall(pkt)
                conn.sock.setblocking(False)
            except OSError as e:
                print(f"[relay] launch send failed to {conn.addr}: {e!r}", flush=True)
        self.launches_sent += 1
        self.launch_at = time.time() + LAUNCH_STEP_SECONDS
        print(f"[relay] RunMission #{self.launches_sent} sent to "
              f"{len(self.console_conns)} pod(s)", flush=True)
        if self.launches_sent == 2:
            self.mission_started_at = time.time()
            # CONSUME the launch: clear the timer + the request so the NEXT
            # LAUNCH press is a fresh, deliberate new-mission request (else the
            # back-to-back block would re-fire while this mission is running).
            self.launch_at = None
            self.launch_requested = False
            if self.mission_length > 0:
                print(f"[relay] mission clock: {self.mission_length:.0f}s "
                      "(StopMission at expiry)", flush=True)

    def _tick_stop(self):
        # THE MISSION CLOCK (console-side, as in 1995): send StopMission when
        # the egg's [mission] length expires, or on the operator's 'stop'
        # command (the End Mission button).
        if self.stop_sent or self.launches_sent < 2:
            return
        expired = (self.mission_length > 0
                   and self.mission_started_at is not None
                   and time.time() >= self.mission_started_at
                       + self.mission_length + STOP_GRACE_SECONDS)
        if not (expired or self.stop_requested):
            return
        pkt = stop_mission_packet()
        for conn in list(self.console_conns):
            try:
                conn.sock.setblocking(True)
                conn.sock.sendall(pkt)
                conn.sock.setblocking(False)
            except OSError as e:
                print(f"[relay] stop send failed to {conn.addr}: {e!r}",
                      flush=True)
        self.stop_sent = True
        why = "operator END MISSION" if self.stop_requested else "time expired"
        print(f"[relay] StopMission sent to {len(self.console_conns)} pod(s) "
              f"({why})", flush=True)

    # ---------------- game TCP side ----------------

    def _accept_game(self, listener):
        sock, addr = listener.accept()
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        sock.setblocking(False)
        conn = RelayGameConn(sock, addr)
        self.game_conns.append(conn)
        self.sel.register(sock, selectors.EVENT_READ, ("game", conn))
        print(f"[relay] {conn.name()} accepted (awaiting HELLO)", flush=True)

    def _game_read(self, conn):
        try:
            data = conn.sock.recv(8192)
        except OSError:
            data = b""
        if not data:
            self._drop_game(conn, "closed")
            return
        conn.buf += data
        while len(conn.buf) >= ENV_TCP_SIZE:
            route, length = struct.unpack_from(ENV_FMT_TCP, conn.buf, 0)
            if length > FRAME_CAP:
                self._drop_game(conn, f"oversize frame ({length})")
                return
            if len(conn.buf) < ENV_TCP_SIZE + length:
                break                                   # partial; wait for more
            payload = conn.buf[ENV_TCP_SIZE:ENV_TCP_SIZE + length]
            conn.buf = conn.buf[ENV_TCP_SIZE + length:]
            self.stats["tcp_rx"] += 1
            self._handle_game_frame(conn, route, payload)
            if conn.sock.fileno() < 0:                  # dropped mid-loop
                return

    def _handle_game_frame(self, conn, route, payload):
        if conn.host_id is None and route == ROUTE_SEAT_REQUEST:
            # SERVER-ASSIGNED SEATS: hand out the lowest roster seat that is
            # neither claimed (by_host) nor reserved; reserve it briefly so
            # two simultaneous joiners can't race onto the same seat.  The
            # requesting connection closes after the reply; the pod re-dials
            # and claims the seat with a normal HELLO.
            now = time.time()
            self.seat_reservations = {h: t for h, t in
                                      self.seat_reservations.items() if t > now}
            assigned = None
            for i, tag in enumerate(self.roster):
                host_id = FIRST_GAME_HOST_ID + i
                if (host_id in self.by_host or host_id in self.seat_reservations
                        or host_id in self.reserved_host_ids):
                    continue                    # claimed / racing / operator-local
                assigned = (host_id, tag)
                break
            if assigned is None:
                self._send_raw(conn, struct.pack(ENV_FMT_TCP,
                                                 ROUTE_SEAT_FULL, 0))
                print(f"[relay] {conn.name()} seat request: ROSTER FULL",
                      flush=True)
                return
            host_id, tag = assigned
            self.seat_reservations[host_id] = now + SEAT_RESERVE_SECONDS
            payload_out = struct.pack("<i", host_id) + tag.encode() + b"\0"
            self._send_raw(conn, struct.pack(ENV_FMT_TCP, ROUTE_SEAT_ASSIGN,
                                             len(payload_out)) + payload_out)
            print(f"[relay] {conn.name()} seat request -> assigned host "
                  f"{host_id} '{tag}' (reserved {SEAT_RESERVE_SECONDS:.0f}s)",
                  flush=True)
            return
        if conn.host_id is None:
            # First frame MUST be a valid HELLO.
            if route != ROUTE_HELLO or len(payload) < 8:
                self._drop_game(conn, f"first frame not HELLO (route={route})")
                return
            magic, host_id = struct.unpack_from("<iI", payload, 0)
            host_id = int(host_id)
            if magic != HELLO_MAGIC:
                self._drop_game(conn, f"bad HELLO magic 0x{magic:x}")
                return
            if host_id not in self.expected_ids:
                self._drop_game(conn, f"HELLO hostID {host_id} not in roster "
                                      f"{sorted(self.expected_ids)}")
                return
            if host_id in self.by_host:
                self._drop_game(conn, f"HELLO hostID {host_id} already registered")
                return
            conn.host_id = host_id
            self.by_host[host_id] = conn
            self.seat_reservations.pop(host_id, None)   # claim clears reserve
            print(f"[relay] {conn.name()} REGISTERED "
                  f"({len(self.by_host)}/{len(self.roster)})", flush=True)
            # PEER_UP exchange: newcomer learns everyone; everyone learns newcomer.
            for other_id in sorted(self.by_host):
                if other_id != host_id:
                    self._send_control(conn, ROUTE_PEER_UP, other_id)
            for other in self.by_host.values():
                if other is not conn:
                    self._send_control(other, ROUTE_PEER_UP, host_id)
            return
        if route == ROUTE_BCAST:
            env = struct.pack(ENV_FMT_TCP, conn.host_id, len(payload))
            for other in list(self.by_host.values()):
                if other is not conn:
                    self._send_raw(other, env + payload)
        elif route >= FIRST_GAME_HOST_ID:
            target = self.by_host.get(route)
            if target is None:
                return                                  # peer gone; drop silently
            env = struct.pack(ENV_FMT_TCP, conn.host_id, len(payload))
            self._send_raw(target, env + payload)
        else:
            print(f"[relay] {conn.name()} bad route {route} -- ignored", flush=True)

    def _send_control(self, conn, route, host_id):
        self._send_raw(conn, struct.pack(ENV_FMT_TCP, route, 4)
                       + struct.pack("<i", host_id))

    def _send_raw(self, conn, data):
        try:
            conn.sock.setblocking(True)
            conn.sock.sendall(data)
            conn.sock.setblocking(False)
            self.stats["tcp_tx"] += 1
        except OSError as e:
            self._drop_game(conn, f"send failed {e!r}")

    def _drop_game(self, conn, why):
        print(f"[relay] {conn.name()} dropped: {why}", flush=True)
        try:
            self.sel.unregister(conn.sock)
        except (KeyError, ValueError):
            pass
        try:
            conn.sock.close()
        except OSError:
            pass
        if conn in self.game_conns:
            self.game_conns.remove(conn)
        if conn.host_id is not None and self.by_host.get(conn.host_id) is conn:
            del self.by_host[conn.host_id]
            self.udp_endpoint.pop(conn.host_id, None)
            for other in self.by_host.values():
                self._send_control(other, ROUTE_PEER_DOWN, conn.host_id)

    # ---------------- game UDP side ----------------

    def _udp_read(self):
        while True:
            try:
                data, endpoint = self.udp_sock.recvfrom(4096)
            except (BlockingIOError, InterruptedError):
                return
            except OSError:
                return
            if len(data) < ENV_UDP_SIZE:
                continue
            route, from_host, _seq = struct.unpack_from(ENV_FMT_UDP, data, 0)
            if from_host not in self.by_host:
                continue                                # not TCP-registered: drop
            self.stats["udp_rx"] += 1
            # NAT-rebind tolerant: every datagram refreshes the endpoint map.
            self.udp_endpoint[from_host] = endpoint
            if route == ROUTE_HELLO:
                ack = struct.pack(ENV_FMT_UDP, ROUTE_UDP_ACK, CONSOLE_HOST_ID, 0)
                try:
                    self.udp_sock.sendto(ack, endpoint)
                except OSError:
                    pass
                continue
            if route == ROUTE_BCAST:
                targets = [h for h in self.by_host if h != from_host]
            elif route >= FIRST_GAME_HOST_ID and route in self.by_host:
                targets = [route]
            else:
                continue
            for host_id in targets:
                if self.udp_drop_pct and random.random() * 100.0 < self.udp_drop_pct:
                    self.stats["udp_dropped"] += 1
                    continue
                ep = self.udp_endpoint.get(host_id)
                if ep is not None:
                    try:
                        self.udp_sock.sendto(data, ep)      # forward VERBATIM
                        self.stats["udp_tx"] += 1
                        continue
                    except OSError:
                        pass
                # Target's UDP endpoint unknown/broken: wrap the frame onto its
                # game TCP connection (asymmetric-UDP-blockage fallback).
                frame = data[ENV_UDP_SIZE:]
                target = self.by_host.get(host_id)
                if target is not None and frame:
                    env = struct.pack(ENV_FMT_TCP, from_host, len(frame))
                    self._send_raw(target, env + frame)
                    self.stats["udp_tcp_fallback"] += 1

    # ---------------- LAN auto-discovery ----------------

    def _disc_read(self):
        while True:
            try:
                data, endpoint = self.disc_sock.recvfrom(64)
            except (BlockingIOError, InterruptedError, OSError):
                return
            if data == DISC_PROBE:
                try:
                    self.disc_sock.sendto(
                        DISC_REPLY + struct.pack("<H", self.console_port),
                        endpoint)
                    print(f"[relay] discovery probe from "
                          f"{endpoint[0]}:{endpoint[1]} -> answered "
                          f"(console port {self.console_port})", flush=True)
                except OSError:
                    pass

    # ---------------- stats ----------------

    def _tick_stats(self):
        if time.time() < self.stats_at:
            return
        self.stats_at = time.time() + STATS_PERIOD
        s = self.stats
        print(f"[relay-stats] tcp rx/tx {s['tcp_rx']}/{s['tcp_tx']} | "
              f"udp rx/tx {s['udp_rx']}/{s['udp_tx']} "
              f"(dropped {s['udp_dropped']}, tcp-fallback {s['udp_tcp_fallback']}) | "
              f"registered {sorted(self.by_host)} udp-known {sorted(self.udp_endpoint)}",
              flush=True)


def relay_main(argv):
    """argv: --relay <consolePort> <egg-file> [--bind ADDR] [--udp-drop PCT]
    [--manual-launch]"""
    args = argv[:]
    bind_addr = "0.0.0.0"
    udp_drop = 0.0
    manual = False
    if "--bind" in args:
        i = args.index("--bind")
        bind_addr = args[i + 1]
        del args[i:i + 2]
    if "--udp-drop" in args:
        i = args.index("--udp-drop")
        udp_drop = float(args[i + 1])
        del args[i:i + 2]
    if "--manual-launch" in args:
        manual = True
        args.remove("--manual-launch")
    reserved_tags = ()
    if "--reserve" in args:
        i = args.index("--reserve")
        reserved_tags = tuple(t for t in args[i + 1].split(",") if t)
        del args[i:i + 2]
    if len(args) != 2:
        print(__doc__)
        return 2
    console_port = int(args[0])
    egg_path = args[1]
    relay = Relay(console_port, egg_path, bind_addr, udp_drop, manual,
                  reserved_tags=reserved_tags)
    if not relay.roster:
        print(f"[relay] ERROR: no [pilots] entries in {egg_path}")
        return 2
    # Operator command channel: a daemon thread reads stdin.  'launch' fires
    # the mission (manual mode only); 'stop' ends it early (End Mission).
    def stdin_reader():
        for line in sys.stdin:
            command = line.strip().lower()
            if command == "launch":
                relay.launch_requested = True
            elif command == "stop":
                relay.stop_requested = True
    threading.Thread(target=stdin_reader, daemon=True).start()
    try:
        relay.run()
    except KeyboardInterrupt:
        print("[relay] shutting down")
    return 0


def main():
    if len(sys.argv) >= 2 and sys.argv[1] == "--relay":
        return relay_main(sys.argv[2:])
    if len(sys.argv) < 3:
        print(__doc__)
        return 2
    egg_bytes = egg_wire_bytes(open(sys.argv[1], "rb").read())
    pods = sys.argv[2:]
    threads = [threading.Thread(target=serve_pod, args=(p, egg_bytes), daemon=True)
               for p in pods]
    for t in threads:
        t.start()
    try:
        while any(t.is_alive() for t in threads):
            time.sleep(0.5)
    except KeyboardInterrupt:
        print("console shutting down")
    return 0


if __name__ == "__main__":
    sys.exit(main())