Add WS support to matches (breaking change)

- Fixed issue where ELO rating could go in the negatives, read more [here](Mercury/dth-pingpong-mobileapp#3 (comment))
- Initial websocket support for multi-screen matchmaking, [#1](Mercury/dth-pingpong-mobileapp#1)

main.py

@@ -1,4 +1,5 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, WebSocket
+from typing import Dict, List
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
 from calls import *
@@ -15,6 +16,9 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+active_connections: Dict[int, List[WebSocket]] = {}
+
+
 class RegisterRequest(BaseModel):
     email: str
     display_name: str
@@ -152,6 +156,9 @@ def create_match(request: CreateMatchRequest):
         )
         match_id = cursor.fetchone()["match_id"]
         conn.commit()
+
+        active_connections[match_id] = []
+
         return {"match_id": match_id}
     except Exception as e:
         conn.rollback()
@@ -160,7 +167,6 @@ def create_match(request: CreateMatchRequest):
         cursor.close()
         conn.close()
 
-
 @app.post("/joinmatch")
 def join_match(request: JoinMatchRequest):
     conn = get_db_connection()
@@ -194,6 +200,7 @@ def join_match(request: JoinMatchRequest):
             (player2_uid, request.match_id)
         )
         conn.commit()
+
         return {"message": "Joined match successfully", "match_id": request.match_id}
     except Exception as e:
         conn.rollback()
@@ -202,6 +209,7 @@ def join_match(request: JoinMatchRequest):
         cursor.close()
         conn.close()
 
+
 def calculate_elo(player1_elo, player2_elo, player1_score, player2_score):
     k_factor = 32
     expected1 = 1 / (1 + 10 ** ((player2_elo - player1_elo) / 400))
@@ -213,7 +221,7 @@ def calculate_elo(player1_elo, player2_elo, player1_score, player2_score):
     return elo_change1, elo_change2
 
 @app.post("/endmatch")
-def end_match(request: EndMatchRequest):
+async def end_match(request: EndMatchRequest):
     conn = get_db_connection()
     cursor = conn.cursor()
     try:
@@ -243,9 +251,7 @@ def end_match(request: EndMatchRequest):
         player1_elo = next(p["current_elo"] for p in players if p["uid"] == player1_uid)
         player2_elo = next(p["current_elo"] for p in players if p["uid"] == player2_uid)
 
-        elo_change1, elo_change2 = calculate_elo(
+        elo_change1, elo_change2 = calculate_elo(player1_elo, player2_elo, request.player1_score, request.player2_score)
-            player1_elo, player2_elo, request.player1_score, request.player2_score
-        )
 
         cursor.execute(
             """
@@ -256,18 +262,28 @@ def end_match(request: EndMatchRequest):
             """,
             (request.player1_score, request.player2_score, elo_change1, elo_change2, request.match_id)
         )
-        new_player1_elo = player1_elo + elo_change1
+
-        new_player2_elo = player2_elo + elo_change2
+        new_player1_elo = max(0, player1_elo + elo_change1)
+        new_player2_elo = max(0, player2_elo + elo_change2)
         cursor.execute(
             """
             UPDATE users
-            SET current_elo = CASE WHEN uid = %s THEN %s WHEN uid = %s THEN %s END
+            SET current_elo = CASE
+                WHEN uid = %s THEN %s
+                WHEN uid = %s THEN %s
+            END
             WHERE uid IN (%s, %s);
             """,
             (player1_uid, new_player1_elo, player2_uid, new_player2_elo, player1_uid, player2_uid)
         )
 
         conn.commit()
+
+        if request.match_id in active_connections:
+            connections = active_connections.pop(request.match_id, [])
+            for websocket in connections:
+                await websocket.close()
+
         return {"message": "Match ended successfully"}
     except Exception as e:
         conn.rollback()
@@ -277,6 +293,40 @@ def end_match(request: EndMatchRequest):
         conn.close()
 
 
+
+@app.websocket("/ws/{match_id}")
+async def websocket_endpoint(websocket: WebSocket, match_id: int):
+    await websocket.accept()
+    if match_id not in active_connections:
+        raise HTTPException(status_code=404, detail="Match not found")
+    active_connections[match_id].append(websocket)
+
+    try:
+        while True:
+            message = await websocket.receive_text()
+            data = json.loads(message)
+
+            if data["type"] == "score_update":
+                query = """
+                    UPDATE matches
+                    SET player1_score = :player1_score, player2_score = :player2_score
+                    WHERE match_id = :match_id;
+                """
+                await database.execute(query, {
+                    "player1_score": data["player1_score"],
+                    "player2_score": data["player2_score"],
+                    "match_id": match_id
+                })
+
+            for connection in active_connections[match_id]:
+                if connection != websocket:
+                    await connection.send_text(message)
+    except WebSocketDisconnect:
+        active_connections[match_id].remove(websocket)
+        if not active_connections[match_id]:
+            del active_connections[match_id]
+
+
 @app.post("/getprofile")
 def get_profile(request: ProfileRequest):
     conn = get_db_connection()