import asyncio
import json
from typing import Any, Dict, Set
from fastapi import WebSocket

class WebSocketManager:
    def __init__(self):
        self._connections: Set[WebSocket] = set()
        self._lock = asyncio.Lock()

    async def connect(self, ws: WebSocket) -> None:
        await ws.accept()
        async with self._lock:
            self._connections.add(ws)

    async def disconnect(self, ws: WebSocket) -> None:
        async with self._lock:
            self._connections.discard(ws)

    async def broadcast(self, message: Dict[str, Any]) -> None:
        payload = json.dumps(message, ensure_ascii=False)
        async with self._lock:
            dead = set()
            for ws in self._connections:
                try:
                    await ws.send_text(payload)
                except Exception:
                    dead.add(ws)
            self._connections -= dead

    async def send_agent_state(self, agent_id: str, state: str, detail: str = "", task_id: str = None) -> None:
        msg = {"type": "agent_state", "agent": agent_id, "state": state, "detail": detail}
        if task_id:
            msg["task_id"] = task_id
        await self.broadcast(msg)

    async def send_task_complete(self, agent_id: str, task_id: str, result: dict) -> None:
        await self.broadcast({
            "type": "task_complete", "agent": agent_id,
            "task_id": task_id, "result": result,
        })

    async def send_agent_move(self, agent_id: str, target: str) -> None:
        await self.broadcast({"type": "agent_move", "agent": agent_id, "target": target})

ws_manager = WebSocketManager()