# app/src/integrations/exchange/ws_client.py

from __future__ import annotations

import asyncio
import json
from typing import AsyncIterator, cast
from uuid import uuid4

import websockets
from websockets.typing import Subprotocol

from src.core.config import load_settings
from src.core.numbers import safe_float
from src.core.types import JsonDict, JsonList, NumericLike


class ExchangeWebSocketClient:
    def __init__(self) -> None:
        self.settings = load_settings()
        self.base_url = self._build_ws_base_url()

    # собрать корректный websocket URL из настроек
    def _build_ws_base_url(self) -> str:
        raw_url = self.settings.exchange_ws_url or self.settings.exchange_base_url

        if raw_url.startswith("https://"):
            raw_url = raw_url.replace("https://", "wss://", 1)
        elif raw_url.startswith("http://"):
            raw_url = raw_url.replace("http://", "ws://", 1)

        raw_url = raw_url.rstrip("/")

        if raw_url.endswith("/connect"):
            return raw_url

        return f"{raw_url}/connect"

    # безопасно нормализовать паузу между websocket-запросами
    def _interval_seconds(self, value: NumericLike | None) -> float:
        interval = safe_float(value)

        if interval is None or interval <= 0:
            return 1.0

        return interval

    # собрать headers для подключения к websocket
    def _headers(self) -> dict[str, str]:
        headers = {
            "Origin": self.settings.exchange_base_url.rstrip("/"),
            "Content-Type": "application/json",
        }

        if self.settings.exchange_api_key:
            headers["X-MBX-APIKEY"] = self.settings.exchange_api_key

        return headers

    # собрать payload запроса стакана
    def _depth_request(self, symbol: str) -> JsonDict:
        return {
            "correlationId": str(uuid4()),
            "destination": "/api/v1/depth",
            "payload": {
                "limit": 5,
                "symbol": symbol,
            },
        }

    # безопасно разобрать JSON от websocket
    def _loads_json(self, raw_message: str | bytes) -> JsonDict | JsonList | None:
        try:
            payload = json.loads(raw_message)
        except json.JSONDecodeError:
            return None

        if isinstance(payload, dict):
            return cast(JsonDict, payload)

        if isinstance(payload, list):
            return cast(JsonList, payload)

        return None

    # поток данных стакана по websocket
    async def stream_depth(
        self,
        symbol: str,
        *,
        interval_seconds: NumericLike = 1.0,
    ) -> AsyncIterator[JsonDict]:
        interval = self._interval_seconds(interval_seconds)
        headers = self._headers()
        timeout_count = 0
        max_timeouts = 3

        async with websockets.connect(
            self.base_url,
            extra_headers=headers,
            subprotocols=[Subprotocol("json")],
            ping_interval=20,
            open_timeout=self.settings.exchange_timeout_sec,
        ) as websocket:
            last_ping_at = 0.0

            while True:
                now = asyncio.get_running_loop().time()

                if now - last_ping_at >= 5.0:
                    pong = await websocket.ping()
                    await asyncio.wait_for(
                        pong,
                        timeout=self.settings.exchange_timeout_sec,
                    )
                    last_ping_at = now

                request = self._depth_request(symbol)
                await websocket.send(json.dumps(request))

                try:
                    raw_message = await asyncio.wait_for(
                        websocket.recv(),
                        timeout=self.settings.exchange_timeout_sec,
                    )
                except asyncio.TimeoutError as exc:
                    timeout_count += 1

                    if timeout_count >= max_timeouts:
                        raise RuntimeError(
                            "WebSocket depth stream timed out repeatedly."
                        ) from exc

                    await asyncio.sleep(interval)
                    continue

                timeout_count = 0

                if not isinstance(raw_message, (str, bytes)):
                    await asyncio.sleep(interval)
                    continue

                payload = self._loads_json(raw_message)

                if isinstance(payload, dict):
                    yield payload

                await asyncio.sleep(interval)