"""
WebSocket connection manager and related runtime helpers.

Extracted from main.py to keep individual modules manageable.
"""

import asyncio
from contextvars import ContextVar

from fastapi import WebSocket


# ── Per-run context ───────────────────────────────────────────────────────────

_current_run_id: ContextVar[str] = ContextVar("_current_run_id", default="")

MAX_CONCURRENT_RUNS = 3


# ── Connection manager ────────────────────────────────────────────────────────

class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []
        self.active_runs: dict[str, bool] = {}
        self.active_run_tasks: dict[str, asyncio.Task] = {}
        self.pending_approvals: dict[str, dict] = {}
        self.pending_clarification: dict | None = None
        self.token_stats: dict = {"total": 0, "by_agent": {}}

    @property
    def is_running(self) -> bool:
        return self.active_run_count() > 0

    def clear_run(self, run_id: str):
        self.active_runs.pop(run_id, None)
        self.active_run_tasks.pop(run_id, None)
        self.pending_approvals.pop(run_id, None)

    def active_run_count(self) -> int:
        count = 0
        for run_id, state in list(self.active_runs.items()):
            task = self.active_run_tasks.get(run_id)
            if task is None:
                if state == 'paused':
                    count += 1
                else:
                    self.clear_run(run_id)
                continue
            if task.done():
                self.clear_run(run_id)
                continue
            if state is True or state == 'paused':
                count += 1
        return count

    def register_run_task(self, run_id: str, task: asyncio.Task, websocket: WebSocket):
        self.active_run_tasks[run_id] = task

        def _done(done_task: asyncio.Task):
            error = None
            try:
                error = done_task.exception()
            except asyncio.CancelledError:
                error = None
            except Exception as exc:
                error = exc

            self.clear_run(run_id)

            if error is None:
                return

            print(f"[RUN-ERROR] {run_id}: {error}")
            try:
                asyncio.create_task(self.send_message({
                    "type": "run_stopped",
                    "reason_code": "INTERNAL_ERROR",
                    "content": "Workflow stopped because of an internal error.",
                    "run_id": run_id,
                }, websocket))
                asyncio.create_task(self.send_message({
                    "type": "error",
                    "content": f"Workflow crashed unexpectedly: {error}",
                    "run_id": run_id,
                }, websocket))
            except RuntimeError:
                pass

        task.add_done_callback(_done)

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def send_message(self, message: dict, websocket: WebSocket):
        rid = _current_run_id.get("")
        if rid:
            message["run_id"] = rid
        try:
            await websocket.send_json(message)
        except Exception as exc:
            print(f"Error sending message: {exc}")

    async def broadcast(self, message: dict):
        for ws in list(self.active_connections):
            try:
                await ws.send_json(message)
            except Exception:
                pass


manager = ConnectionManager()


def _browser_update_callback(state: dict):
    """Push browser state to all connected WebSocket clients."""
    msg = {"type": "browser_update", **state}
    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            asyncio.run_coroutine_threadsafe(manager.broadcast(msg), loop)
        else:
            loop.run_until_complete(manager.broadcast(msg))
    except Exception:
        pass


async def _reset_runtime_state(close_ws: bool = True):
    for run_id in list(manager.active_runs):
        manager.active_runs[run_id] = False
    manager.pending_approvals = {}
    manager.pending_clarification = None
    manager.token_stats = {"total": 0, "by_agent": {}}
    if close_ws:
        for ws in list(manager.active_connections):
            try:
                await ws.close(code=1001)
            except Exception:
                pass
        manager.active_connections.clear()
