Wizard/wizard_server/app/lobby.py

"""Lobbies: a group of seats (humans and AI players) that play one game."""

from __future__ import annotations

import asyncio
import logging
import secrets
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from enum import Enum
from typing import Any

from wizard_game import (
    MAX_PLAYERS,
    MIN_PLAYERS,
    Event,
    GameFinished,
    GameOptions,
    IllegalActionError,
    TrickCompleted,
    WizardGame,
    action_from_dict,
)

from .ai_bridge import AIBridge
from .config import Settings
from .protocol import ERROR_BUSY
from .sessions import Session

logger = logging.getLogger(__name__)

LOBBY_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"  # no 0/O, 1/I
LOBBY_CODE_LENGTH = 6

AI_NAMES = (
    "Merlin", "Morgana", "Gandalf", "Radagast", "Rincewind", "Circe",
    "Prospero", "Elminster", "Saruman", "Medea", "Zatanna", "Raistlin",
)


class LobbyError(Exception):
    """An error that is reported to the client that caused it."""

    def __init__(self, message: str, code: str | None = None) -> None:
        super().__init__(message)
        # Optional machine readable code, see ``protocol.ERROR_*``.
        self.code = code


class LobbyStatus(str, Enum):
    WAITING = "waiting"
    PLAYING = "playing"
    FINISHED = "finished"


@dataclass
class Seat:
    player_id: str
    name: str
    is_ai: bool = False
    session: Session | None = None

    @property
    def connected(self) -> bool:
        return self.is_ai or (self.session is not None and self.session.connected)


class Lobby:
    def __init__(
        self,
        code: str,
        bridge: AIBridge,
        settings: Settings,
        on_status_change: Callable[[], None] | None = None,
    ) -> None:
        self.code = code
        # Called when a game starts or ends (the number of running games changes).
        self._on_status_change = on_status_change
        self.seats: list[Seat] = []
        self.host_id: str | None = None
        self.options = GameOptions()
        self.status = LobbyStatus.WAITING
        self.game: WizardGame | None = None
        self._bridge = bridge
        self._settings = settings
        self._ai_task: asyncio.Task | None = None
        self._pause_until = 0.0
        self._closed = False

    # --- Queries ------------------------------------------------------------

    @property
    def has_humans(self) -> bool:
        return any(not seat.is_ai for seat in self.seats)

    def seat_index(self, player_id: str) -> int | None:
        for index, seat in enumerate(self.seats):
            if seat.player_id == player_id:
                return index
        return None

    def view(self) -> dict[str, Any]:
        return {
            "code": self.code,
            "status": self.status.value,
            "host_id": self.host_id,
            "options": self.options.to_dict(),
            "players": [
                {
                    "id": seat.player_id,
                    "name": seat.name,
                    "is_ai": seat.is_ai,
                    "connected": seat.connected,
                }
                for seat in self.seats
            ],
            "min_players": MIN_PLAYERS,
            "max_players": MAX_PLAYERS,
            "ai_available": self._bridge.available,
        }

    def state_for(self, session: Session, events: Iterable[Event] = ()) -> dict[str, Any]:
        game_view = None
        if self.game is not None:
            you = self.seat_index(session.player_id)
            observation = self.game.observation(you if you is not None else 0)
            data = observation.to_dict()
            if you is None:  # not seated (should not happen): hide private info
                data["hand"], data["legal_actions"] = [], []
            game_view = {"you": you, "observation": data}
        return {
            "type": "state",
            "lobby": self.view(),
            "game": game_view,
            "events": [event.to_dict() for event in events],
        }

    # --- Membership ---------------------------------------------------------

    def add_human(self, session: Session) -> None:
        index = self.seat_index(session.player_id)
        if index is not None:  # already seated, e.g. after a reconnect
            self.seats[index].session = session
            return
        if self.status is not LobbyStatus.WAITING:
            raise LobbyError("The game in this lobby has already started")
        if len(self.seats) >= MAX_PLAYERS:
            raise LobbyError("This lobby is full")
        assert session.name is not None
        self.seats.append(Seat(session.player_id, self._unique_name(session.name), session=session))
        if self.host_id is None:
            self.host_id = session.player_id

    def add_ai(self, by: Session) -> None:
        self._require_host(by)
        self._require_waiting()
        if len(self.seats) >= MAX_PLAYERS:
            raise LobbyError("This lobby is full")
        taken = {seat.name for seat in self.seats}
        name = next((n for n in AI_NAMES if n not in taken), f"Bot {len(self.seats) + 1}")
        self.seats.append(Seat(f"ai-{secrets.token_hex(4)}", name, is_ai=True))

    def remove_player(self, by: Session, player_id: str) -> Seat:
        self._require_host(by)
        self._require_waiting()
        if player_id == by.player_id:
            raise LobbyError("Use 'leave' to leave the lobby")
        index = self.seat_index(player_id)
        if index is None:
            raise LobbyError("No such player")
        return self.seats.pop(index)

    def leave(self, session: Session) -> None:
        index = self.seat_index(session.player_id)
        if index is None:
            return
        seat = self.seats[index]
        if self.status is LobbyStatus.WAITING:
            self.seats.pop(index)
        else:
            # Keep the seat so the game can go on: an AI player takes over.
            seat.is_ai = True
            seat.session = None
            seat.player_id = f"ai-{secrets.token_hex(4)}"
            self._schedule_ai()
        if self.host_id == session.player_id:
            self.host_id = next((s.player_id for s in self.seats if not s.is_ai), None)

    # --- Host actions -------------------------------------------------------

    def set_options(self, by: Session, options: GameOptions) -> None:
        self._require_host(by)
        self._require_waiting()
        self.options = options

    def start(self, by: Session) -> None:
        self._require_host(by)
        self._require_waiting()
        if not MIN_PLAYERS <= len(self.seats) <= MAX_PLAYERS:
            raise LobbyError(f"Wizard needs {MIN_PLAYERS} to {MAX_PLAYERS} players")
        self.game = WizardGame(len(self.seats), self.options)
        self.status = LobbyStatus.PLAYING
        self._pause_until = 0.0
        logger.info("Lobby %s: game started with %d players", self.code, len(self.seats))
        self._notify_status_change()
        self._schedule_ai()

    def return_to_lobby(self, by: Session) -> None:
        self._require_host(by)
        if self.status is not LobbyStatus.FINISHED:
            raise LobbyError("The game is not finished yet")
        self.status = LobbyStatus.WAITING
        self.game = None

    # --- Game actions -------------------------------------------------------

    def apply_action(self, session: Session, action_data: dict[str, Any]) -> None:
        if self.status is not LobbyStatus.PLAYING or self.game is None:
            raise LobbyError("No game is running")
        index = self.seat_index(session.player_id)
        if index is None:
            raise LobbyError("You are not playing in this game")
        try:
            action = action_from_dict(action_data)
            events = self.game.apply(index, action)
        except (ValueError, IllegalActionError) as exc:
            raise LobbyError(str(exc)) from exc
        self._after_action(events)

    # --- Broadcasting -------------------------------------------------------

    def broadcast(self, events: Iterable[Event] = ()) -> None:
        events = list(events)
        for seat in self.seats:
            if seat.session is not None:
                seat.session.send(self.state_for(seat.session, events))

    def close(self) -> None:
        self._closed = True
        if self._ai_task and not self._ai_task.done():
            self._ai_task.cancel()
        if self.status is LobbyStatus.PLAYING:
            self.status = LobbyStatus.FINISHED
            self._notify_status_change()

    # --- Internals ----------------------------------------------------------

    def _require_host(self, session: Session) -> None:
        if session.player_id != self.host_id:
            raise LobbyError("Only the host can do that")

    def _require_waiting(self) -> None:
        if self.status is not LobbyStatus.WAITING:
            raise LobbyError("The game has already started")

    def _unique_name(self, name: str) -> str:
        taken = {seat.name for seat in self.seats}
        candidate, n = name, 2
        while candidate in taken:
            candidate = f"{name} ({n})"
            n += 1
        return candidate

    def _after_action(self, events: list[Event]) -> None:
        loop = asyncio.get_running_loop()
        if any(isinstance(e, TrickCompleted) for e in events):
            self._pause_until = loop.time() + self._settings.trick_pause
        finished = any(isinstance(e, GameFinished) for e in events)
        if finished:
            self.status = LobbyStatus.FINISHED
            logger.info("Lobby %s: game finished", self.code)
        self.broadcast(events)
        if finished:
            self._notify_status_change()
        self._schedule_ai()

    def _notify_status_change(self) -> None:
        if self._on_status_change is not None:
            self._on_status_change()

    def _current_ai_seat(self) -> int | None:
        game = self.game
        if self._closed or self.status is not LobbyStatus.PLAYING or game is None:
            return None
        index = game.current_player
        if index is None or not self.seats[index].is_ai:
            return None
        return index

    def _schedule_ai(self) -> None:
        if self._ai_task and not self._ai_task.done():
            return  # the running loop picks up the new state
        if self._current_ai_seat() is not None:
            self._ai_task = asyncio.create_task(self._ai_loop())

    async def _ai_loop(self) -> None:
        loop = asyncio.get_running_loop()
        try:
            while (index := self._current_ai_seat()) is not None:
                game = self.game
                assert game is not None
                started = loop.time()
                wait_until = max(started + self._settings.ai_move_delay, self._pause_until)
                action = await self._bridge.decide(game.observation(index))
                if (remaining := wait_until - loop.time()) > 0:
                    await asyncio.sleep(remaining)
                # The state may have changed while waiting (e.g. game restarted).
                if self.game is not game or self._current_ai_seat() != index:
                    continue
                if not game.is_legal(index, action):
                    continue
                self._after_action(game.apply(index, action))
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception("Lobby %s: AI loop crashed", self.code)


