From 70f9d24b73432f8fe69454607c1e4b3ac8bbf29f Mon Sep 17 00:00:00 2001
From: Scott Idem <stidem@gmail.com>
Date: Wed, 2 Nov 2022 18:25:50 -0400
Subject: [PATCH] Update to websockets

---
 app/routers/websockets.py | 126 +++++++++++++++++++++++++++++++++++---
 1 file changed, 116 insertions(+), 10 deletions(-)

diff --git a/app/routers/websockets.py b/app/routers/websockets.py
index 8a2feef..99f2528 100644
--- a/app/routers/websockets.py
+++ b/app/routers/websockets.py
@@ -30,7 +30,6 @@ html = """
             var client_id = Date.now();
             document.querySelector("#ws-id").textContent = client_id;
             var ws = new WebSocket(`ws://localhost:5005/ws/${client_id}`);
-            // var ws = new WebSocket(`ws://localhost:5005/ws`);
             // var ws = new WebSocket("ws://localhost:8000/ws");
             // var ws = new WebSocket("ws://fastapi.localhost/ws");
             ws.onmessage = function(event) {
@@ -47,7 +46,6 @@ html = """
                 var data_json_str = JSON.stringify(data);
                 ws.send(data_json_str);
 
-                //ws.send(input.value);
                 input.value = '';
                 event.preventDefault();
             }
@@ -64,19 +62,127 @@ async def get(response: Response = Response):
     return HTMLResponse(html)
 
 
-@router.websocket("/ws/{client_id}")
-async def websocket_endpoint(
+
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: List[WebSocket] = []
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections.append(websocket)
+        log.debug(self.active_connections)
+
+    def disconnect(self, websocket: WebSocket):
+        self.active_connections.remove(websocket)
+
+    async def send_personal_message(self, message: str, websocket: WebSocket):
+        log.setLevel(logging.DEBUG)
+        # log.debug(dir(websocket))
+        log.debug(vars(websocket))
+        log.debug(websocket.url)
+        log.debug(websocket.client)
+        log.debug(websocket.client_state)
+        log.debug(websocket.headers['sec-websocket-key'])
+        await websocket.send_text(message)
+
+    async def broadcast(self, message: str):
+        log.setLevel(logging.DEBUG)
+        for connection in self.active_connections:
+            log.debug(connection)
+            await connection.send_text(message)
+
+
+manager = ConnectionManager()
+
+
+
+
+@router.websocket('/ws/room/{room_id}')
+async def ws_room_id(
+        websocket: WebSocket,
+        room_id: str,
+        ):
+    await manager.connect(websocket)
+    await manager.broadcast(f'Welcome to room "{room_id}"!')
+    try:
+        while True:
+            data = await websocket.receive_json()
+            log.debug(data)
+            data_dict = data
+            # data_dict = json.loads(data)
+            log.debug(data_dict['client_id'])
+            client_id = data_dict['client_id']
+            await manager.send_personal_message(f'You wrote: {data}', websocket)
+            await manager.broadcast(f'Client #{client_id} says: {data}')
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        await manager.broadcast(f'Client #{client_id} left the chat')
+
+
+    # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+    # time.sleep(3.5) # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+    # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+
+
+
+
+@router.websocket('/ws/looping')
+async def ws_looping(
+        websocket: WebSocket,
+        ):
+    await manager.connect(websocket)
+    # await manager.broadcast(f'Welcome to looping')
+    try:
+        while True:
+            # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+            # await time.sleep(3.5) # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+            # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+            # data = await websocket.receive_json()
+            # log.debug(data)
+            # data_dict = data
+            # data_dict = json.loads(data)
+            # log.debug(data_dict['client_id'])
+            # await manager.send_personal_message(f'You wrote: {data}', websocket)
+            await manager.broadcast(f'Loop!!!')
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        await manager.broadcast(f'Client left')
+
+
+
+@router.websocket('/ws/{client_id}')
+async def ws_client_id(
         websocket: WebSocket,
         client_id: int,
-        response: Response = Response,
         ):
-    log.setLevel(logging.DEBUG)
-    log.debug(locals())
+    await manager.connect(websocket)
+    try:
+        while True:
+            data = await websocket.receive_json()
+            log.debug(data)
+            data_dict = data
+            # data_dict = json.loads(data)
+            log.debug(data_dict['client_id'])
+            await manager.send_personal_message(f'You wrote: {data}', websocket)
+            await manager.broadcast(f'Client #{client_id} says: {data}')
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        await manager.broadcast(f'Client #{client_id} left the chat')
 
-    log.info('Root of ws. Waiting to accept a websocket and then the redis_connector')
 
-    await websocket.accept()
-    await redis_connector(websocket)
+# @router.websocket("/ws/{client_id}")
+# async def websocket_endpoint(
+#         websocket: WebSocket,
+#         client_id: int,
+#         response: Response = Response,
+#         ):
+#     log.setLevel(logging.DEBUG)
+#     log.debug(locals())
+
+#     log.info('Root of ws. Waiting to accept a websocket and then the redis_connector')
+
+#     await websocket.accept()
+#     await redis_connector(websocket)
 
 
 async def redis_connector(