"""Async WebSocket client for /api/v2/agent/ws."""

from __future__ import annotations

import asyncio
import json
import uuid
from dataclasses import dataclass
from typing import Any, Optional
from urllib.parse import urlencode

import websockets
from websockets.asyncio.client import ClientConnection


@dataclass
class WsClientConfig:
    """Configuration for agent websocket client connections."""

    base_url: str
    user_id: str
    user_token: str
    thread_id: Optional[str] = None
    explicit_thread_id: bool = False
    user_language: Optional[str] = None
    debug: bool = False
    connect_timeout_seconds: float = 30.0
    reconnect_attempts: int = 3
    reconnect_backoff_seconds: float = 1.5


@dataclass(frozen=True)
class SessionBindingUpdate:
    """Describes how the server bound the current websocket session."""

    requested_thread_id: Optional[str]
    active_thread_id: str
    explicit_thread_id: bool
    is_new: Optional[bool]
    server_overrode_requested_thread: bool


class AgentWsV3Client:
    """Client wrapper for the agent websocket message contract."""

    def __init__(self, config: WsClientConfig) -> None:
        self.config = config
        self._ws: Optional[ClientConnection] = None
        self._recv_task: Optional[asyncio.Task[None]] = None
        self._incoming_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self._closed = False
        self._thread_id_ready = asyncio.Event()
        self._requested_thread_id = self.config.thread_id
        if self.config.thread_id:
            self._thread_id_ready.set()

    def _debug_print_raw(self, direction: str, raw: Any) -> None:
        if not self.config.debug:
            return

        if isinstance(raw, bytes):
            try:
                rendered = raw.decode("utf-8")
            except UnicodeDecodeError:
                rendered = repr(raw)
        else:
            rendered = str(raw)

        print(f"[DEBUG][{direction}] {rendered}")

    def _debug_print_session_block(self) -> None:
        if not self.config.debug:
            return
        block = {
            "endpoint": self.config.base_url,
            "user_id": self.config.user_id,
            "session_id": self.config.thread_id,
            "language": self.config.user_language,
            "explicit_thread_id": self.config.explicit_thread_id,
        }
        lines = "\n".join(f"  {k}: {v}" for k, v in block.items())
        print(f"[DEBUG][SESSION]\n{lines}")

    @property
    def thread_id(self) -> str:
        return self.config.thread_id or ""

    @property
    def requested_thread_id(self) -> Optional[str]:
        return self._requested_thread_id

    @property
    def explicit_thread_id(self) -> bool:
        return bool(self.config.explicit_thread_id)

    def set_thread_id(self, thread_id: str) -> None:
        if not thread_id:
            return
        self.config.thread_id = thread_id
        self._thread_id_ready.set()

    def apply_metadata(self, payload: dict[str, Any]) -> Optional[SessionBindingUpdate]:
        """Update client session state from a metadata message."""
        thread_id = str(payload.get("thread_id") or "").strip()
        if not thread_id:
            return None

        is_new_value = payload.get("is_new")
        is_new = is_new_value if isinstance(is_new_value, bool) else None
        requested_thread_id = self._requested_thread_id
        server_overrode = bool(
            self.config.explicit_thread_id
            and requested_thread_id
            and requested_thread_id != thread_id
        )

        self.set_thread_id(thread_id)
        return SessionBindingUpdate(
            requested_thread_id=requested_thread_id,
            active_thread_id=thread_id,
            explicit_thread_id=self.config.explicit_thread_id,
            is_new=is_new,
            server_overrode_requested_thread=server_overrode,
        )

    @property
    def incoming_queue(self) -> asyncio.Queue[dict[str, Any]]:
        return self._incoming_queue

    def _build_ws_url(self) -> str:
        query_params: dict[str, str] = {
            "userToken": self.config.user_token,
        }
        if self.config.thread_id:
            query_params["thread_id"] = self.config.thread_id
        if self.config.user_language:
            query_params["User_language"] = self.config.user_language
        base = self.config.base_url.rstrip("/")
        if base.endswith("/v3"):
            query_params["user_id"] = self.config.user_id
            return f"{base}?{urlencode(query_params)}"
        return f"{base}/{self.config.user_id}?{urlencode(query_params)}"

    async def wait_for_thread_id(self, timeout_seconds: float = 10.0) -> str:
        if self.config.thread_id:
            return self.config.thread_id
        try:
            await asyncio.wait_for(self._thread_id_ready.wait(), timeout=timeout_seconds)
        except asyncio.TimeoutError as exc:
            raise RuntimeError(
                f"Timed out waiting for thread_id from websocket metadata ({timeout_seconds}s)"
            ) from exc
        if not self.config.thread_id:
            raise RuntimeError("thread_id is still unavailable")
        return self.config.thread_id

    async def connect(self) -> None:
        """Open websocket and start background receive loop."""
        if self._ws is not None:
            return

        self._closed = False
        url = self._build_ws_url()
        attempt = 0
        last_error: Optional[Exception] = None

        while attempt < max(1, self.config.reconnect_attempts):
            attempt += 1
            try:
                self._ws = await asyncio.wait_for(
                    websockets.connect(url, ping_interval=None),
                    timeout=self.config.connect_timeout_seconds,
                )
                self._recv_task = asyncio.create_task(self._recv_loop())
                self._debug_print_session_block()
                return
            except Exception as exc:  # pragma: no cover - network dependent
                last_error = exc
                if attempt >= self.config.reconnect_attempts:
                    break
                await asyncio.sleep(self.config.reconnect_backoff_seconds * attempt)

        if last_error is not None:
            raise last_error
        raise RuntimeError("Failed to connect to websocket endpoint")

    async def reconnect(
        self, *, new_thread: bool = False, thread_id: Optional[str] = None
    ) -> None:
        """Reconnect with optional new or explicit thread_id."""
        await self.close(send_terminate=False)
        if thread_id is not None:
            self.config.thread_id = thread_id
            self.config.explicit_thread_id = True
            self._requested_thread_id = thread_id
            self._thread_id_ready.set()
        elif new_thread:
            self.config.thread_id = str(uuid.uuid4())
            self.config.explicit_thread_id = False
            self._requested_thread_id = self.config.thread_id
            self._thread_id_ready.set()
        await self.connect()

    async def close(self, *, send_terminate: bool = True) -> None:
        """Close websocket and receiver task."""
        if self._closed:
            return
        self._closed = True

        if send_terminate and self._ws is not None:
            try:
                await self.send_terminate()
            except Exception:
                pass

        if self._recv_task is not None and not self._recv_task.done():
            self._recv_task.cancel()
            try:
                await self._recv_task
            except asyncio.CancelledError:
                pass
            except Exception:
                pass

        if self._ws is not None:
            await self._ws.close()
            self._ws = None

    async def send_raw(self, payload: dict[str, Any]) -> None:
        if self._ws is None:
            raise RuntimeError("WebSocket is not connected")
        raw_text = json.dumps(payload, ensure_ascii=False)
        self._debug_print_raw("SEND", raw_text)
        await self._ws.send(raw_text)

    async def send_query(
        self,
        query: str,
        *,
        attachments: Optional[list[dict[str, Any]]] = None,
        scenario: Optional[str] = None,
        config: Optional[dict[str, Any]] = None,
        system_prompt: Optional[str] = None,
        language: Optional[str] = None,
    ) -> None:
        """Send query payload expected by agent websocket."""
        payload: dict[str, Any] = {
            "type": "query",
            "query": query,
            "user_id": self.config.user_id,
            "attachments": attachments or [],
        }
        if scenario:
            payload["scenario"] = scenario
        if config is not None:
            payload["config"] = config
        if system_prompt:
            payload["system_prompt"] = system_prompt
        if language:
            payload["language"] = language
        await self.send_raw(payload)

    async def send_interrupt(self) -> None:
        await self.send_raw({"type": "interrupt"})

    async def send_terminate(self) -> None:
        await self.send_raw({"type": "terminate"})

    async def send_answer(
        self,
        *,
        message_id: str,
        answers: dict[str, str],
        as_type: str = "ask_message",
    ) -> None:
        await self.send_raw(
            {
                "type": as_type,
                "message_id": message_id,
                "answers": answers,
            }
        )

    async def _recv_loop(self) -> None:
        """Continuously read websocket frames and enqueue JSON messages."""
        assert self._ws is not None
        try:
            async for raw in self._ws:
                self._debug_print_raw("RECV", raw)
                try:
                    payload = json.loads(raw)
                except json.JSONDecodeError:
                    payload = {"type": "_raw", "raw": raw}

                # Keep-alive frames should not spam the main render loop.
                if payload.get("type") == "heartbeat":
                    continue

                await self._incoming_queue.put(payload)
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            await self._incoming_queue.put(
                {
                    "type": "_connection_error",
                    "error_message": str(exc),
                }
            )
