api/cah/websocket_conn.py

#!/usr/bin/env python3.12

import json
import time

from fastapi import WebSocket
from cah.constructors import CAHClient

class ConnectionManager:
    def __init__(self):
        self.active_connections: dict = {}

    def get_connection_by_ws(self, websocket: WebSocket) -> WebSocket:
        return self.active_connections.get(websocket)
    
    def get_connection_by_csid(self, csid: str) -> WebSocket:
        for connection in self.active_connections:
            if connection.get('csid') == csid:
                return connection
            
    async def send_client_and_game_lists(self, state, websocket: WebSocket):
        clients = []
        games = [game.__dict__ for game in state.games]
        
        for ws, client in self.active_connections.items():
            print(f"Client: {client}")
            _client = client.get('client')
            clients.append({
                'resource': _client.resource,
                'platform': _client.platform,
                'connected_at': _client.connected_at,
                'current_game': _client.current_game if _client.current_game else None,
            })
        
        await websocket.send_json({
            "event": "client_list",
            "ts": int(time.time()),
            "data": {
                "clients": clients
            }
        })
        await websocket.send_json({
            "event": "game_list",
            "ts": int(time.time()),
            "data":
            {
                "games": state.get_games()
            }
        })
    
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[websocket] = {
            'client': None,
            'websocket': websocket,
        }
 

    async def handshake_complete(self,
                                 state,
                                 websocket: WebSocket,
                                 csid: str,
                                 handshakedClient: CAHClient):
          
          self.active_connections[websocket] = {
              'websocket': websocket,
              'csid': csid,
              'client': handshakedClient,
              }
          
          await self.broadcast({
                "event": "client_connected",
                "ts": int(time.time()),
                "data": {
                    "connected_resource": handshakedClient.resource,
                    "connected_platform": handshakedClient.platform, 
                }
            })
          
          await self.send_client_and_game_lists(state, 
                                                websocket)        

    def disconnect(self, websocket: WebSocket, csid: str = None):
        self.active_connections.pop(websocket)

    async def send(self, message: str, websocket: WebSocket):
        await websocket.send_json(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_json(message)