Update main.py

main.py

@@ -1,6 +1,5 @@
-from fastapi import FastAPI, WebSocket
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse
+from fastapi import FastAPI, Request, WebSocket
+from fastapi.templating import Jinja2Templates
 import subprocess
 import asyncio
 import os
@@ -10,6 +9,12 @@ from queue import Queue, Empty
 
 os.environ['LD_LIBRARY_PATH'] = '/engines/maia/lib:' + os.environ.get('LD_LIBRARY_PATH', '')
 
+user_playing = {
+    'total': 0,
+    'stockfish': 0,
+    'maia': 0
+}
+
 def enqueue_output(out, queue):
     for line in iter(out.readline, b''):
         queue.put(line)
@@ -29,9 +34,9 @@ class EngineChess:
         self.thread = Thread(target=enqueue_output, args=(self._stockfish.stdout, self.queueOutput))
         self.thread.daemon = True # thread dies with the program
         self.thread.start()
-        
+
         self._has_quit_command_been_sent = False
-        self._debug_view = True
+        self._debug_view = False
 
     def _put(self, command):
         if not self._stockfish.stdin:
@@ -41,150 +46,93 @@ class EngineChess:
             self._stockfish.stdin.flush()
             if command == "quit":
                 self._has_quit_command_been_sent = True
-    
+
     def _read_line(self) -> str:
         if not self._stockfish.stdout:
             raise BrokenPipeError()
         if self._stockfish.poll() is not None:
             raise StockfishException("The Stockfish process has crashed")
-            
+
         try:
             line = self.queueOutput.get_nowait() # or q.get(timeout=.1)
         except Empty:
             return ""
-        
+
         if self._debug_view:
             print(f"Engine {self.path_engine[0]}:", line.strip())
         return line.strip()
-    
+
     def _is_ready(self) -> None:
         self._put("isready")
         while self._read_line() != "readyok":
             pass
-    
+
     def put(self, cmd):
         return self._put(cmd)
-        
+
     def read_line(self) -> str:
         return self._read_line()
 
 app = FastAPI()
+templates = Jinja2Templates(directory="templates")
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
+# WebSocket endpoints
 @app.websocket("/stockfish-{version}")
 async def websocket_endpoint(websocket: WebSocket, version: str):
     await websocket.accept()
-    
-    stockfish = EngineChess([f"./engines/stockfish/stockfish-{version}-uci"])
+
+    user_playing['total'] += 1
+    user_playing['stockfish'] += 1
 
     async def read_from_socket(websocket: WebSocket):
         async for data in websocket.iter_text():
             print(f"Stockfish Client: {data}")
             stockfish.put(data)
-                    
+
     asyncio.create_task(read_from_socket(websocket))
 
-    while True:
+    try:
         while True:
             res = stockfish.read_line()
             if res:
                 await websocket.send_text(f"{res}")
             else:
                 break
-        await asyncio.sleep(0.1)
+    finally:
+        user_playing['total'] -= 1
+        user_playing['stockfish'] -= 1
 
 @app.websocket("/maia-{elo}")
 async def websocket_endpoint(websocket: WebSocket, elo: str):
     await websocket.accept()
-    
-    stockfish = EngineChess([f"./engines/maia/maia-{elo}/lc0", "--backend=trivial"])
+
+    user_playing['total'] += 1
+    user_playing['maia'] += 1
 
     async def read_from_socket(websocket: WebSocket):
         async for data in websocket.iter_text():
             print(f"Maia Client: {data}")
             stockfish.put(data)
-                    
+
     asyncio.create_task(read_from_socket(websocket))
 
-    while True:
+    try:
         while True:
             res = stockfish.read_line()
             if res:
                 await websocket.send_text(f"{res}")
             else:
                 break
-        await asyncio.sleep(0.1)
-
-
-
-
-
-########################
-async def get_prompt():
-    # Get current working directory
-    cwd = os.getcwd()
-    # Shorten path if it is too long
-    if len(cwd) > 30:
-        cwd = os.path.basename(cwd)
-        cwd = f".../{cwd}"
-    return f"{cwd}$ "
-
-async def exec_command(command):
-    try:
-        # Execute command using subprocess
-        process = await asyncio.create_subprocess_shell(
-            command,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            shell=True,
-            executable='/bin/bash'
-        )
-
-        # Wait for command to complete
-        stdout, stderr = await process.communicate()
-
-        # Prepare output to send back
-        output = stdout.decode() + stderr.decode()
-
-    except Exception as e:
-        output = str(e)
-
-    return output
-
-async def send_prompt(websocket):
-    # Send initial prompt to client
-    prompt = await get_prompt()
-    await websocket.send_text(prompt)
-
-@app.websocket("/shell")
-async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
-    await send_prompt(websocket)
-
-    try:
-        while True:
-            # Receive command from client
-            command = await websocket.receive_text()
-
-            if command == "exit":
-                await websocket.send_text("Goodbye!")
-                break
-
-            # Execute command
-            output = await exec_command(command)
-
-            # Send output to client
-            await websocket.send_text(output)
-
-            # Send updated prompt to client
-            await send_prompt(websocket)
-
-    except websockets.exceptions.ConnectionClosedError:
-        pass
+    finally:
+        user_playing['total'] -= 1
+        user_playing['maia'] -= 1
+
+
+@app.get("/")
+async def root(request: Request):
+    return templates.TemplateResponse("index.html", {
+        "request": request,
+        "total_users": user_playing['total'],
+        "stockfish_users": user_playing['stockfish'],
+        "maia_users": user_playing['maia']
+    })