# app/tools/dzengi_probe/probe.py

from __future__ import annotations

import argparse
import asyncio

from app.tools.dzengi_probe.config import load_config
from app.tools.dzengi_probe.connection_probe import DzengiConnectionProbe
from app.tools.dzengi_probe.rest_probe import DzengiRestProbe
from app.tools.dzengi_probe.stream_market_data_probe import MarketDataStreamProbe
from app.tools.dzengi_probe.websocket_probe import DzengiWebSocketProbe


def run_rest_probe() -> None:
    config = load_config()
    probe = DzengiRestProbe(config)

    results = probe.run_all()

    print("Dzengi REST probe results:")
    for result in results:
        status = "OK" if result.ok else "FAIL"
        print(f"- {status} {result.endpoint} -> {result.output_file}")

        if result.error:
            print(f"  error: {result.error}")


async def run_connection_probe() -> None:
    config = load_config()
    probe = DzengiConnectionProbe(config)

    results = await probe.run_all()

    print("Dzengi connection probe results:")
    for result in results:
        status = "OK" if result.ok else "FAIL"
        print(f"- {status} {result.name}")

        if result.error:
            print(f"  error: {result.error}")


async def run_ws_depth_probe() -> None:
    config = load_config()
    probe = DzengiWebSocketProbe(config)

    result = await probe.probe_depth_request()

    status = "OK" if result.ok else "FAIL"
    print("Dzengi WebSocket Request probe results:")
    print(f"- {status} {result.stream} -> {result.output_file}")
    print(f"  messages saved: {result.messages_saved}")

    if result.error:
        print(f"  error: {result.error}")


async def run_market_data_stream_probe(
    duration: int | None = None,
    scenario: str = "valid",
    repeat: int = 1,
) -> None:
    config = load_config()
    probe = MarketDataStreamProbe(config)

    result = await probe.run(
        max_messages=None if duration else 10,
        duration_seconds=duration,
        scenario=scenario,
        repeat=repeat,
    )

    status = "OK" if result.ok else "FAIL"
    print("Dzengi Stream probe results:")
    print(f"- {status} {result.stream} -> {result.output_file}")
    print(f"  messages saved: {result.messages_saved}")

    if result.error:
        print(f"  error: {result.error}")


def main() -> None:
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "target",
        choices=[
            "rest",
            "ws-depth",
            "stream-market-data",
            "connection",
        ],
    )

    parser.add_argument(
        "--duration",
        type=int,
        default=None,
    )

    parser.add_argument(
        "--repeat",
        type=int,
        default=1,
    )

    parser.add_argument(
        "--scenario",
        default="valid",
        choices=[
            "valid",
            "invalid-symbol",
            "empty-symbols",
            "missing-symbols",
            "missing-payload",
            "invalid-destination",
            "unsubscribe",
            "multi-symbol",
        ],
    )

    args = parser.parse_args()

    if args.target == "rest":
        run_rest_probe()
        return

    if args.target == "ws-depth":
        asyncio.run(run_ws_depth_probe())
        return

    if args.target == "stream-market-data":
        asyncio.run(
            run_market_data_stream_probe(
                duration=args.duration,
                scenario=args.scenario,
                repeat=args.repeat,
            )
        )
        return

    if args.target == "connection":
        asyncio.run(run_connection_probe())
        return


if __name__ == "__main__":
    main()