From 4f9fdb911eaa06ff614bb295e75c08768e4cfd30 Mon Sep 17 00:00:00 2001
From: erik <erik@mosswart.snakedesert.se>
Date: Thu, 1 May 2025 19:21:34 +0000
Subject: [PATCH] Websocket test

---
 main.py | 40 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index 0892dcd8..fdb48928 100644
--- a/main.py
+++ b/main.py
@@ -3,7 +3,7 @@ import json
 import sqlite3
 from typing import Dict
 
-from fastapi import FastAPI, Header, HTTPException, Query
+from fastapi import FastAPI, Header, HTTPException, Query, WebSocket, WebSocketDisconnect
 from fastapi.responses import JSONResponse
 from fastapi.routing import APIRoute
 from fastapi.staticfiles import StaticFiles
@@ -11,6 +11,8 @@ from pydantic import BaseModel
 from typing import Optional
 
 from db import init_db, save_snapshot, DB_FILE
+import asyncio
+from starlette.concurrency import run_in_threadpool
 
 # ------------------------------------------------------------------
 app = FastAPI()
@@ -194,6 +196,42 @@ def get_trails(
     ]
     return JSONResponse(content={"trails": trails})
 
+# --------------------  WebSocket endpoints -----------------------
+browser_conns: set[WebSocket] = set()
+
+async def _broadcast_to_browser_clients(snapshot: dict):
+    for ws in list(browser_conns):
+        try:
+            await ws.send_json(snapshot)
+        except WebSocketDisconnect:
+            browser_conns.remove(ws)
+
+@app.websocket("/ws/position")
+async def ws_receive_snapshots(websocket: WebSocket, secret: str = Query(...)):
+    await websocket.accept()
+    if secret != SHARED_SECRET:
+        await websocket.close(code=1008)
+        return
+    try:
+        while True:
+            data = await websocket.receive_json()
+            snap = TelemetrySnapshot.parse_obj(data)
+            live_snapshots[snap.character_name] = snap.dict()
+            await run_in_threadpool(save_snapshot, snap.dict())
+            await _broadcast_to_browser_clients(snap.dict())
+    except WebSocketDisconnect:
+        pass
+
+@app.websocket("/ws/live")
+async def ws_live_updates(websocket: WebSocket):
+    await websocket.accept()
+    browser_conns.add(websocket)
+    try:
+        while True:
+            await asyncio.sleep(3600)
+    except WebSocketDisconnect:
+        browser_conns.remove(websocket)
+
 
 # --------------------  static frontend  ---------------------------
 app.mount("/", StaticFiles(directory="static", html=True), name="static")