"""WebSocket progress publisher – decouples task manager from WebSocket."""

import json
from typing import Dict, Set

from fastapi import WebSocket

# Active WebSocket connections
active_connections: Dict[str, Set[WebSocket]] = {}


def publish_progress(task_id: str, progress: int, message: str, status: str = "processing"):
    if task_id not in active_connections:
        return

    data = {
        "type": "progress",
        "data": {
            "task_id": task_id,
            "status": status,
            "progress": progress,
            "message": message
        }
    }

    to_remove = set()
    for websocket in active_connections.get(task_id, set()):
        try:
            websocket.send_json(data)
        except Exception:
            to_remove.add(websocket)

    for websocket in to_remove:
        active_connections[task_id].discard(websocket)
    if task_id in active_connections and not active_connections[task_id]:
        del active_connections[task_id]


def register_connection(task_id: str, websocket: WebSocket):
    if task_id not in active_connections:
        active_connections[task_id] = set()
    active_connections[task_id].add(websocket)


def unregister_connection(task_id: str, websocket: WebSocket):
    if task_id in active_connections:
        active_connections[task_id].discard(websocket)
        if not active_connections[task_id]:
            del active_connections[task_id]