# app/tools/dzengi_probe/websocket_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
from app.tools.dzengi_probe.response_store import append_jsonl, save_json


@dataclass(frozen=True)
class WebSocketProbeResult:
    stream: str
    output_file: str
    ok: bool
    messages_saved: int
    error: str | None = None


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

    async def probe_market_data(self, *, max_messages: int = 10) -> WebSocketProbeResult:
        output_file = (
            f"websocket/marketData.subscribe/"
            f"{self.config.symbol_file_name}.jsonl"
        )

        request_file = (
            f"websocket/marketData.subscribe/"
            f"{self.config.symbol_file_name}.request.json"
        )

        subscribe_message: dict[str, Any] = {
            "correlationId": "probe-market-data-subscribe",
            "destination": "marketData.subscribe",
            "payload": {
                "symbols": [self.config.symbol],
            },
        }

        save_json(self.config.output_dir / request_file, subscribe_message)

        return await self._subscribe_and_save(
            stream="marketData.subscribe",
            subscribe_message=subscribe_message,
            output_file=output_file,
            max_messages=max_messages,
        )

    async def _subscribe_and_save(
        self,
        *,
        stream: str,
        subscribe_message: dict[str, Any],
        output_file: str,
        max_messages: int,
    ) -> WebSocketProbeResult:
        output_path = self.config.output_dir / output_file
        messages_saved = 0

        if output_path.exists():
            output_path.unlink()

        try:
            async with websockets.connect(
                self.config.ws_url,
                extra_headers={
                    "Origin": self.config.base_url,
                    "Content-Type": "application/json",
                },
                subprotocols=[Subprotocol("json")],
                ping_interval=20,
                ping_timeout=20,
                close_timeout=5,
            ) as websocket:
                await websocket.send(json.dumps(subscribe_message))

                while messages_saved < max_messages:
                    raw_message = await asyncio.wait_for(
                        websocket.recv(),
                        timeout=self.config.timeout_seconds,
                    )

                    try:
                        payload = json.loads(raw_message)
                    except json.JSONDecodeError:
                        payload = {"raw": raw_message}

                    print(payload)
                    append_jsonl(output_path, payload)
                    messages_saved += 1

            return WebSocketProbeResult(
                stream=stream,
                output_file=str(output_path),
                ok=True,
                messages_saved=messages_saved,
            )

        except Exception as exc:
            return WebSocketProbeResult(
                stream=stream,
                output_file=str(output_path),
                ok=False,
                messages_saved=messages_saved,
                error=str(exc),
            )

    async def probe_depth_request(self) -> WebSocketProbeResult:
        output_file = f"websocket/depth/{self.config.symbol_file_name}.jsonl"

        request_message: dict[str, Any] = {
            "correlationId": "probe-depth-request",
            "destination": f"/api/{self.config.api_version}/depth",
            "payload": {
                "symbol": self.config.symbol,
                "limit": int(self.config.depth_limit),
            },
        }

        return await self._subscribe_and_save(
            stream="wss:/api/v1/depth",
            subscribe_message=request_message,
            output_file=output_file,
            max_messages=1,
        )