---

import asyncio
import logging
import time
from typing import Any, Callable, List, Optional

import msgpack
import requests
import websockets
from websockets import InvalidStatus

from constants.architecture_constants import WebSocketClientConstants as WS_Const
from utils.msgpack_utils.message_filters import is_desired_invocation_id, is_desired_type
from utils.msgpack_utils.msgpack_utils import encode_with_varint_prefix, parse_message

logger = logging.getLogger(__name__)


class WebSocketClient:
    """
    Асинхронный ws-клиент для api-gateway по протоколу Async Api
    """

    def __init__(
        self,
        host: str,
        access_token: str,
        x_user_id: str,
        reconnect_interval: float = WS_Const.DEFAULT_RECONNECT_INTERVAL,
    ):
        self._host = host
        self._access_token = access_token
        self._x_user_id = x_user_id
        self._reconnect_interval = reconnect_interval
        self._ws_url = f"wss://{host.rstrip('/')}{WS_Const.WS_HUBS}"
        self._negotiate_url = f"https://{host.rstrip('/')}{WS_Const.NEGOTIATE_ENDPOINT}"
        self._buffer = b""
        self._next_id = WS_Const.START_INVOCATION_ID

        self._ws: websockets.ClientConnection | None = None
        self.recv_queue: asyncio.Queue[Any] = asyncio.Queue()
        self._recv_task: asyncio.Task | None = None
        self._ping_task: asyncio.Task | None = None
        self._stop_event = asyncio.Event()
        self._invocation_id: Optional[str] = None
        self.suppress_recv_logging: bool = False
        self.received_messages_count: int = 0
        self.sent_ping_count: int = 0

    @property
    def invocation_id(self):
        return self._invocation_id

    @property
    def is_connected(self) -> bool:
        """
        True, если wss-соединение установлено и приём сообщений работает.

        Соединение считается живым, пока открыт сокет и фоновый приёмник
        (_recv_task) не завершился. При обрыве связи _recv_loop возвращается,
        и задача переходит в состояние done - это признак потери соединения.
        """
        if self._ws is None:
            return False
        if self._recv_task is None:
            return False
        return not self._recv_task.done()

    def clear_queue(self):
        """
        Очищает очередь путем пересоздания экземпляра класса очереди
        """
        while not self.recv_queue.empty():
            try:
                self.recv_queue.get_nowait()
                self.recv_queue.task_done()
            except asyncio.QueueEmpty as message_empty:
                logger.info(f"Очередь сообщений очищена: {message_empty}")
                break

    def update_credentials(self, access_token: str, x_user_id: str) -> None:
        """
        Обновляет креды для последующего reconnect/negotiate.

        Нужно при смене датасета: токен и x_user_id обновляются в group_state,
        а соединение переиспользуется (пересоздаётся через reconnect).
        """
        self._access_token = access_token
        self._x_user_id = x_user_id

    async def __aenter__(self):
        await self._connect_loop()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        self._stop_event.set()
        if self._ping_task:
            self._ping_task.cancel()
            try:
                await self._ping_task
            except asyncio.CancelledError:
                pass
        if self._ws:
            await self._ws.close()
        if self._recv_task:
            await self._recv_task

        # TODO: дергать ручку завершения сессии в LDS-4083

    async def _handshake(self) -> None:
        payload = WS_Const.HANDSHAKE_MESSAGE + WS_Const.RS.decode()
        logger.debug(
            f"Отправлен handshake: {payload}",
        )
        await self._ws.send(payload)

        buf = self._buffer
        finish = time.monotonic() + WS_Const.HANDSHAKE_WAITING
        while time.monotonic() < finish:
            chunk = await self._ws.recv()
            logger.info(f"Ответ на handshake: {chunk}")
            buf += chunk.encode() if isinstance(chunk, str) else chunk
            if WS_Const.RS in buf:
                return

        raise TimeoutError("Handshake timeout: не получили сообщение с разделителем RS за указанное время")

    async def negotiate(self) -> str:
        """
        SignalR-negotiate: получает connectionToken для ws-подключения.

        connectionToken подставляется в параметр id= URL вебсокета. Без negotiate
        api-gateway не регистрирует соединение как полноценного SignalR-клиента,
        поэтому он обязателен для корректного подключения.
        """
        params = {
            "token": self._access_token,
            "xUserId": self._x_user_id,
            "negotiateVersion": WS_Const.NEGOTIATE_VERSION,
        }
        response = await asyncio.to_thread(
            requests.post,
            self._negotiate_url,
            params=params,
            timeout=WS_Const.NEGOTIATE_TIMEOUT_SECONDS,
        )
        response.raise_for_status()
        data = response.json()
        connection_token = data.get("connectionToken")
        logger.info(
            "[NEGOTIATE] POST %s -> %s (connectionId=%s, connectionToken=%s)",
            WS_Const.NEGOTIATE_ENDPOINT,
            response.status_code,
            data.get("connectionId"),
            connection_token,
        )
        if not connection_token:
            raise RuntimeError("negotiate не вернул connectionToken")
        return connection_token

    async def _connect_loop(self, timeout: float = WS_Const.WS_CONNECT_TIMEOUT_SECONDS) -> None:
        """
        Цикл подключения с повторными попытками до наступления stop_event
        или истечения WS_CONNECT_TIMEOUT_SECONDS.
        """
        deadline = time.monotonic() + timeout
        attempt = 0
        transient_errors = (ConnectionError, OSError, asyncio.TimeoutError, InvalidStatus)

        while not self._stop_event.is_set():
            attempt += 1
            try:
                connection_token = await self.negotiate()
                self.ws_request = (
                    f"{self._ws_url}/?token={self._access_token}"
                    f"&xUserId={self._x_user_id}&id={connection_token}"
                )
                logger.info(f"Попытка подключения по wss: {self._ws_url}/?token=...&xUserId=...&id=...")
                self._ws = await websockets.connect(
                    self.ws_request,
                    ping_interval=WS_Const.PING_INTERVAL,
                    ping_timeout=WS_Const.PING_TIMEOUT,
                    close_timeout=WS_Const.CLOSE_TIMEOUT,
                )
                # Handshake
                await self._handshake()
                # Запускаем приём и SignalR keepalive в фоне
                self._recv_task = asyncio.create_task(self._recv_loop())
                self._ping_task = asyncio.create_task(self._ping_loop())
                logger.info("Websocket connected")
                return
            except transient_errors as exc:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise TimeoutError(
                        f"WSS hub не готов за {WS_Const.WS_CONNECT_TIMEOUT_SECONDS} с " f"(попыток: {attempt}): {exc}"
                    ) from exc
                status_info = ""
                if isinstance(exc, InvalidStatus):
                    response = getattr(exc, "response", None)
                    if response is not None:
                        status_info = f", HTTP {response.status_code}"
                logger.warning(
                    "WSS подключение не установлено (попытка %s%s): %s. " "Повтор через %s с, осталось %.0f с",
                    attempt,
                    status_info,
                    exc,
                    self._reconnect_interval,
                    remaining,
                )
                await asyncio.sleep(self._reconnect_interval)

    async def _recv_loop(self) -> None:
        """
        Прием сообщений, парсинг и отправка в очередь.
        """
        assert self._ws is not None

        while not self._stop_event.is_set():
            try:
                chunk = await self._ws.recv()
                if not self.suppress_recv_logging:
                    logger.debug(f"Сырые биты до обработки: {chunk[:100]}")
            except websockets.ConnectionClosed as e:
                logger.warning(f"WebSocket соединение разорвано: {e}")
                return

            result_message = parse_message(chunk)
            await self.recv_queue.put(result_message)
            self.received_messages_count += 1
            if not self.suppress_recv_logging:
                str_message = str(result_message)
                logger.info(
                    f"Обработанное сообщение от api-gateway: {str_message[:200]}... полное сообщение в attach",
                )
            continue

    async def send_ping(self) -> None:
        """Отправляет SignalR Ping [6] для прикладного keepalive."""
        if not self._ws:
            return
        packet = encode_with_varint_prefix(msgpack.packb([WS_Const.PING_MESSAGE_TYPE], use_bin_type=True))
        await self._ws.send(packet)
        self.sent_ping_count += 1

    async def _ping_loop(self) -> None:
        """Периодически шлёт SignalR Ping, пока соединение живо."""
        while not self._stop_event.is_set():
            await asyncio.sleep(WS_Const.SIGNALR_PING_INTERVAL_SECONDS)
            if self._stop_event.is_set() or self._ws is None:
                return
            try:
                await self.send_ping()
            except (websockets.ConnectionClosed, OSError):
                return

    async def invoke(self, target: str, args: list) -> None:
        """
        Отправляет удаленный вызов invocation о websocket соединению

        Метод формирует сообщение по протоколу SignalR, включая:
        - типа сообщения
        - заголовки
        - уникальный идентификатор запроса
        - имя целевого метода
        - аргументы запроса

        Сообщение запаковывается в messagepack и отправляется через текущее websocket соединение
        """

        if not self._ws:
            raise websockets.WebSocketException("Не установлено подключение по wss")
        self._invocation_id = str(self._next_id)
        self._next_id += 1
        invocation = [
            WS_Const.DEFAULT_SIGNALR_MESSAGE_TYPE,
            WS_Const.DEFAULT_SIGNALR_MAP_HEADERS,
            self._invocation_id,
            target,
            [args],
        ]
        logger.info(f"Сообщение подготовлено к отправке: {invocation}")
        payload = msgpack.packb(invocation, use_bin_type=True)
        packet = encode_with_varint_prefix(payload)
        logger.debug(f"Отправляем сообщение: {packet}")
        await self._ws.send(packet)

    async def invoke_stream(self, target: str, args: list) -> None:
        """
        Отправляет streaming-вызов (StreamInvocation) по протоколу SignalR.
        """
        if not self._ws:
            raise websockets.WebSocketException("Не установлено подключение по wss")
        self._invocation_id = str(self._next_id)
        self._next_id += 1
        invocation = [
            WS_Const.STREAM_INVOCATION_MESSAGE_TYPE,
            WS_Const.DEFAULT_SIGNALR_MAP_HEADERS,
            self._invocation_id,
            target,
            [args],
        ]
        logger.info(f"Streaming-сообщение подготовлено к отправке: {invocation}")
        payload = msgpack.packb(invocation, use_bin_type=True)
        packet = encode_with_varint_prefix(payload)
        await self._ws.send(packet)

    async def receive_by_type(self, message_type: str, timeout: float = WS_Const.FILTERING_TIMEOUT) -> List[Any]:
        """
        Фильтрует сообщения по message_type
        """
        try:
            return await self._receive_by(filter_func=lambda msg: is_desired_type(msg, message_type), timeout=timeout)
        except websockets.WebSocketException:
            raise websockets.WebSocketException(f"Ошибка при фильтрации сообщений по {message_type}")

    async def receive_by_invocation_id(
        self, invocation_id: str, timeout: float = WS_Const.FILTERING_TIMEOUT
    ) -> List[Any]:
        """
        Фильтрует сообщения по invocation_id
        """
        try:
            return await self._receive_by(
                filter_func=lambda msg: is_desired_invocation_id(msg, invocation_id), timeout=timeout
            )
        except websockets.WebSocketException:
            raise websockets.WebSocketException("Ошибка при фильтрации сообщений по invocation_id")

    async def _receive_by(self, filter_func: Callable[[list], bool], timeout: float) -> List[Any]:
        """
        Ждет и фильтрует сообщение по filter_func
        """
        # 1) Единая точка вычисления дедлайна
        deadline = time.monotonic() + timeout

        while True:
            # 2) Остаток времени до таймаута
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise asyncio.TimeoutError(f"Timeout при фильтрации сообщений {timeout:.1f} секунд")

            try:
                # 3) Получает сообщение
                msg = await asyncio.wait_for(self.recv_queue.get(), timeout=remaining)
            except asyncio.TimeoutError:
                # 4) Явно перехватывает и пробрасывает свой TimeoutError
                raise asyncio.TimeoutError(f"Timeout при фильтрации сообщений {timeout:.1f} секунд")

            # 5) Фильтрация по filter_func
            if isinstance(msg, list) and filter_func(msg):
                return msg

    async def reconnect(self):
        """
        Переподключение websocket соединения
        """
        try:
            if self._ping_task:
                self._ping_task.cancel()
                try:
                    await self._ping_task
                except asyncio.CancelledError:
                    pass
                self._ping_task = None

            if self._recv_task:
                try:
                    self._recv_task.cancel()
                    await self._recv_task
                except asyncio.CancelledError:
                    pass
                self._recv_task = None

            if self._ws:
                await self._ws.close()
                self._ws = None
            await self._connect_loop(WS_Const.WS_RECONNECT_TIMEOUT_SECONDS)
            return self
        except (asyncio.TimeoutError, ConnectionError, ConnectionResetError, OSError) as error:
            raise RuntimeError(f"Ошибка при попытке переподключения {error}")