Update app.py

app.py

@@ -1,63 +1,95 @@
 import asyncio
-import websockets
-import random
-import string
-import sqlite3
-
-# Database setup (using SQLite for simplicity)
-conn = sqlite3.connect('pairing_database.db')
-cursor = conn.cursor()
-cursor.execute('CREATE TABLE IF NOT EXISTS pairs (key TEXT, client_ws TEXT)')
-conn.commit()
-
-async def handle_client(websocket, path):
+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)
+
+# Set of connected users
+connected_users = set()
+
+
+async def broadcast_changes(websocket, user, change):
+    """
+    Broadcast code editing changes to all connected clients.
+    """
+    event = {
+        "type": "edit",
+        "user": user,
+        "change": change,
+    }
+
+    # Broadcast the change to all connected clients
+    for user_socket in connected_users:
+        if user_socket != websocket:
+            await user_socket.send(json.dumps(event))
+
+
+async def handle_code_editing(websocket, user):
+    """
+    Handle code editing events from a user.
+    """
+    connected_users.add(websocket)
+
     try:
-        # Receive the client's request (create a pair or connect to a key)
-        request_type = await websocket.recv()
+        # Send the initial state of the document to the new user
+        await websocket.send(json.dumps({"type": "init", "content": editor_content[user]}))
+
+        async for message in websocket:
+            event = json.loads(message)
+            assert event["type"] == "edit"
+            change = event["change"]
 
-        if request_type == "create_pair":
-            # Generate a unique key
-            unique_key = generate_key()
+            # Apply the change to the document
+            editor_content[user] = change["text"]
 
-            # Save the key and associate it with the current client in the database
-            cursor.execute('INSERT INTO pairs VALUES (?, ?)', (unique_key, str(websocket)))
-            conn.commit()
+            # Broadcast the change to all connected clients
+            await broadcast_changes(websocket, user, change)
 
-            # Inform the client about the created key
-            await websocket.send(f"Pair key created: {unique_key}")
+    except ConnectionClosed:
+        pass
+    finally:
+        # Remove the user when the WebSocket connection is closed
+        connected_users.remove(websocket)
 
-        elif request_type == "connect_to_key":
-            # Receive the key from the client
-            key_to_connect = await websocket.recv()
 
-            # Retrieve all clients associated with the key from the database
-            cursor.execute('SELECT client_ws FROM pairs WHERE key=?', (key_to_connect,))
-            results = cursor.fetchall()
+async def editor_handler(websocket, path):
+    """
+    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"
 
-            if results:
-                # Pair the current client with all clients associated with the key
-                for result in results:
-                    other_client_ws = result[0]
-                    if other_client_ws != str(websocket):
-                        await websocket.send(f"Connected to key {key_to_connect}. You are paired with {other_client_ws}")
-                        await other_client_ws.send(f"Connected to key {key_to_connect}. You are paired with {str(websocket)}")
-            else:
-                # Inform the client that the key does not exist
-                await websocket.send(f"Key {key_to_connect} does not exist.")
+    user = event.get("user", None)
 
-    except websockets.exceptions.ConnectionClosed:
-        pass  # Connection closed, handle appropriately
+    if not user:
+        # Assign a random username for simplicity
+        user = secrets.token_urlsafe(8)
 
-def generate_key():
-    # Generate a random key (you might want to implement a more robust key generation logic)
-    return ''.join(random.choices(string.ascii_uppercase + string.digits, k=8))
+    try:
+        await handle_code_editing(websocket, user)
+    except Exception as e:
+        print(f"Error: {e}")
+    finally:
+        connected_users.remove(websocket)
 
-start_server = websockets.serve(handle_client, "localhost", 8765)
 
 async def main():
-    print("WebSocket server is running. Waiting for connections...")
-    await start_server
-    await asyncio.Event().wait()
+    # 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):
+        await stop
+
 
 if __name__ == "__main__":
     asyncio.run(main())