back-python/src/controller/lobbyClient.py

"""Handles communication between one WebSocket tunnel and the game"""

from typing import Callable, Any, TYPE_CHECKING, Union
from pydantic import ValidationError

from websockets.server import WebSocketServerProtocol
from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError

from netcode.models import Message, LoginMessage, MessageType

from .auth import discord_process_code

from controller import ClientController

from .session_storage import generate_cookie, session_storage

if TYPE_CHECKING:
    from controller import LobbyController


def vibe_check(t):
    """Enforces the selected type for the message"""

    def decorator(func):
        async def new_f(self, message):
            try:
                message = t.model_validate_json(message)
                return await func(self, message)
            except ValidationError as e:
                print(e)
                await self.send_error(e.errors())

        return new_f

    return decorator


class LobbyClientController:
    """Handles communication between one WebSocket client and the lobby"""

    websocket: WebSocketServerProtocol
    lobby_controller: "LobbyController"
    track_connection: Callable

    discord_uid: str | None
    username: str | None
    discord_avatar: str | None

    game_client: Union["ClientController", None]

    def __init__(
        self,
        lobby: "LobbyController",
        websocket: WebSocketServerProtocol,
    ):
        self.lobby_controller = lobby
        self.websocket = websocket
        self.on_message = self.login_context
        self.lobby_controller.clients.append(self)
        self.discord_uid = None
        self.username = None
        self.discord_avatar = None
        self.game_client = None

    async def listen(self) -> None:
        """Hooks the client to the correct listen callbacks"""
        await self.send(
            Message(
                type=MessageType.CONTEXT,
                data_type="game_context_id",
                data="login",
            )
        )
        try:
            async for message in self.websocket:
                await self.on_message(message)
        except ConnectionClosedError:
            print("Connection Closed Error")

    async def close(self):
        """Remove self from Lobby broadcasting list"""
        if self.game_client is not None:
            await self.game_client.close()
        self.lobby_controller.clients.remove(self)

    async def send(self, message: Message) -> None:
        """Send a Message object to the connected client"""
        try:
            await self.websocket.send(message.model_dump_json())
        except ConnectionClosedOK:
            print("Connection Closed")

    async def send_error(self, error: str) -> None:
        """Send an error string to the connected client"""
        await self.send(Message(type=MessageType.ERROR, data_type="error", data=error))

    async def login_user(self, user: dict):
        self.discord_uid = user["id"]
        self.username = user["username"]
        self.discord_avatar = user["avatar"]
        self.on_message = self.lobby_context
        cookie = generate_cookie(user)
        await self.send(
            Message(type=MessageType.SET, data_type="session_cookie", data=cookie)
        )
        await self.send(
            Message(
                type=MessageType.CREATE,
                data_type="object",
                data=self.lobby_controller.lobby.serialize(),
            )
        )
        await self.send(
            Message(
                type=MessageType.CONTEXT,
                data_type="game_context_id",
                data="lobby",
            )
        )

    @vibe_check(LoginMessage)
    async def login_context(self, message: LoginMessage) -> None:
        if message.type == MessageType.DISCORD_LOGIN:
            user = discord_process_code(message.data)
            if user is None:
                return
            if "id" in user:
                await self.login_user(user)
        if message.type == MessageType.COOKIE_LOGIN:
            if message.data in session_storage:
                await self.login_user(session_storage[message.data])
        else:
            print("user not connected")

    @vibe_check(Message)
    async def lobby_context(self, message: Message) -> None:
        if message.type == MessageType.CREATE:
            if message.data_type == "game":
                await self.lobby_controller.create_game()
        if message.type == MessageType.JOIN:
            if message.data_type == "game":
                game_ctrler = next(
                    (
                        g
                        for g in self.lobby_controller.games
                        if str(g.game.uuid) == message.data
                    ),
                    None,
                )
                if game_ctrler is None:
                    return print("Game not found")
                self.game_client = ClientController(game_ctrler, self.websocket)
                self.on_message = self.game_context
                await self.game_client.game_controller.announce_game(self.websocket)
                await self.game_client.join_game(
                    game_ctrler.game,
                    self.username,
                    self.discord_uid,
                    self.discord_avatar,
                )

    async def game_context(self, message: Message) -> None:
        if self.game_client is None:
            return
        await self.game_client.on_message(message)