"""WebSocket client for ARBI real-time notifications.

Depends on the ``websockets`` library.  All incoming messages are parsed into
the autogenerated typed models (``TaskUpdateMessage``, ``AuthResultMessage``,
etc.) so that consumers get full type safety.
"""

from __future__ import annotations

import asyncio
import json
import time
from collections.abc import Callable
from typing import Any, TypeAlias

from arbi_client.models.auth_result_message import AuthResultMessage
from arbi_client.models.batch_complete_message import BatchCompleteMessage
from arbi_client.models.connection_closed_message import ConnectionClosedMessage
from arbi_client.models.error_message import ErrorMessage
from arbi_client.models.notification_response import NotificationResponse
from arbi_client.models.presence_update_message import PresenceUpdateMessage
from arbi_client.models.response_complete_message import ResponseCompleteMessage
from arbi_client.models.task_update_message import TaskUpdateMessage
from arbi_client.models.task_update_message_status import TaskUpdateMessageStatus

#: Union of all known server→client WebSocket message types.
WsMessage: TypeAlias = (
    AuthResultMessage
    | TaskUpdateMessage
    | BatchCompleteMessage
    | ErrorMessage
    | ConnectionClosedMessage
    | PresenceUpdateMessage
    | ResponseCompleteMessage
    | NotificationResponse
)

#: Mapping of ``type`` field value → model class for dispatching.
_TYPE_DISPATCH: dict[str, type] = {
    "auth_result": AuthResultMessage,
    "task_update": TaskUpdateMessage,
    "batch_complete": BatchCompleteMessage,
    "error": ErrorMessage,
    "connection_closed": ConnectionClosedMessage,
    "presence_update": PresenceUpdateMessage,
    "response_complete": ResponseCompleteMessage,
}

#: Terminal statuses for document processing.
TERMINAL_STATUSES: frozenset[TaskUpdateMessageStatus] = frozenset(
    {
        TaskUpdateMessageStatus.COMPLETED,
        TaskUpdateMessageStatus.FAILED,
        TaskUpdateMessageStatus.SKIPPED,
    }
)


def parse_ws_message(data: dict[str, Any]) -> WsMessage:
    """Parse a raw JSON dict into the appropriate typed WS message model.

    Dispatches on the ``type`` field.  ``NotificationResponse`` messages don't
    carry a fixed ``type`` literal — they use ``NotificationType`` enum values
    (e.g. ``"user_message"``, ``"contact_accepted"``).  These are detected by
    the presence of ``external_id`` + ``sender`` fields.

    Falls back to ``NotificationResponse.from_dict`` for notification-shaped
    payloads.  Raises :class:`ValueError` for truly unknown message shapes.
    """
    msg_type = data.get("type")

    # Direct dispatch for messages with a known literal type.
    model_cls = _TYPE_DISPATCH.get(msg_type) if msg_type else None
    if model_cls is not None:
        return model_cls.from_dict(data)

    # NotificationResponse — identified by structure rather than a fixed type.
    if "external_id" in data and "sender" in data:
        return NotificationResponse.from_dict(data)

    raise ValueError(f"Unknown WebSocket message type: {msg_type!r}")


def print_ws_status(msg: TaskUpdateMessage) -> None:
    """Print a formatted WebSocket task_update for diagnostics (visible with ``pytest -s``).

    Shows elapsed time, document ID, filename, status, and progress bar.

    Example output::

          2.3s [doc-abc123] ARBI_Cover_Page.pdf  parsing  [####------] 40%
          8.7s [doc-abc123] ARBI_Cover_Page.pdf  completed [##########] 100%
    """
    elapsed = msg.additional_properties.get("elapsed")
    ts = f"{elapsed:>7.1f}s" if elapsed is not None else "       "

    status_str = msg.status.value
    pct = msg.progress

    # Progress bar: 10-char block
    filled = pct // 10
    bar = "#" * filled + "-" * (10 - filled)

    # Colour-code status for terminal readability
    status_display = status_str.upper() if status_str in ("completed", "failed", "skipped") else status_str

    print(f"  {ts} [{msg.doc_ext_id}] {msg.file_name:<40s} {status_display:<12s} [{bar}] {pct}%")


