Update app.py

app.py

@@ -1,95 +1,198 @@
+#!/usr/bin/env python
+
 import asyncio
 import json
 import os
 import secrets
 import signal
-from collections import defaultdict
-from websockets import serve, ConnectionClosed
 
-# Dictionary to store editor content for each user
-editor_content = defaultdict(str)
+import websockets
+
+from connect4 import PLAYER1, PLAYER2, Connect4
+
+
+JOIN = {}
 
-# Set of connected users
-connected_users = set()
+WATCH = {}
 
 
-async def broadcast_changes(websocket, user, change):
+async def error(websocket, message):
     """
-    Broadcast code editing changes to all connected clients.
+    Send an error message.
+
     """
     event = {
-        "type": "edit",
-        "user": user,
-        "change": change,
+        "type": "error",
+        "message": message,
     }
+    await websocket.send(json.dumps(event))
+
+
+async def replay(websocket, game):
+    """
+    Send previous moves.
+
+    """
+    # Make a copy to avoid an exception if game.moves changes while iteration
+    # is in progress. If a move is played while replay is running, moves will
+    # be sent out of order but each move will be sent once and eventually the
+    # UI will be consistent.
+    for player, column, row in game.moves.copy():
+        event = {
+            "type": "play",
+            "player": player,
+            "column": column,
+            "row": row,
+        }
+        await websocket.send(json.dumps(event))
+
+
+async def play(websocket, game, player, connected):
+    """
+    Receive and process moves from a player.
+
+    """
+    async for message in websocket:
+        # Parse a "play" event from the UI.
+        event = json.loads(message)
+        assert event["type"] == "play"
+        column = event["column"]
+
+        try:
+            # Play the move.
+            row = game.play(player, column)
+        except RuntimeError as exc:
+            # Send an "error" event if the move was illegal.
+            await error(websocket, str(exc))
+            continue
+
+        # Send a "play" event to update the UI.
+        event = {
+            "type": "play",
+            "player": player,
+            "column": column,
+            "row": row,
+        }
+        websockets.broadcast(connected, json.dumps(event))
+
+        # If move is winning, send a "win" event.
+        if game.winner is not None:
+            event = {
+                "type": "win",
+                "player": game.winner,
+            }
+            websockets.broadcast(connected, json.dumps(event))
+
+
+async def start(websocket):
+    """
+    Handle a connection from the first player: start a new game.
+
+    """
+    # Initialize a Connect Four game, the set of WebSocket connections
+    # receiving moves from this game, and secret access tokens.
+    game = Connect4()
+    connected = {websocket}
 
-    # Broadcast the change to all connected clients
-    for user_socket in connected_users:
-        if user_socket != websocket:
-            await user_socket.send(json.dumps(event))
+    join_key = secrets.token_urlsafe(12)
+    JOIN[join_key] = game, connected
 
+    watch_key = secrets.token_urlsafe(12)
+    WATCH[watch_key] = game, connected
 
-async def handle_code_editing(websocket, user):
+    try:
+        # Send the secret access tokens to the browser of the first player,
+        # where they'll be used for building "join" and "watch" links.
+        event = {
+            "type": "init",
+            "join": join_key,
+            "watch": watch_key,
+        }
+        await websocket.send(json.dumps(event))
+        # Receive and process moves from the first player.
+        await play(websocket, game, PLAYER1, connected)
+    finally:
+        del JOIN[join_key]
+        del WATCH[watch_key]
+
+
+async def join(websocket, join_key):
     """
-    Handle code editing events from a user.
+    Handle a connection from the second player: join an existing game.
+
     """
-    connected_users.add(websocket)
+    # Find the Connect Four game.
+    try:
+        game, connected = JOIN[join_key]
+    except KeyError:
+        await error(websocket, "Game not found.")
+        return
 
+    # Register to receive moves from this game.
+    connected.add(websocket)
     try:
-        # Send the initial state of the document to the new user
-        await websocket.send(json.dumps({"type": "init", "content": editor_content[user]}))
+        # Send the first move, in case the first player already played it.
+        await replay(websocket, game)
+        # Receive and process moves from the second player.
+        await play(websocket, game, PLAYER2, connected)
+    finally:
+        connected.remove(websocket)
 
-        async for message in websocket:
-            event = json.loads(message)
-            assert event["type"] == "edit"
-            change = event["change"]
 
-            # Apply the change to the document
-            editor_content[user] = change["text"]
+async def watch(websocket, watch_key):
+    """
+    Handle a connection from a spectator: watch an existing game.
 
-            # Broadcast the change to all connected clients
-            await broadcast_changes(websocket, user, change)
+    """
+    # Find the Connect Four game.
+    try:
+        game, connected = WATCH[watch_key]
+    except KeyError:
+        await error(websocket, "Game not found.")
+        return
 
-    except ConnectionClosed:
-        pass
+    # Register to receive moves from this game.
+    connected.add(websocket)
+    try:
+        # Send previous moves, in case the game already started.
+        await replay(websocket, game)
+        # Keep the connection open, but don't receive any messages.
+        await websocket.wait_closed()
     finally:
-        # Remove the user when the WebSocket connection is closed
-        connected_users.remove(websocket)
+        connected.remove(websocket)
 
 
-async def editor_handler(websocket, path):
+async def handler(websocket):
     """
     Handle a connection and dispatch it according to who is connecting.
+
     """
     # Receive and parse the "init" event from the UI.
     message = await websocket.recv()
     event = json.loads(message)
     assert event["type"] == "init"
 
-    user = event.get("user", None)
-
-    if not user:
-        # Assign a random username for simplicity
-        user = secrets.token_urlsafe(8)
-
-    try:
-        await handle_code_editing(websocket, user)
-    except Exception as e:
-        print(f"Error: {e}")
-    finally:
-        connected_users.remove(websocket)
+    if "join" in event:
+        # Second player joins an existing game.
+        await join(websocket, event["join"])
+    elif "watch" in event:
+        # Spectator watches an existing game.
+        await watch(websocket, event["watch"])
+    else:
+        # First player starts a new game.
+        await start(websocket)
 
 
 async def main():
-    # Set the stop condition when receiving SIGTERM
+    # Set the stop condition when receiving SIGTERM.
     loop = asyncio.get_running_loop()
     stop = loop.create_future()
     loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
 
-    
-    async with serve(editor_handler, "0.0.0.0", 7860):
+    port = int(os.environ.get("PORT", "7860"))
+    async with websockets.serve(handler, "0.0.0.0", port):
         await stop
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(main())