magician/server_tcp.py

import json
import random
import sys

import ssl
import asyncio
from websockets.asyncio.server import serve

MAX_LOBBY_CLIENT_COUNT = 6

# Generates a 6-character code from the letters A-Z
def generate_code() -> str:
    result = ""
    for i in range(6):
        result += chr(random.randint(0, 25) + 65)
    return result

class Lobby:
    
    def __init__(self, lobby_id, host_client_id):
        self.lobby_id = lobby_id
        self.host_client_id = host_client_id
        self.connected_clients = {}
        self.client_playing_order = []
        self.client_bets = {} # client_id: bet
        # self.currently_playing_client = 0 # index of client_playing_order
        self.current_round = 0

class Client:

    def __init__(self, client_id, client_name, websocket):
        self.client_id = client_id
        self.client_name = client_name
        self.websocket = websocket

# Class that holds active lobbies and manages communication and information exchange
class LobbyManager:

    def __init__(self):
        self.active_lobbies = {}

    #region lobby moderation

    # Creates a new lobby and assigns a host client
    def create_lobby(self, lobby_id, host_client_id) -> None:
        # Add lobby
        self.active_lobbies[lobby_id] = Lobby(lobby_id, host_client_id)
        print(f"Created lobby {lobby_id}")

    # Joins a client to an existing lobby
    async def join_lobby(self, lobby_id, client_id, client_name, websocket) -> str:
        if not self.lobby_id_exists(lobby_id):
            # Cannot join; lobby does not exist
            return "lobby_not_found"  

        if self.get_lobby_player_count(lobby_id) >= MAX_LOBBY_CLIENT_COUNT:
            # Lobby full! Cannot join
            return "lobby_full"  
        
        # Add client to lobby
        self.active_lobbies[lobby_id].connected_clients[client_id] = Client(
            client_id,
            client_name,
            websocket,
        )
        print(f"Client {client_id} joined lobby {lobby_id}")

        # Notify the lobby of the newly-joined client
        await self.broadcast_message(
            lobby_id,
            json.dumps({
                "response": "client_joined_lobby",
                "new_client_id": client_id,
                "new_client_name": client_name,
                "connected_client_ids": self.get_lobby_client_ids(lobby_id),
                "connected_client_names": self.get_lobby_client_names(lobby_id),
            })
        )

        return "ok"
    
    # Removes a client from a lobby
    async def leave_lobby(self, lobby_id, client_id) -> None:
        try:
            # Get client name before deleting them
            client_name = self.get_lobby_client_name(lobby_id, client_id)

            # Delete client
            del self.active_lobbies[lobby_id].connected_clients[client_id]
            
            # Notify remaining clients of leaving client and remaining clients
            await self.broadcast_message(
                lobby_id,
                json.dumps({
                    "response": "client_left_lobby",
                    "status": "ok",
                    "client_name": client_name,
                    "connected_client_ids": self.get_lobby_client_ids(lobby_id),
                    "connected_client_names": self.get_lobby_client_names(lobby_id),
                }),
            )
            print(f"Deleted client {client_id} from lobby {lobby_id}")
        except KeyError:
            print(f"Client with ID {client_id} in lobby {lobby_id} not found, could not be deleted")

    # Deletes a lobby. May only be called after ensuring that all clients have left
    def delete_lobby(self, lobby_id) -> None:
        try:
            del self.active_lobbies[lobby_id]
            print(f"Deleted lobby {lobby_id}")
        except KeyError:
            print(f"Lobby with ID {lobby_id} not found, could not be deleted")

    #endregion

    #region lobby info

    # Get lobby by its ID
    def get_lobby(self, lobby_id) -> Lobby:
        return self.active_lobbies[lobby_id]

    # Checks if a lobby exists under the given ID
    def lobby_id_exists(self, lobby_id) -> bool:
        return lobby_id in self.active_lobbies

    #endregion

    #region lobby client info

    # Get count of clients currently connected to a lobby
    def get_lobby_player_count(self, lobby_id) -> int:
        return len(self.active_lobbies[lobby_id].connected_clients)
    
    # Get name of a client connected to a given lobby
    def get_lobby_client_name(self, lobby_id, client_id) -> str:
        return self.active_lobbies[lobby_id].connected_clients[client_id].client_name
    
    # Get names of all clients connected to a given lobby
    def get_lobby_client_names(self, lobby_id) -> list[str]:
        names = []
        lobby = self.get_lobby(lobby_id)
        for client_id in lobby.connected_clients:
            names.append(lobby.connected_clients[client_id].client_name)
        return names

    def get_lobby_client_ids(self, lobby_id) -> list[str]:
        return list(self.get_lobby(lobby_id).connected_clients.keys())
    
    # Checks if a client is connected to a given lobby under the given ID
    def client_id_exists_in_lobby(self, lobby_id, client_id) -> bool:
        if not self.lobby_id_exists(lobby_id):
            return False
        return client_id in self.active_lobbies[lobby_id].connected_clients

    # Gets the ID of the host client of a given lobby
    def get_lobby_host_client_id(self, lobby_id) -> str:
        return self.active_lobbies[lobby_id].host_client_id
    
    #endregion

    #region communication

    # Sends a message to a single client in a given lobby.
    async def send_message(self, lobby_id, client_id, message) -> None:
        await self.active_lobbies[lobby_id].connected_clients[client_id].websocket.send(message)

    # Sends a message to all clients in a given lobby.
    async def broadcast_message(self, lobby_id, message) -> None:
        connected_clients = self.active_lobbies[lobby_id].connected_clients
        for client_id in connected_clients:
            await connected_clients[client_id].websocket.send(message)
    
    #endregion
    
    #region gameplay

    async def spread_cards(self, lobby_id, round, cards, playing_order) -> None:
        client_ids = self.get_lobby_client_ids(lobby_id)
        
        trump = None
        if cards["trump"]:
            trump = cards["trump"]
        cards.pop("trump")
        
        for client_id in client_ids:
            await self.send_message(
                lobby_id,
                client_id,
                json.dumps({
                    "response": "new_round_started",
                    "current_round": round,
                    "cards": cards,
                    "trump": trump,
                    "playing_order": playing_order,
                })
            )
    
    async def request_bet(self, lobby_id, client_id, disallowed_bet) -> None:
        await self.send_message(
            lobby_id,
            client_id,
            json.dumps({
                "response": "request_bet",
                "disallowed_bet": disallowed_bet,
                "round": self.get_lobby(lobby_id).current_round,
            })
        )

    async def spread_bet(self, lobby_id, client_id, bet) -> None:
        await self.broadcast_message(
            lobby_id,
            json.dumps({
                "response": "client_submitted_bet",
                "client_id": client_id,
                "bet": bet,
            }),
        )
    
    async def spread_played_card(self, lobby_id, client_id, played_card) -> None:
        await self.broadcast_message(
            lobby_id,
            json.dumps({
                "response": "client_played_card",
                "client_id": client_id,
                "played_card": played_card,
            }),
        )
    
    async def spread_play_winner(self, lobby_id, client_id) -> None:
        await self.broadcast_message(
            lobby_id,
            json.dumps({
                "response": "client_won_play",
                "client_id": client_id,
            }),
        )

    #endregion

    #region miscellaneous
    
    # Generates an ID and ensures uniqueness among all lobbies.
    def get_unique_lobby_id(self) -> str:
        id = generate_code()
        if self.lobby_id_exists(id):
            return self.get_unique_lobby_id()
        return id
    
    # Generates an ID and ensures uniqueness among all clients in a given lobby.
    def get_unique_client_id(self, lobby_id) -> str:
        id = generate_code()
        if self.client_id_exists_in_lobby(lobby_id, id):
            return self.get_unique_client_id(lobby_id)
        return id

    #endregion

