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

import asyncio
from typing import Dict, Set

from fastapi import WebSocket

# This will be set by main.py during startup
MAIN_LOOP = None

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


def publish_progress(task_id: str, progress: int, message: str, status: str = "processing"):
    """
    Publish progress update to all connected WebSocket clients for a task.
    """
    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()
    # Get the loop to use: either the stored one or try to get the current loop
    loop = MAIN_LOOP
    if loop is None:
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            # No running loop, fallback to default event loop
            loop = asyncio.get_event_loop()

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

    # Clean up disconnected clients
    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):
    """Register a WebSocket connection for a task."""
    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):
    """Unregister a WebSocket connection for a task."""
    if task_id in active_connections:
        active_connections[task_id].discard(websocket)
        if not active_connections[task_id]:
            del active_connections[task_id]