"""WebSocket delivery manager: server-issued ws_id, terminal replay, delivery fan-out.

Uses SubscriberRegistry with SubscriptionCategory.WS_DELIVERY only.
Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com
"""

from __future__ import annotations

import asyncio
import copy
import uuid
from typing import Optional

from mcp_proxy_adapter.core.delivery.terminal_replay_store import TerminalReplayStore
from mcp_proxy_adapter.core.job_push.subscriptions import (
    Subscriber,
    SubscriberRegistry,
    SubscriptionCategory,
)
from mcp_proxy_adapter.core.logging import get_global_logger


class WsDeliveryManager:
    """Orchestrate ws_id issuance, replay, and fan-out to WS_DELIVERY subscribers."""

    def __init__(
        self,
        registry: SubscriberRegistry,
        replay_store: Optional[TerminalReplayStore] = None,
    ) -> None:
        self._registry = registry
        self._replay = replay_store or TerminalReplayStore()
        self._logger = get_global_logger()
        self._lock = asyncio.Lock()
        self._job_to_ws: dict[str, str] = {}

    def issue_ws_id(self) -> str:
        """Return a new unique ws_id (UUID4 string)."""
        return str(uuid.uuid4())

    async def notify_terminal_delivery(self, ws_id: str, payload: dict) -> None:
        """Store terminal; fan-out to WS_DELIVERY subscribers; log send failures."""
        wid = str(ws_id).strip()
        if wid == "":
            return
        await self._replay.store_terminal(wid, payload)
        subs = await self._registry.get_subscribers(
            wid, SubscriptionCategory.WS_DELIVERY
        )
        msg = copy.deepcopy(payload)
        for sub in subs:
            try:
                await sub.send(msg)
            except Exception as e:
                self._logger.warning(
                    "WsDeliveryManager: send failed for ws_id=%s: %s", wid, e
                )

    async def flush_replay_to_subscriber(
        self, ws_id: str, subscriber: Subscriber
    ) -> None:
        """If replay exists for ws_id, send one copy to subscriber."""
        wid = str(ws_id).strip()
        if wid == "":
            return
        terminal = await self._replay.get_terminal_copy(wid)
        if terminal is None:
            return
        try:
            await subscriber.send(copy.deepcopy(terminal))
        except Exception as e:
            self._logger.warning(
                "WsDeliveryManager: send failed for ws_id=%s: %s", wid, e
            )

    async def register_queue_job_mapping(self, job_id: str, ws_id: str) -> None:
        """Register job_id -> ws_id for queued terminal handoff. Last-writer wins."""
        jid = str(job_id).strip()
        wid = str(ws_id).strip()
        if jid == "" or wid == "":
            return
        async with self._lock:
            self._job_to_ws[jid] = wid

    async def resolve_ws_id_for_job(self, job_id: str) -> Optional[str]:
        """Return ws_id for queue job_id if registered; else None."""
        jid = str(job_id).strip()
        if jid == "":
            return None
        async with self._lock:
            return self._job_to_ws.get(jid)

    async def clear_queue_job_mapping(self, job_id: str) -> None:
        """Remove job_id from mapping after terminal delivery."""
        jid = str(job_id).strip()
        if jid == "":
            return
        async with self._lock:
            self._job_to_ws.pop(jid, None)