lobby_manager = LobbyManager()

# Listen for incoming connections
async def listen_websocket(run_local: bool) -> None:
    ssl_context = None
    if not run_local:
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ssl_context.load_cert_chain("/etc/letsencrypt/live/denizk0461.dev/fullchain.pem", "/etc/letsencrypt/live/denizk0461.dev/privkey.pem")
    async with serve(start_listen_client_websocket, "127.0.0.1", 9974, ssl=ssl_context) as server:
        await server.serve_forever()

# Connect to a websocket and listen to their messages
async def start_listen_client_websocket(websocket) -> None:
    async for message in websocket:
        # await 
        await process_input(message, websocket)

# Retrieves and processes an input from a connected websocket
async def process_input(message: str, websocket) -> None:
    data_object = json.loads(message)

    # Create a new lobby
    if data_object["response"] == "create_lobby":
        new_lobby_id = lobby_manager.get_unique_lobby_id()
        new_client_id = lobby_manager.get_unique_client_id(new_lobby_id)
        
        # Create lobby with requesting client as host
        lobby_manager.create_lobby(
            new_lobby_id,
            new_client_id,
        )

        # Join requesting client to new lobby.
        # No error handling here, since lobby was just created, so it
        # must exist and be empty
        await lobby_manager.join_lobby(
            new_lobby_id,
            new_client_id,
            data_object["client_name"],
            websocket,
        )

        # Send success confirmation
        await lobby_manager.send_message(
            new_lobby_id,
            new_client_id,
            json.dumps({
                "response": "create_lobby",
                "status": "ok",
                "lobby_id": new_lobby_id,
                "client_id": new_client_id,
            }),
        )

        print(f"Client {new_client_id} created lobby {new_lobby_id}")
    
    # Join client to existing lobby
    elif data_object["response"] == "join_lobby":
        lobby_id = data_object["lobby_id"]
        client_id = lobby_manager.get_unique_client_id(lobby_id)
        client_name = data_object["client_name"]

        # Join client to lobby
        response = await lobby_manager.join_lobby(
            lobby_id,
            client_id,
            client_name,
            websocket,
        )

        match response:
            case "lobby_not_found":
                await websocket.send(
                    json.dumps({
                        "response": "join_lobby_response",
                        "status": "lobby_not_found",
                    })
                )
                
                print(f"Client {client_name} could not connect; lobby {lobby_id} does not exist")
            
            case "lobby_full":
                await websocket.send(
                    json.dumps({
                        "response": "join_lobby_response",
                        "status": "lobby_full",
                    })
                )

                print(f"Client {client_name} could not connect; lobby {lobby_id} is full")
            
            case "ok":
                # Send success confirmation
                await lobby_manager.send_message(
                    lobby_id,
                    client_id,
                    json.dumps({
                        "response": "join_lobby_response",
                        "status": "ok",
                        "client_id": client_id,
                    }),
                )

                print(f"Client {client_id} joined lobby {lobby_id}")

    # Disconnect client from lobby
    elif data_object["response"] == "leave_lobby":
        lobby_id = data_object["lobby_id"]
        client_id = data_object["client_id"]
        host_client_id = lobby_manager.get_lobby_host_client_id(lobby_id)

        if host_client_id == client_id:
            # Client is host; close lobby
            await lobby_manager.leave_lobby(lobby_id, client_id)
            await lobby_manager.broadcast_message(
                lobby_id,
                json.dumps({
                    "response": "leave_lobby_request",
                })
            )
            print(f"Client {client_id} was host; broadcasting leave request")
        else:
            # Client is not host; client leaves only
            await lobby_manager.leave_lobby(lobby_id, client_id)
        
        # Delete empty lobby
        if lobby_manager.get_lobby_player_count(lobby_id) == 0:
            lobby_manager.delete_lobby(lobby_id)
            print(f"Lobby {lobby_id} was disbanded")
    
    # Relays a chat message to all clients connected to a lobby
    elif data_object["response"] == "send_chat_message":
        client_name = lobby_manager.get_lobby_client_name(data_object["lobby_id"], data_object["client_id"])
        await lobby_manager.broadcast_message(
            data_object["lobby_id"],
            json.dumps({
                "response": "receive_chat_message",
                "message": data_object["message"],
                "client_name": client_name,
            })
        )
        print(f"Relayed chat message by {client_name}")
    
    # Start game (host only)
    elif data_object["response"] == "start_game":
        lobby_id = data_object["lobby_id"]
        lobby = lobby_manager.get_lobby(lobby_id)
        lobby.client_playing_order = list(lobby.connected_clients.keys())
        random.shuffle(lobby.client_playing_order)
        
        await lobby_manager.spread_cards(
            data_object["lobby_id"],
            1,
            data_object["cards"],
            lobby.client_playing_order,
        )

        lobby_manager.get_lobby(data_object["lobby_id"]).current_round = 1
    
        print(f"Lobby {lobby_id} is starting the game, playing order is: {lobby.client_playing_order}")

        await lobby_manager.request_bet(
            lobby_id,
            lobby_manager.get_lobby(lobby_id).client_playing_order[0],
            -1,
        )

    elif data_object["response"] == "start_next_play":
        lobby_id = data_object["lobby_id"]
        await lobby_manager.broadcast_message(
            lobby_id,
            json.dumps({
                "response": "begin_next_play",
            }),
        )
        await lobby_manager.send_message(
            lobby_id,
            lobby_manager.get_lobby(lobby_id).client_playing_order[0],
            json.dumps({
                "response": "request_card",
            }),
        )

    # Start next round (host only)
    elif data_object["response"] == "start_next_round":
        lobby_id = data_object["lobby_id"]

        # Reset bets
        lobby_manager.get_lobby(lobby_id).client_bets = {}
        lobby_manager.get_lobby(data_object["lobby_id"]).current_round += 1

        await lobby_manager.spread_cards(
            lobby_id,
            data_object["round"],
            data_object["cards"],
            None,
        )

        await lobby_manager.request_bet(
            lobby_id,
            lobby_manager.get_lobby(lobby_id).client_playing_order[0],
            -1,
        )

    elif data_object["response"] == "submit_bet":
        lobby_id = data_object["lobby_id"]
        client_id = data_object["client_id"]
        bet = data_object["bet"]

        lobby_manager.get_lobby(lobby_id).client_bets[client_id] = bet

        await lobby_manager.spread_bet(
            lobby_id,
            client_id,
            bet,
        )

        last_client_index = lobby_manager.get_lobby(lobby_id).client_playing_order.index(client_id)
        player_count = lobby_manager.get_lobby_player_count(lobby_id)

        if last_client_index == (player_count - 1):
            # all have bet; start game
            await lobby_manager.broadcast_message(
                lobby_id,
                json.dumps({
                    "response": "begin_playing_cards",
                }),
            )

            print("Sent game start")

            playing_client_id = lobby_manager.get_lobby(lobby_id).client_playing_order[0]

            await lobby_manager.send_message(
                lobby_id,
                playing_client_id,
                json.dumps({
                    "response": "request_card",
                    "round": lobby_manager.get_lobby(lobby_id).current_round,
                }),
            )

            print(f"Playing client is {playing_client_id}")
        else:
            # bets are not finished; continue asking for bets
            disallowed_bet = -1
            if last_client_index == (player_count - 2):
                # last client to bet; limit bet options
                total_bet = sum(lobby_manager.get_lobby(lobby_id).client_bets.values())
                disallowed_bet = data_object["round"] - total_bet
                print(f"bets: {data_object["round"]} / {total_bet}")
                if disallowed_bet < 0:
                    disallowed_bet = -1

            await lobby_manager.request_bet(
                lobby_id,
                lobby_manager.get_lobby(lobby_id).client_playing_order[last_client_index + 1],
                disallowed_bet,
            )

    elif data_object["response"] == "play_card":
        lobby_id = data_object["lobby_id"]
        client_id = data_object["client_id"]

        await lobby_manager.spread_played_card(
            lobby_id,
            client_id,
            data_object["played_card"],
        )
        playing_order = lobby_manager.get_lobby(lobby_id).client_playing_order
        client_index = playing_order.index(client_id)
        if client_index == (len(playing_order) - 1):
            # end round
            await lobby_manager.send_message(
                lobby_id,
                lobby_manager.get_lobby_host_client_id(lobby_id),
                json.dumps({
                    "response": "determine_play_winner_request",
                })
            )
        else:
            # not all clients have played a card yet; continue
            playing_client_id = playing_order[client_index + 1]

            await lobby_manager.send_message(
                lobby_id,
                playing_client_id,
                json.dumps({
                    "response": "request_card",
                }),
            )
    
    elif data_object["response"] == "submit_play_winner":
        await lobby_manager.spread_play_winner(
            data_object["lobby_id"],
            data_object["client_id"],
        )

        print(f"DING DING DING we have a winner!!! {data_object['client_id']}")
    
    elif data_object["response"] == "end_game":
        await lobby_manager.broadcast_message(
            data_object["lobby_id"],
            json.dumps({
                "response": "end_game_request",
            })
        )

if __name__ == '__main__':
    run_local = False
    # Run locally without SSL certificate
    if len(sys.argv) > 1 and sys.argv[1] == "local":
        run_local = True
    
    print('Listening on *:9974')
    asyncio.run(listen_websocket(run_local))