relay.py

"""WebSocket match rooms — fan-out relay, no game simulation."""

from __future__ import annotations

import asyncio
import json
from typing import Any

from fastapi import WebSocket, WebSocketDisconnect

from lobbies import board


class MatchRoom:
    def __init__(self, lobby_id: str, mode: str, match_id: str) -> None:
        self.lobby_id = lobby_id
        self.mode = mode
        self.match_id = match_id
        self.clients: dict[str, WebSocket] = {}  # seated players
        self.spectators: dict[str, WebSocket] = {}  # read-only viewers
        self.lock = asyncio.Lock()
        self.started = False

    async def add(self, username: str, ws: WebSocket, *, spectate: bool = False) -> None:
        async with self.lock:
            bucket = self.spectators if spectate else self.clients
            old = bucket.get(username)
            if old and old is not ws:
                try:
                    await old.close()
                except Exception:
                    pass
            bucket[username] = ws

    async def remove(self, username: str, ws: WebSocket, *, spectate: bool = False) -> None:
        async with self.lock:
            bucket = self.spectators if spectate else self.clients
            if bucket.get(username) is ws:
                del bucket[username]

    async def broadcast(self, msg: dict[str, Any], exclude: str | None = None) -> None:
        data = json.dumps(msg)
        dead_players: list[str] = []
        dead_specs: list[str] = []
        async with self.lock:
            players = list(self.clients.items())
            specs = list(self.spectators.items())
        for user, client in players + specs:
            if exclude and user == exclude:
                continue
            try:
                await client.send_text(data)
            except Exception:
                if user in self.clients:
                    dead_players.append(user)
                else:
                    dead_specs.append(user)
        for u in dead_players:
            board.leave(u)
            async with self.lock:
                self.clients.pop(u, None)
        async with self.lock:
            for u in dead_specs:
                self.spectators.pop(u, None)


rooms: dict[str, MatchRoom] = {}
rooms_lock = asyncio.Lock()


async def get_or_create_room(lobby_id: str, mode: str, match_id: str) -> MatchRoom:
    key = f"{mode}:{lobby_id}"
    async with rooms_lock:
        room = rooms.get(key)
        if room is None or room.match_id != match_id:
            room = MatchRoom(lobby_id, mode, match_id)
            rooms[key] = room
        return room


async def handle_match_ws(
    ws: WebSocket,
    mode: str,
    lobby_id: str,
    username: str,
    role: str = "play",
) -> None:
    await ws.accept()
    username = (username or "").strip()
    spectate = (role or "play").lower() in ("spectate", "spec", "watch")
    if not username:
        await ws.send_text(json.dumps({"type": "error", "error": "username required"}))
        await ws.close()
        return

    # Spectators use a distinct key so they never collide with a seated player name
    if spectate:
        username = f"spec:{username}"

    lobby = board.get(lobby_id)
    if not lobby or lobby["mode"] != mode:
        await ws.send_text(json.dumps({"type": "error", "error": "lobby not found"}))
        await ws.close()
        return

    if not spectate:
        seat_info = board.by_user.get(username)
        if not seat_info or seat_info[0] != lobby_id:
            await ws.send_text(json.dumps({"type": "error", "error": "not seated in lobby"}))
            await ws.close()
            return
        board.heartbeat(username)

    # Wait until lobby is starting/live with a match_id
    for _ in range(180):
        lobby = board.get(lobby_id)
        if lobby and not spectate:
            # Keep seat alive while parked in the wait loop (HTTP HB can lag).
            board.heartbeat(username)
            board._maybe_start(lobby)
            lobby = board.get(lobby_id)
            # Seat was wiped (stale / leave) — stop waiting instead of crashing later.
            if not board.by_user.get(username) or board.by_user.get(username)[0] != lobby_id:
                await ws.send_text(json.dumps({"type": "error", "error": "seat lost — rejoin lobby"}))
                await ws.close()
                return
        if lobby and lobby.get("match_id") and lobby["status"] in ("starting", "live"):
            # Only proceed if this player is still seated (or spectating).
            if spectate or (board.by_user.get(username) and board.by_user[username][0] == lobby_id):
                break
        if spectate and lobby and lobby["status"] == "open" and not lobby.get("match_id"):
            await ws.send_text(
                json.dumps({"type": "waiting", "lobby": board._public(lobby), "spectating": True})
            )
            await asyncio.sleep(0.5)
            continue
        await asyncio.sleep(0.5)
        try:
            pub = board._public(lobby) if lobby else None
            await ws.send_text(
                json.dumps({"type": "waiting", "lobby": pub, "spectating": spectate})
            )
        except Exception:
            return
    else:
        await ws.send_text(json.dumps({"type": "error", "error": "match did not start"}))
        await ws.close()
        return

    lobby = board.get(lobby_id)
    if not lobby or not lobby.get("match_id"):
        await ws.send_text(json.dumps({"type": "error", "error": "lobby reset — rejoin"}))
        await ws.close()
        return

    seat = None
    if not spectate:
        seat_info = board.by_user.get(username)
        if not seat_info or seat_info[0] != lobby_id:
            await ws.send_text(json.dumps({"type": "error", "error": "not seated in lobby"}))
            await ws.close()
            return
        seat = seat_info[1]

    match_id = lobby["match_id"]
    room = await get_or_create_room(lobby_id, mode, match_id)
    await room.add(username, ws, spectate=spectate)

    players = board.players_in(lobby_id)
    start_payload = {
        "type": "match_start",
        "matchId": match_id,
        "mode": mode,
        "lobbyId": lobby_id,
        "username": username,
        "seat": seat,
        "players": players,
        "spectating": spectate,
    }
    await ws.send_text(json.dumps(start_payload))
    if not spectate:
        board.mark_live(lobby_id)
        await room.broadcast(
            {
                "type": "player_joined",
                "username": username,
                "seat": seat,
                "players": board.players_in(lobby_id),
            },
            exclude=username,
        )

    try:
        while True:
            raw = await ws.receive_text()
            if not spectate:
                board.heartbeat(username)
            try:
                msg = json.loads(raw)
            except json.JSONDecodeError:
                continue
            if not isinstance(msg, dict):
                continue
            # Spectators are receive-only (except ping)
            if spectate:
                if msg.get("type") == "ping":
                    await ws.send_text(json.dumps({"type": "pong"}))
                continue
            msg["from"] = username
            if msg.get("type") == "ping":
                await ws.send_text(json.dumps({"type": "pong"}))
                continue
            await room.broadcast(msg, exclude=username)
    except WebSocketDisconnect:
        pass
    except Exception:
        pass
    finally:
        await room.remove(username, ws, spectate=spectate)
        if not spectate:
            # HF WS often blips right at match_start. Instant leave was wiping BOTH
            # seats and bootstrapping players back to an empty lobby.
            await asyncio.sleep(3.0)
            still = board.by_user.get(username)
            reconnected = username in room.clients
            if still and still[0] == lobby_id and not reconnected:
                board.leave(username)
                await room.broadcast(
                    {
                        "type": "player_left",
                        "username": username,
                        "players": board.players_in(lobby_id),
                    }
                )
        async with rooms_lock:
            if room.lobby_id == lobby_id and not room.clients and not room.spectators:
                rooms.pop(f"{mode}:{lobby_id}", None)