class ArbiWebSocket:
    """Async WebSocket client for ARBI notifications.

    All received messages are parsed into typed models from
    ``arbi_client.models``.

    Usage::

        async with ArbiWebSocket(base_url="https://app.arbi.city", access_token=token) as ws:
            result = await ws.wait_for_docs(doc_ext_ids, timeout=120)
    """

    def __init__(self, base_url: str, access_token: str) -> None:
        self._base_url = base_url.rstrip("/")
        self._access_token = access_token
        self._ws: Any = None

    # ------------------------------------------------------------------
    # Connection lifecycle
    # ------------------------------------------------------------------

    async def connect(self) -> None:
        """Open the WebSocket connection and authenticate."""
        import websockets

        ws_base = self._base_url.replace("https://", "wss://").replace("http://", "ws://")
        ws_url = f"{ws_base}/v1/notifications/ws"

        self._ws = await websockets.connect(ws_url)

        # Authenticate
        await self._ws.send(json.dumps({"type": "auth", "token": self._access_token}))
        raw = json.loads(await self._ws.recv())
        auth_resp = AuthResultMessage.from_dict(raw)
        if not auth_resp.success:
            reason = auth_resp.reason if hasattr(auth_resp, "reason") else "unknown"
            await self._ws.close()
            self._ws = None
            raise RuntimeError(f"WebSocket auth failed: {reason}")

    async def close(self) -> None:
        """Close the WebSocket connection."""
        if self._ws is not None:
            await self._ws.close()
            self._ws = None

    # ------------------------------------------------------------------
    # Messaging
    # ------------------------------------------------------------------

    async def recv(self) -> WsMessage:
        """Receive, JSON-decode, and parse the next message into a typed model."""
        if self._ws is None:
            raise RuntimeError("WebSocket not connected")
        raw = json.loads(await self._ws.recv())
        return parse_ws_message(raw)

    async def send(self, message: dict) -> None:
        """JSON-encode and send a message."""
        if self._ws is None:
            raise RuntimeError("WebSocket not connected")
        await self._ws.send(json.dumps(message))

    # ------------------------------------------------------------------
    # High-level helpers
    # ------------------------------------------------------------------

    async def wait_for_docs(
        self,
        doc_ext_ids: list[str],
        target_statuses: set[TaskUpdateMessageStatus] | None = None,
        timeout: int = 120,
        on_status: Callable[[TaskUpdateMessage], Any] | None = None,
    ) -> dict[str, TaskUpdateMessageStatus]:
        """Wait for documents to reach target status via WebSocket notifications.

        Args:
            doc_ext_ids: Document external IDs to wait for.
            target_statuses: Statuses to wait for (default: completed, failed,
                skipped).  Use :class:`TaskUpdateMessageStatus` enum values.
            timeout: Maximum seconds to wait.
            on_status: Optional callback invoked with a :class:`TaskUpdateMessage`
                for **every** ``task_update`` matching a tracked document, not
                just terminal ones.  The message is enriched with an ``elapsed``
                key in ``additional_properties`` (seconds since tracking started).

                Use this for progress reporting, logging, or UI updates.  The
                typed model gives you ``msg.status``, ``msg.progress``,
                ``msg.file_name``, ``msg.doc_ext_id`` directly.

        Returns:
            Dict mapping ``doc_ext_id`` -> final :class:`TaskUpdateMessageStatus`.

        Raises:
            TimeoutError: If timeout is reached.
        """
        if self._ws is None:
            raise RuntimeError("WebSocket not connected")

        if target_statuses is None:
            target_statuses = TERMINAL_STATUSES

        pending = set(doc_ext_ids)
        tracked = set(doc_ext_ids)
        results: dict[str, TaskUpdateMessageStatus] = {}
        t0 = time.monotonic()

        async def _listen():
            while pending:
                raw = json.loads(await self._ws.recv())
                if raw.get("type") != "task_update":
                    continue
                msg = TaskUpdateMessage.from_dict(raw)
                if msg.doc_ext_id not in tracked:
                    continue

                msg["elapsed"] = round(time.monotonic() - t0, 2)

                if on_status is not None:
                    on_status(msg)

                if msg.doc_ext_id in pending and msg.status in target_statuses:
                    results[msg.doc_ext_id] = msg.status
                    pending.discard(msg.doc_ext_id)

        await asyncio.wait_for(_listen(), timeout=timeout)
        return results

    async def listen(
        self,
        on_message: Callable[[WsMessage], Any],
        timeout: float | None = None,
    ) -> None:
        """Listen for messages, calling ``on_message`` for each.

        Args:
            on_message: Callback receiving the parsed typed message model.
            timeout: Optional timeout in seconds (``None`` for indefinite).
        """
        if self._ws is None:
            raise RuntimeError("WebSocket not connected")

        async def _loop():
            async for raw_bytes in self._ws:
                msg = parse_ws_message(json.loads(raw_bytes))
                on_message(msg)

        if timeout is not None:
            await asyncio.wait_for(_loop(), timeout=timeout)
        else:
            await _loop()

    # ------------------------------------------------------------------
    # Context manager
    # ------------------------------------------------------------------

    async def __aenter__(self) -> ArbiWebSocket:
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.close()
