from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
from threading import Thread
import uvicorn
import asyncio

#only rx
class WSTransport:
    def __init__(self, port, t, path = "/ws/0"):
        self.port = port
        self.token = t
        self.path = path
        self._recv_q = asyncio.Queue()
        self.thread = None

        self.app = FastAPI()
        
        @self.app.websocket(self.path)
        async def websocket_endpoint(*, websocket: WebSocket, token: str = Query("")):
            if token != self.token:
                return
            
            await websocket.accept()
            while True:
                data = await websocket.receive_bytes()
                await self._recv_q.put(data)

    async def start(self):
        if self.thread != None:
            return

        def run():
            uvicorn.run(self.app, host="0.0.0.0", port = int(self.port))

        self.thread = Thread(target=run, daemon=True)
        self.thread.start()

    async def recv(self):
        item = await self._recv_q.get()
        if item is None:
            raise ConnectionError(self._error or "tcp transport error")
        return item
    
    async def send(self, payload):
        pass
    
    async def close(self):
        pass