from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass

import websockets

from yuxi.utils.logging_config import logger


@dataclass
class ProbeResult:
    url: str
    connected: bool
    latency_ms: float | None = None
    error: str | None = None


async def probe_relay(url: str, timeout: float = 10.0) -> ProbeResult:
    start = asyncio.get_event_loop().time()
    try:
        ws = await asyncio.wait_for(
            websockets.connect(url, ping_interval=20, ping_timeout=10),
            timeout=timeout,
        )
        latency_ms = (asyncio.get_event_loop().time() - start) * 1000
        sub_id = "forcepilot_probe"
        req = json.dumps(["REQ", sub_id, {"kinds": [1], "limit": 1}])
        await asyncio.wait_for(ws.send(req), timeout=5)
        eose_received = False
        try:
            while not eose_received:
                raw = await asyncio.wait_for(ws.recv(), timeout=5)
                data = json.loads(raw)
                if isinstance(data, list) and len(data) >= 2 and data[0] == "EOSE":
                    eose_received = True
        except TimeoutError:
            pass
        close_req = json.dumps(["CLOSE", sub_id])
        try:
            await ws.send(close_req)
        except Exception:
            pass
        await ws.close()
        return ProbeResult(url=url, connected=True, latency_ms=latency_ms)
    except Exception as e:
        latency_ms = (asyncio.get_event_loop().time() - start) * 1000
        logger.debug(f"Relay 探测失败 {url}: {e}")
        return ProbeResult(url=url, connected=False, latency_ms=latency_ms, error=str(e))