dzentra_bot/app/scripts/check_trades_websocket.py

# app/scripts/check_trades_websocket.py

from __future__ import annotations

import asyncio
import json
from typing import Any
from uuid import uuid4

import websockets
from websockets.typing import Subprotocol

from src.core.config import load_settings


SYMBOLS = (
    "BTC/USD_LEVERAGE",
)

MAX_MESSAGES = 100
TOTAL_TIMEOUT_SECONDS = 3000.0
RECEIVE_TIMEOUT_SECONDS = 15.0


def build_ws_url() -> str:
    settings = load_settings()

    raw_url = settings.exchange_ws_url or 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"


def build_headers() -> dict[str, str]:
    settings = load_settings()

    headers = {
        "Origin": settings.exchange_base_url.rstrip("/"),
        "Content-Type": "application/json",
    }

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

    return headers


def build_subscribe_request(
    symbols: tuple[str, ...],
) -> dict[str, Any]:
    return {
        "correlationId": str(uuid4()),
        "destination": "trades.subscribe",
        "payload": {
            "symbols": list(symbols),
        },
    }


def format_json(
    raw_message: str | bytes | bytearray | memoryview,
) -> str:
    if isinstance(raw_message, str):
        text = raw_message
    elif isinstance(raw_message, memoryview):
        text = raw_message.tobytes().decode(
            "utf-8",
            errors="replace",
        )
    else:
        text = bytes(raw_message).decode(
            "utf-8",
            errors="replace",
        )

    try:
        document = json.loads(text)
    except json.JSONDecodeError:
        return text

    return json.dumps(
        document,
        ensure_ascii=False,
        indent=2,
        sort_keys=True,
    )


def decode_document(
    raw_message: str | bytes | bytearray | memoryview,
) -> object | None:
    if isinstance(raw_message, str):
        text = raw_message
    elif isinstance(raw_message, memoryview):
        text = raw_message.tobytes().decode(
            "utf-8",
            errors="replace",
        )
    else:
        text = bytes(raw_message).decode(
            "utf-8",
            errors="replace",
        )

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return None


def is_trade_event(document: object) -> bool:
    if not isinstance(document, dict):
        return False

    return document.get("destination") == "internal.trade"


def elapsed_seconds(
    loop: asyncio.AbstractEventLoop,
    started_at: float,
) -> float:
    return max(0.0, loop.time() - started_at)


async def receive_trades() -> None:
    settings = load_settings()

    ws_url = build_ws_url()
    headers = build_headers()
    request = build_subscribe_request(SYMBOLS)

    print(f"WebSocket URL: {ws_url}")
    print(f"Symbols: {', '.join(SYMBOLS)}")
    print(f"Maximum messages: {MAX_MESSAGES}")
    print(f"Total timeout: {TOTAL_TIMEOUT_SECONDS:.0f} seconds")
    print()

    print("Subscription request:")
    print(
        json.dumps(
            request,
            ensure_ascii=False,
            indent=2,
        )
    )
    print()

    async with websockets.connect(
        ws_url,
        extra_headers=headers,
        subprotocols=[Subprotocol("json")],
        ping_interval=20,
        ping_timeout=float(settings.exchange_timeout_sec),
        open_timeout=float(settings.exchange_timeout_sec),
        close_timeout=float(settings.exchange_timeout_sec),
    ) as websocket:
        await websocket.send(
            json.dumps(
                request,
                ensure_ascii=False,
            )
        )

        print("Subscription request sent.")
        print("Waiting for acknowledgement and trade events...")
        print()

        loop = asyncio.get_running_loop()
        started_at = loop.time()
        deadline = started_at + TOTAL_TIMEOUT_SECONDS

        received_count = 0
        trade_event_count = 0
        timeout_count = 0

        while received_count < MAX_MESSAGES:
            remaining_seconds = deadline - loop.time()

            if remaining_seconds <= 0:
                print(
                    "Total timeout reached after "
                    f"{elapsed_seconds(loop, started_at):.1f} seconds."
                )
                break

            receive_timeout = min(
                RECEIVE_TIMEOUT_SECONDS,
                remaining_seconds,
            )

            try:
                raw_message = await asyncio.wait_for(
                    websocket.recv(),
                    timeout=receive_timeout,
                )
            except asyncio.TimeoutError:
                timeout_count += 1

                print(
                    "No trade event yet; connection is alive. "
                    f"Elapsed: {elapsed_seconds(loop, started_at):.1f}s, "
                    f"remaining: {max(0.0, remaining_seconds):.1f}s."
                )
                continue

            if not isinstance(
                raw_message,
                (str, bytes, bytearray, memoryview),
            ):
                print(
                    "Skipped unsupported WebSocket message type: "
                    f"{type(raw_message).__name__}"
                )
                continue

            received_count += 1
            document = decode_document(raw_message)

            if is_trade_event(document):
                trade_event_count += 1
                message_kind = "TRADE EVENT"
            else:
                message_kind = "CONTROL / ACKNOWLEDGEMENT"

            print("=" * 80)
            print(
                f"Message #{received_count} — {message_kind} — "
                f"elapsed {elapsed_seconds(loop, started_at):.1f}s"
            )
            print(format_json(raw_message))
            print()

        print("=" * 80)
        print("Diagnostic summary")
        print(f"Received messages: {received_count}")
        print(f"Trade events: {trade_event_count}")
        print(f"Receive timeouts: {timeout_count}")
        print(
            "Elapsed time: "
            f"{elapsed_seconds(loop, started_at):.1f} seconds"
        )


def main() -> None:
    try:
        asyncio.run(receive_trades())
    except KeyboardInterrupt:
        print()
        print("Stopped by user.")
    except Exception as exc:
        print()
        print(
            "WebSocket diagnostic failed: "
            f"{type(exc).__name__}: {exc}"
        )
        raise


if __name__ == "__main__":
    main()