import json
import asyncio
import websockets
import threading
import queue

from utils.logger import get_logger
logger = get_logger(__name__)

# ========================= WebSocket 服务线程 =========================
class WebSocketSender(threading.Thread):
    def __init__(self, send_queue: queue.Queue, stop_event: threading.Event, ws_host: str, ws_port: int):
        super().__init__(daemon=True)
        self.send_queue = send_queue
        self.stop_event = stop_event
        self.ws_clients = set()
        self.ws_host = ws_host
        self.ws_port = ws_port

    async def _ws_handler(self, websocket):
        self.ws_clients.add(websocket)
        try:
            async for _ in websocket:
                pass
        finally:
            self.ws_clients.discard(websocket)

    async def _broadcaster(self):
        while not self.stop_event.is_set():
            try:
                msg = await asyncio.to_thread(self.send_queue.get, timeout=0.5)
            except queue.Empty:
                continue
            data = json.dumps(msg)
            dead = []
            for ws in list(self.ws_clients):
                try:
                    await ws.send(data)
                except:
                    dead.append(ws)
            for ws in dead:
                self.ws_clients.discard(ws)
            self.send_queue.task_done()

    async def _run_async(self):
        async with websockets.serve(self._ws_handler, self.ws_host, self.ws_port):
            logger.info(f"[INFO] WebSocket server started at ws://{self.ws_host}:{self.ws_port}")
            await self._broadcaster()

    def run(self):
        asyncio.run(self._run_async())