from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
from threading import Thread
import uvicorn
import asyncio
from protobuf_decoder.protobuf_decoder import Parser
from logger import logger

#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()
            try:
                while True:
                    data = await websocket.receive_bytes()
                    await self._recv_q.put(data)
            except WebSocketDisconnect:
                pass

    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, packetParse = False):
        item = await self._recv_q.get()
        if item is None:
            raise ConnectionError(self._error or "tcp transport error")
        
        if packetParse:
            to_parse = ' '.join(f'{byte:02x}' for byte in item)
            res = Parser().parse(to_parse)
            logger.info("[WS PACKET]", res.to_dict())

        return item
    
    async def send(self, payload):
        pass
    
    async def close(self):
        pass