class LobbyManager:
    def __init__(
        self,
        bridge: AIBridge,
        settings: Settings,
        on_games_changed: Callable[[], None] | None = None,
    ) -> None:
        self._bridge = bridge
        self._settings = settings
        self._on_games_changed = on_games_changed
        self._lobbies: dict[str, Lobby] = {}

    def __len__(self) -> int:
        return len(self._lobbies)

    def __iter__(self):
        return iter(list(self._lobbies.values()))

    def active_games(self) -> int:
        return sum(lobby.status is LobbyStatus.PLAYING for lobby in self._lobbies.values())

    def create(self) -> Lobby:
        if len(self._lobbies) >= self._settings.max_lobbies:
            raise LobbyError("Too many open lobbies, please try again later", ERROR_BUSY)
        while True:
            code = "".join(secrets.choice(LOBBY_CODE_ALPHABET) for _ in range(LOBBY_CODE_LENGTH))
            if code not in self._lobbies:
                break
        lobby = Lobby(code, self._bridge, self._settings, self._on_games_changed)
        self._lobbies[code] = lobby
        logger.info("Lobby %s created (%d open)", code, len(self._lobbies))
        return lobby

    def get(self, code: str | None) -> Lobby | None:
        if not code:
            return None
        return self._lobbies.get(code.upper())

    def discard_if_abandoned(self, lobby: Lobby) -> bool:
        if lobby.has_humans:
            return False
        lobby.close()
        self._lobbies.pop(lobby.code, None)
        logger.info("Lobby %s closed (%d open)", lobby.code, len(self._lobbies))
        return True