import json
from typing import Any

from fastapi import WebSocket


class ConnectionManager:
    def __init__(self):
        self._connections: dict[str, list[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, session_id: str) -> None:
        await websocket.accept()
        if session_id not in self._connections:
            self._connections[session_id] = []
        self._connections[session_id].append(websocket)

    def disconnect(self, websocket: WebSocket, session_id: str) -> None:
        if session_id in self._connections:
            self._connections[session_id].remove(websocket)
            if not self._connections[session_id]:
                del self._connections[session_id]

    async def send_message(self, session_id: str, message: dict[str, Any]) -> None:
        if session_id in self._connections:
            for connection in self._connections[session_id]:
                await connection.send_json(message)

    async def broadcast(self, message: dict[str, Any]) -> None:
        for connections in self._connections.values():
            for connection in connections:
                await connection.send_json(message)


manager = ConnectionManager()


async def send_progress(
    session_id: str,
    current_agent: str | None,
    status: str,
    progress: float,
    message: str,
) -> None:
    await manager.send_message(
        session_id,
        {
            "type": "progress",
            "session_id": session_id,
            "data": {
                "current_agent": current_agent,
                "status": status,
                "progress": progress,
                "message": message,
            },
        },
    )


async def send_result(session_id: str, result: dict[str, Any]) -> None:
    await manager.send_message(
        session_id,
        {
            "type": "result",
            "session_id": session_id,
            "data": result,
        },
    )


async def send_error(session_id: str, error: str) -> None:
    await manager.send_message(
        session_id,
        {
            "type": "error",
            "session_id": session_id,
            "data": {"error": error},
        },
    )