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

import asyncio
import time

import uvloop
from aiohttp import ClientSession, ClientWebSocketResponse
from curl_cffi import AsyncSession, AsyncWebSocket

from server import binary_data_generator

TOTAL_GB: float = 2
CHUNK_SIZE: int = 65536


async def ws_counter_curl_cffi(ws: AsyncWebSocket) -> None:
    """
    Simple task which counts how many bytes were received for curl-cffi.
    """
    recvd_len: int = 0
    try:
        while True:
            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")

            print("Connection established (curl-cffi). Starting data transmission...")
            start_time: float = time.perf_counter()
            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_bytes(payload=data_chunk)

            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
    except Exception as e:
        print(f"curl-cffi benchmark failed: {e}")
        print("--- curl-cffi Benchmark Failed ---")


async def ws_counter_aiohttp(ws: ClientWebSocketResponse) -> None:
    """
    Simple task which counts how many bytes were received for aiohttp.
    """
    recvd_len: int = 0
    try:
        while True:
            recvd_len += len(await ws.receive_bytes())
    except asyncio.CancelledError:
        print(f"Counter (aiohttp) returned: {recvd_len / (1024**3):.2f} GB")


async def run_benchmark_aiohttp() -> None:
    """Simple client which sends binary messages using aiohttp."""
    print("--- Starting aiohttp Benchmark ---")
    try:
        async with ClientSession() as session:
            async with session.ws_connect(url="wss://127.0.0.1:4443/ws", ssl=False) as ws:
                print("Connection established (aiohttp). Starting data transmission...")
                start_time: float = time.perf_counter()
                counter_task: asyncio.Task[None] = asyncio.create_task(ws_counter_aiohttp(ws))
                async for data_chunk in binary_data_generator(
                    total_gb=TOTAL_GB, chunk_size=CHUNK_SIZE
                ):
                    await ws.send_bytes(data_chunk)

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

                print("\n--- aiohttp 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
    except Exception as e:
        print(f"aiohttp benchmark failed: {e}")
        print("--- aiohttp Benchmark Failed ---")


async def main() -> None:
    """Entrypoint to run both benchmarks sequentially."""
    await run_benchmark_curl_cffi()
    print("\n" + "=" * 40 + "\n")
    await asyncio.sleep(2)  # Brief pause between benchmarks
    await run_benchmark_aiohttp()


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