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

class Accelerator:
  def __del__(self): self.ws.close()

  ws: Union[WebSocket, None] = None

  def connected(self):
    if self.ws == None: return False
    
    print(f"===== client_state: {self.ws.client_state}")

    if self.ws.client_state == WebSocketState.CONNECTED:
      return True
    else:
      self.ws = None
      return False

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

  async def accelerate(self, input):    
    await self.ws.send_text(input)
    return await self.ws.receive_text()