import asyncio
import json
import websockets
import logging

class DashboardWS:
    def __init__(self, host="0.0.0.0", port=8001):
        self.host = host
        self.port = port
        self.clients = set()

    async def register(self, websocket):
        self.clients.add(websocket)
        logging.info(f"Dashboard client connected: {websocket.remote_address}")

    async def unregister(self, websocket):
        self.clients.remove(websocket)
        logging.info(f"Dashboard client disconnected")

    async def broadcast(self, data):
        if not self.clients:
            return
        message = json.dumps(data)
        await asyncio.gather(*[client.send(message) for client in self.clients])

    async def ws_handler(self, websocket):
        await self.register(websocket)
        try:
            async for message in websocket:
                # Handle commands if needed
                pass
        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            await self.unregister(websocket)

    async def start(self):
        async with websockets.serve(self.ws_handler, self.host, self.port):
            logging.info(f"WebSocket server started on ws://{self.host}:{self.port}")
            await asyncio.Future()  # run forever