from typing import Union
from fastapi.websockets import WebSocket, WebSocketState

class Accelerator:
  ws: Union[WebSocket, None]

  def connected(self): return self.ws != None

  async def connect(self, ws: WebSocket):
    await ws.accept()
    self.ws = ws

  async def accelerate(self, input):
    while ws.client_state == WebSocketState.CONNECTING: pass
    
    if ws.client_state == WebSocketState.CONNECTED:
      await self.ws.send_text(input)
      return await self.ws.receive_text()
    elif ws.client_state == WebSocketState.DISCONNECTED:
      ws = None