magician/server.py

import json
import random
import socket
import sys
import threading

import asyncio
from websockets.asyncio.server import serve

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 = {}

class Client:

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

class LobbyManager:

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

    def get_lobby(self, lobby_id) -> Lobby:
        return self.active_lobbies[lobby_id]

    def lobby_id_exists(self, lobby_id) -> bool:
        return lobby_id in self.active_lobbies
    
    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
    
    def get_unique_lobby_id(self) -> str:
        id = generate_code()
        if self.lobby_id_exists(id):
            return self.get_unique_lobby_id()
        return id
    
    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

    def create_lobby(self, lobby_id, host_client_id) -> None:
        self.active_lobbies[lobby_id] = Lobby(lobby_id, host_client_id)
        print(f"Created lobby {lobby_id}")

    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")

    async def join_lobby(self, lobby_id, client_id, client_name, websocket) -> None:
        self.active_lobbies[lobby_id].connected_clients[client_id] = Client(
            client_id,
            client_name,
            websocket,
        )
        print(f"Client {client_id} joined lobby {lobby_id}")

        await self.broadcast_message(
            lobby_id,
            json.dumps({
                "response": "client_joined_lobby",
                "new_client_id": client_id,
                "new_client_name": client_name,
            })
        )
    
    async def leave_lobby(self, lobby_id, client_id) -> None:
        try:
            # get client name BEFORE deleting the client...
            client_name = self.get_client_name(lobby_id, client_id)
            del self.active_lobbies[lobby_id].connected_clients[client_id]
            await self.broadcast_message(
                lobby_id,
                json.dumps({
                    "response": "client_left_lobby",
                    "status": "ok",
                    "client_name": client_name,
                }),
            )
            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")

    def get_lobby_player_count(self, lobby_id) -> int:
        return len(self.active_lobbies[lobby_id].connected_clients)
    
    def get_client_name(self, lobby_id, client_id) -> str:
        return self.active_lobbies[lobby_id].connected_clients[client_id].client_name

    def get_lobby_host_client_id(self, lobby_id) -> str:
        return self.active_lobbies[lobby_id].host_client_id

    async def send_message(self, lobby_id, client_id, message) -> None:
        await self.active_lobbies[lobby_id].connected_clients[client_id].websocket.send(message)

    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)

# class ClientConnection:

#     def __init__(self, connection, address, lobby_manager):
#         self.connection = connection
#         self.address = address
#         self.lobby_manager = lobby_manager

#         self.listen_client()

#     # Listens to a connected TCP client.
#     def listen_client(self) -> None:
#         try:
#             print(f"Connected to {self.address[0]}:{self.address[1]}")

#             received_data = ""
#             while True:
#                 data = self.connection.recv(1)
#                 received_data += data.decode("utf-8")
#                 if received_data[-1] == "}":
#                     print(received_data)
#                     self.process_input(received_data)
#                     received_data = ""
                
#                 if not data:
#                     break
#         except IndexError:
#             print("Client force-disconnected")
#         except ConnectionResetError:
#             print("Client force-disconnected")
#         finally:
#             self.connection.close()
#             print(f"Closing connection to {self.address[0]}:{self.address[1]}")

    

lobby_manager = LobbyManager()

async def listen_websocket() -> None:
    async with serve(start_listen_client_websocket, "localhost", 9974) as server:
        await server.serve_forever()

async def start_listen_client_websocket(websocket) -> None:
    async for message in websocket:
        # await 
        await process_input(message, websocket)

async def process_input(message: str, websocket) -> None:
    data_object = json.loads(message)

    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)
        print(new_client_id + " is creating lobby " + new_lobby_id)

        lobby_manager.create_lobby(
            new_lobby_id,
            new_client_id,
        )
        await lobby_manager.join_lobby(
            new_lobby_id,
            new_client_id,
            data_object["client_name"],
            websocket,
        )
        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,
            }),
        )
    
    elif data_object["response"] == "join_lobby":
        print("joining lobby")
        lobby_id = data_object["lobby_id"]
        new_client_id = lobby_manager.get_unique_client_id(lobby_id)

        await lobby_manager.join_lobby(
            lobby_id,
            new_client_id,
            data_object["client_name"],
            websocket,
        )
        await lobby_manager.send_message(
            lobby_id,
            new_client_id,
            json.dumps({
                "response": "join_lobby",
                "status": "ok",
                "client_id": new_client_id,
            }),
        )

        # send message to all connected clients
    
    elif data_object["response"] == "leave_lobby":
        print("leaving lobby")
        lobby_id = data_object["lobby_id"]
        host_client_id = lobby_manager.get_lobby_host_client_id(lobby_id)

        if host_client_id == data_object["client_id"]:
            # client is host; close lobby
            await lobby_manager.leave_lobby(lobby_id, data_object["client_id"])
            await lobby_manager.broadcast_message(
                lobby_id,
                json.dumps({
                    "response": "leave_lobby_request",
                })
            )
            if lobby_manager.get_lobby_player_count(lobby_id) == 0:
                lobby_manager.delete_lobby(lobby_id)
        else:
            # client is not host; client leaves only
            await lobby_manager.leave_lobby(lobby_id, data_object["client_id"])
    
    elif data_object["response"] == "send_chat_message":
        print("relaying chat message")
        await lobby_manager.broadcast_message(
            data_object["lobby_id"],
            json.dumps({
                "response": "receive_chat_message",
                "message": data_object["message"],
                "client_name": lobby_manager.get_client_name(data_object["lobby_id"], data_object["client_id"]),
            })
        )

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: ./server.py PORT")
        sys.exit(1)

    port = int(sys.argv[1])
    # tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # server_address = ('168.119.97.227', port)
    # tcp_socket.bind(server_address)
    # tcp_socket.listen(1)
    # tcp_socket.settimeout(0.5)

    print('Listening on *:%d' % (port))

    # listen()
    asyncio.run(listen_websocket())