# app/tools/dzengi_probe/connection_probe.py

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass
from typing import Any

import websockets
from websockets.typing import Subprotocol

from app.tools.dzengi_probe.config import DzengiProbeConfig


@dataclass(frozen=True)
class ConnectionProbeResult:
    name: str
    ok: bool
    error: str | None = None


class DzengiConnectionProbe:
    def __init__(self, config: DzengiProbeConfig) -> None:
        self.config = config

    async def run_all(self) -> list[ConnectionProbeResult]:
        return [
            await self._check(
                name="baseline",
                headers={
                    "Origin": self.config.base_url,
                    "Content-Type": "application/json",
                },
                subprotocols=[Subprotocol("json")],
            ),
            await self._check(
                name="without_origin",
                headers={
                    "Content-Type": "application/json",
                },
                subprotocols=[Subprotocol("json")],
            ),
            await self._check(
                name="without_content_type",
                headers={
                    "Origin": self.config.base_url,
                },
                subprotocols=[Subprotocol("json")],
            ),
            await self._check(
                name="without_headers",
                headers=None,
                subprotocols=[Subprotocol("json")],
            ),
            await self._check(
                name="without_subprotocol",
                headers={
                    "Origin": self.config.base_url,
                    "Content-Type": "application/json",
                },
                subprotocols=None,
            ),
            await self._check(
                name="minimal_connection",
                headers=None,
                subprotocols=None,
            ),
        ]

    async def _check(
        self,
        *,
        name: str,
        headers: dict[str, str] | None,
        subprotocols: list[Subprotocol] | None,
    ) -> ConnectionProbeResult:
        subscribe_message: dict[str, Any] = {
            "correlationId": f"connection-probe-{name}",
            "destination": "marketData.subscribe",
            "payload": {
                "symbols": [self.config.symbol],
            },
        }

        try:
            async with websockets.connect(
                self.config.ws_url,
                extra_headers=headers,
                subprotocols=subprotocols,
                ping_interval=20,
                ping_timeout=20,
                close_timeout=5,
            ) as websocket:
                await websocket.send(json.dumps(subscribe_message))

                raw_message = await asyncio.wait_for(
                    websocket.recv(),
                    timeout=self.config.timeout_seconds,
                )

                payload = json.loads(raw_message)

                if payload.get("status") != "OK":
                    return ConnectionProbeResult(
                        name=name,
                        ok=False,
                        error=json.dumps(payload, ensure_ascii=False),
                    )

                return ConnectionProbeResult(name=name, ok=True)

        except Exception as exc:
            return ConnectionProbeResult(
                name=name,
                ok=False,
                error=str(exc),
            )