#!/usr/bin/env python3
"""
Example WebSocket client benchmark
"""

import asyncio
import time

import uvloop
from curl_cffi import AsyncSession, AsyncWebSocket

from server_ssl import binary_data_generator

TOTAL_GB: float = 10
CHUNK_SIZE: int = 65536


async def ws_counter_curl_cffi(ws: AsyncWebSocket) -> None:
    """
    Simple task which counts how many bytes were received.
    """
    recvd_len: int = 0
    try:
        async for msg in ws:
            recvd_len += len(msg)
    except asyncio.CancelledError:
        print(f"Counter (curl-cffi) returned: {recvd_len / (1024**3):.2f} GB")


async def run_benchmark_curl_cffi() -> None:
    """Simple client which sends binary messages using curl-cffi."""
    print("--- Starting curl-cffi Benchmark ---")
    try:
        async with AsyncSession(impersonate="chrome", verify=False) as session:
            ws: AsyncWebSocket = await session.ws_connect(
                "wss://127.0.0.1:4443/ws", coalesce_frames=True
            )

            print("Connection established (curl-cffi). Starting data transmission...")
            start_time: float = time.perf_counter()
            # Uncomment for send/recv or both
            # counter_task: asyncio.Task[None] = asyncio.create_task(ws_counter_curl_cffi(ws))
            # async for data_chunk in binary_data_generator(total_gb=TOTAL_GB, chunk_size=CHUNK_SIZE):
            #     _ = await ws.send(payload=data_chunk)

            # Recv only test
            await ws_counter_curl_cffi(ws)

            end_time: float = time.perf_counter()
            duration: float = end_time - start_time
            rate_gbps: float = (TOTAL_GB * 8) / duration

            print("\n--- curl-cffi Benchmark Complete ---")
            print(f"Sent {TOTAL_GB:.2f} GB of data in {duration:.2f} seconds.")
            print(f"Average throughput: {rate_gbps:.2f} Gbps.")

            # try:
            #     _ = counter_task.cancel()
            #     await counter_task
            # except asyncio.CancelledError:
            #     pass  # This is expected
            await ws.close()

    except Exception as e:
        print(f"curl-cffi benchmark failed: {e}")
        print("--- curl-cffi Benchmark Failed ---")


async def heartbeat():
    """A simple coroutine that prints a dot to prove the event loop is alive."""
    counter = 0
    while True:
        await asyncio.sleep(0.05)
        print(".", end="", flush=True)
        counter += 1
        if counter % 100 == 0:
            print("")


async def main() -> None:
    print("Starting sanity check. You should see a continuous stream of dots '.'")
    print("If the dots stop for a long time, the event loop is blocked.")

    # Create the heartbeat task
    heartbeat_task = asyncio.create_task(heartbeat())

    try:
        # Run your benchmark. The heartbeat will run in the background.
        await run_benchmark_curl_cffi()
    finally:
        # Stop the heartbeat task when the benchmark is done
        heartbeat_task.cancel()
        await asyncio.gather(heartbeat_task, return_exceptions=True)
        print("\nSanity check complete.")


if __name__ == "__main__":
    uvloop.run(main())