# app/tools/dzengi_probe/stream_market_data_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


MarketDataScenario = str


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


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

    async def run(
        self,
        *,
        max_messages: int | None = 10,
        duration_seconds: int | None = None,
        scenario: MarketDataScenario = "valid",
        repeat: int = 1,
    ) -> StreamProbeResult:
        output_file = (
            f"websocket/marketData.subscribe/"
            f"{self.config.symbol_file_name}.{scenario}.jsonl"
        )

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

        subscribe_message = self._build_subscribe_message(scenario)

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

        return await self._subscribe_and_save(
            subscribe_message=subscribe_message,
            output_file=output_file,
            max_messages=max_messages,
            duration_seconds=duration_seconds,
            repeat=repeat,
        )

    def _build_subscribe_message(
        self,
        scenario: MarketDataScenario,
    ) -> dict[str, Any]:
        if scenario == "valid":
            return {
                "correlationId": "probe-market-data-subscribe",
                "destination": "marketData.subscribe",
                "payload": {
                    "symbols": [self.config.symbol],
                },
            }

        if scenario == "invalid-symbol":
            return {
                "correlationId": "probe-market-data-invalid-symbol",
                "destination": "marketData.subscribe",
                "payload": {
                    "symbols": ["INVALID/SYMBOL"],
                },
            }

        if scenario == "empty-symbols":
            return {
                "correlationId": "probe-market-data-empty-symbols",
                "destination": "marketData.subscribe",
                "payload": {
                    "symbols": [],
                },
            }

        if scenario == "missing-symbols":
            return {
                "correlationId": "probe-market-data-missing-symbols",
                "destination": "marketData.subscribe",
                "payload": {},
            }

        if scenario == "missing-payload":
            return {
                "correlationId": "probe-market-data-missing-payload",
                "destination": "marketData.subscribe",
            }

        if scenario == "invalid-destination":
            return {
                "correlationId": "probe-market-data-invalid-destination",
                "destination": "invalid.subscribe",
                "payload": {
                    "symbols": [self.config.symbol],
                },
            }

        if scenario == "unsubscribe":
            return {
                "correlationId": "probe-market-data-unsubscribe",
                "destination": "marketData.unsubscribe",
                "payload": {
                    "symbols": [self.config.symbol],
                },
            }

        if scenario == "multi-symbol":
            return {
                "correlationId": "probe-market-data-multi-symbol",
                "destination": "marketData.subscribe",
                "payload": {
                    "symbols": [
                        "BTC/USD_LEVERAGE",
                        "ETH/USD_LEVERAGE",
                    ],
                },
            }

        raise ValueError(f"Unknown marketData.subscribe scenario: {scenario}")

    async def _subscribe_and_save(
        self,
        *,
        subscribe_message: dict[str, Any],
        output_file: str,
        max_messages: int | None,
        duration_seconds: int | None,
        repeat: int,
    ) -> StreamProbeResult:
        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:
                for index in range(repeat):
                    message = dict(subscribe_message)
                    message["correlationId"] = f"{subscribe_message.get('correlationId')}-{index + 1}"
                    await websocket.send(json.dumps(message))

                started_at = asyncio.get_running_loop().time()

                while True:
                    if max_messages is not None and messages_saved >= max_messages:
                        break

                    if duration_seconds is not None:
                        elapsed = asyncio.get_running_loop().time() - started_at
                        if elapsed >= duration_seconds:
                            break

                    try:
                        raw_message = await asyncio.wait_for(
                            websocket.recv(),
                            timeout=self.config.timeout_seconds,
                        )
                    except asyncio.TimeoutError:
                        if duration_seconds is not None:
                            continue

                        raise

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

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

            return StreamProbeResult(
                stream="marketData.subscribe",
                output_file=str(output_path),
                ok=True,
                messages_saved=messages_saved,
            )

        except Exception as exc:
            return StreamProbeResult(
                stream="marketData.subscribe",
                output_file=str(output_path),
                ok=False,
                messages_saved=messages_saved,
                error=str(exc),
            )