pip install websockets tenacity pandas matplotlib seaborn numpy certifi

import asyncio
import json
import time
import logging
import random
import ssl
import certifi
from collections import deque
from typing import Dict, Any, Optional

import websockets
from tenacity import retry, stop_after_attempt, wait_exponential, wait_random, retry_if_exception_type
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def create_client_state(websocket_url: str, max_latency_history: int = 1000) -> Dict[str, Any]:
    """
    Initializes the client state dictionary.

    Parameters
    ----------
    websocket_url : str
        The URL of the WebSocket endpoint.
    max_latency_history : int, optional
        The maximum number of latency measurements to store, defaults to 1000.

    Returns
    -------
    dict
        An initialized client state dictionary.
    """
    state = {
        "websocket_url": websocket_url,
        "websocket": None,  # Will hold the websocket connection object
        "connected": False,
        "message_queue": asyncio.Queue(),
        "received_data": [],  # To store processed market data
        "latency_history": deque(maxlen=max_latency_history), # Rolling window for latency
        "last_ping_time": None,
        "last_pong_time": None,
        "subscriptions": [],
        "sequence_numbers": {}, # To track message sequence numbers per symbol
        "state_lock": asyncio.Lock() # For protecting state modifications in async context
    }
    logger.info(f"Client state initialized for URL: {websocket_url}")
    return state

@retry(
    stop=stop_after_attempt(10),
    wait=wait_exponential(multiplier=1, min=4, max=60) + wait_random(0, 5),
    retry=retry_if_exception_type(websockets.WebSocketException)
)
async def connect_websocket(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Establishes a WebSocket connection with retry logic.

    Parameters
    ----------
    state : dict
        Current client state dictionary.

    Returns
    -------
    dict
        Updated state with the websocket connection object.
    """
    async with state["state_lock"]:
        if state["websocket"] and state["connected"]:
            logger.info("WebSocket already connected.")
            return state

        logger.info(f"Attempting to connect to {state['websocket_url']}...")
        ssl_context = ssl.create_default_context(cafile=certifi.where())
        try:
            websocket = await websockets.connect(state["websocket_url"], ssl=ssl_context)
            state["websocket"] = websocket
            state["connected"] = True
            logger.info("WebSocket connection established successfully.")
        except websockets.WebSocketException as e:
            logger.error(f"WebSocket connection failed: {e}")
            state["connected"] = False
            raise # Re-raise to trigger retry
        except Exception as e:
            logger.critical(f"An unexpected error occurred during connection: {e}")
            state["connected"] = False
            raise # Re-raise for tenacity to handle
    return state

async def subscribe_to_market_data(state: Dict[str, Any], symbols: list) -> Dict[str, Any]:
    """
    Sends a subscription message to the WebSocket server for market data.
    Assumes a typical crypto exchange subscription format.

    Parameters
    ----------
    state : dict
        Current client state dictionary.
    symbols : list
        A list of cryptocurrency symbols to subscribe to (e.g., ['BTC/USDT', 'ETH/USDT']).

    Returns
    -------
    dict
        Updated state with active subscriptions.
    """
    async with state["state_lock"]:
        if not state["connected"] or not state["websocket"]:
            logger.warning("Not connected to WebSocket. Cannot subscribe.")
            return state

        for symbol in symbols:
            # Example subscription message for a generic crypto exchange
            # Adjust this payload based on the actual exchange API documentation
            subscription_message = {
                "method": "SUBSCRIBE",
                "params": [
                    f"{symbol.lower()}@trade", # Example: BTC/USDT@trade for trades
                    f"{symbol.lower()}@kline_1m" # Example: BTC/USDT@kline_1m for 1-minute candles
                ],
                "id": int(time.time() * 1000) # Unique ID for the request
            }
            try:
                await state["websocket"].send(json.dumps(subscription_message))
                state["subscriptions"].append(symbol)
                state["sequence_numbers"][symbol] = 0 # Initialize sequence number tracker
                logger.info(f"Subscribed to market data for {symbol}")
            except websockets.WebSocketException as e:
                logger.error(f"Failed to send subscription for {symbol}: {e}")
            except Exception as e:
                logger.error(f"An unexpected error occurred during subscription for {symbol}: {e}")
    return state

async def receive_message(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Continuously receives messages from the WebSocket and puts them into a queue.

    Parameters
    ----------
    state : dict
        Current client state dictionary.

    Returns
    -------
    dict
        Updated state (primarily message_queue).
    """
    if not state["connected"] or not state["websocket"]:
        logger.warning("Not connected to WebSocket. Cannot receive messages.")
        return state

    try:
        async for message in state["websocket"]:
            receive_timestamp = time.perf_counter() # High-resolution timer
            await state["message_queue"].put((message, receive_timestamp))
            # logger.debug(f"Received message and queued. Queue size: {state['message_queue'].qsize()}")
    except websockets.exceptions.ConnectionClosedOK:
        logger.info("WebSocket connection closed gracefully.")
        state["connected"] = False
    except websockets.exceptions.ConnectionClosedError as e:
        logger.error(f"WebSocket connection closed with error: {e}")
        state["connected"] = False
    except Exception as e:
        logger.critical(f"An unexpected error occurred while receiving messages: {e}")
        state["connected"] = False
    finally:
        # Ensure 'connected' is false if loop exits unexpectedly
        if state["connected"]:
            async with state["state_lock"]:
                state["connected"] = False
        logger.info("Message receiving task stopped.")
    return state

async def process_trade_data(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Processes messages from the queue, extracts trade data, and tracks latency.

    Parameters
    ----------
    state : dict
        Current client state dictionary.

    Returns
    -------
    dict
        Updated state with processed data and latency history.
    """
    while state["connected"] or not state["message_queue"].empty():
        try:
            message_str, receive_timestamp = await asyncio.wait_for(state["message_queue"].get(), timeout=1.0)
            # logger.debug(f"Dequeued message. Queue size: {state['message_queue'].qsize()}")

            try:
                message = json.loads(message_str)
            except json.JSONDecodeError:
                logger.warning(f"Could not decode JSON: {message_str[:100]}...")
                continue

            # Example processing for a generic crypto trade message or Binance direct trade event
            trade_data = None
            if isinstance(message, dict):
                if "data" in message and isinstance(message["data"], dict) and message["data"].get("e") == "trade":
                    trade_data = message["data"]
                elif message.get("e") == "trade":
                    trade_data = message

            if trade_data is not None:
                symbol = trade_data.get("s") # e.g., "BTCUSDT"
                server_timestamp_ms = trade_data.get("T") # Server timestamp in milliseconds
                price = float(trade_data.get("p"))
                quantity = float(trade_data.get("q"))

                if symbol and server_timestamp_ms is not None:
                    latency_ms = (receive_timestamp - (server_timestamp_ms / 1000)) * 1000

                    async with state["state_lock"]:
                        state["latency_history"].append(latency_ms)
                        state["received_data"].append({
                            "timestamp": pd.to_datetime(server_timestamp_ms, unit='ms'),
                            "symbol": symbol,
                            "price": price,
                            "quantity": quantity,
                            "latency_ms": latency_ms
                        })

                        current_seq = trade_data.get("L") # Example: last trade ID
                        if symbol in state["sequence_numbers"] and current_seq is not None:
                            expected_seq = state["sequence_numbers"][symbol] + 1
                            if current_seq != expected_seq and state["sequence_numbers"][symbol] != 0:
                                logger.warning(f"[{symbol}] Possible message gap detected! Expected {expected_seq}, Got {current_seq}")
                            state["sequence_numbers"][symbol] = current_seq
                        elif current_seq is not None:
                            state["sequence_numbers"][symbol] = current_seq

                    logger.debug(f"Processed trade: {symbol} @ {price}, Latency: {latency_ms:.2f}ms")
                else:
                    logger.debug(f"Received trade event without timestamp or symbol: {message}")
            else:
                logger.debug(f"Received non-trade or incomplete message: {message}")

        except asyncio.TimeoutError:
            # logger.debug("No messages in queue for 1 second.")
            await asyncio.sleep(0.1) # Yield control to avoid busy waiting
        except Exception as e:
            logger.error(f"Error processing message: {e}")
            await asyncio.sleep(0.5) # Wait before next attempt
    logger.info("Message processing task stopped.")
    return state

async def track_latency(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Tracks and logs the current latency statistics.

    Parameters
    ----------
    state : dict
        Current client state dictionary.

    Returns
    -------
    dict
        Unchanged state dictionary.
    """
    while state["connected"]:
        async with state["state_lock"]:
            if state["latency_history"]:
                current_latencies = list(state["latency_history"])
                avg_latency = np.mean(current_latencies)
                min_latency = np.min(current_latencies)
                max_latency = np.max(current_latencies)
                logger.info(f"Current Latency (ms) - Avg: {avg_latency:.2f}, Min: {min_latency:.2f}, Max: {max_latency:.2f} (from {len(current_latencies)} samples)")
            else:
                logger.info("No latency data yet.")
        await asyncio.sleep(5) # Log every 5 seconds
    logger.info("Latency tracking task stopped.")
    return state

async def shutdown_websocket(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Gracefully shuts down the WebSocket connection.

    Parameters
    ----------
    state : dict
        Current client state dictionary.

    Returns
    -------
    dict
        Updated state with the connection closed.
    """
    async with state["state_lock"]:
        if state["websocket"] and state["connected"]:
            logger.info("Closing WebSocket connection...")
            try:
                await state["websocket"].close()
                state["connected"] = False
                state["websocket"] = None
                logger.info("WebSocket connection closed.")
            except Exception as e:
                logger.error(f"Error closing WebSocket: {e}")
        else:
            logger.info("WebSocket not active or already closed.")
    return state

# Define the WebSocket URL for a public crypto exchange (e.g., Binance Spot Testnet)
# For production, replace with a real endpoint, e.g., 'wss://stream.binance.com:9443/ws'
# This example uses a mock URL for demonstration purposes unless a real public one is available.
# For a real public websocket, e.g., Binance's testnet: 'wss://testnet.binance.vision/ws/btcusdt@trade'
# Or a combined stream: 'wss://testnet.binance.vision/ws-api/v3'
# For local testing, you might use 'ws://localhost:8765'

# NOTE: For this demonstration, we will simulate a connection with a real Binance public stream.
# Be mindful of rate limits if running for extended periods.
# Binance Spot Testnet (example for future reference if you want to connect to a full stream)
# WEBSOCKET_URL = "wss://testnet.binance.vision/ws-api/v3"

# Using a Binance public trade stream endpoint.
# The live `stream.binance.com` URL might be geo-restricted or blocked for Colab IPs.
# The correct Binance stream path is under `/ws/<symbol>@trade`.
WEBSOCKET_URL = "wss://stream.binance.com:9443/ws/btcusdt@trade" # Example: BTC/USDT 1-minute kline data stream

# Initialize client state
client_state = create_client_state(websocket_url=WEBSOCKET_URL, max_latency_history=5000)

async def run_client_demo():
    global client_state # Allow modification of the global state object

    # 1. Connect to WebSocket
    try:
        client_state = await connect_websocket(client_state)
    except Exception as e:
        logger.error(f"Initial connection failed: {e}. Exiting demo.")
        return

    if not client_state["connected"]:
        logger.error("Failed to connect to WebSocket. Exiting demo.")
        return

    # 2. Subscribe to market data (if applicable for the chosen URL)
    # The 'btcusdt@trade' stream is direct, so no explicit SUBSCRIBE message is often needed
    # If using a combined stream (like 'ws-api/v3'), you'd call subscribe_to_market_data
    # await subscribe_to_market_data(client_state, ['BTCUSDT', 'ETHUSDT'])
    logger.info("Assuming direct trade stream, no explicit subscription needed for this URL.")

    # 3. Start concurrent tasks
    try:
        receive_task = asyncio.create_task(receive_message(client_state))
        process_task = asyncio.create_task(process_trade_data(client_state))
        latency_task = asyncio.create_task(track_latency(client_state))

        logger.info("Running client for 30 seconds to collect data...")
        await asyncio.sleep(30) # Run for a specified duration

    except asyncio.CancelledError:
        logger.info("Client demo cancelled.")
    except Exception as e:
        logger.critical(f"An error occurred during client runtime: {e}")
    finally:
        logger.info("Shutting down client tasks...")
        receive_task.cancel()
        process_task.cancel()
        latency_task.cancel()

        # Wait for tasks to complete cancellation
        await asyncio.gather(receive_task, process_task, latency_task, return_exceptions=True)

        # 4. Shutdown WebSocket connection
        client_state = await shutdown_websocket(client_state)
        logger.info("Client demo finished.")

# Run the demo
# In a Colab environment, use nest_asyncio to run async code directly.
# This is often needed if the event loop is already running (e.g., by Colab's internal machinery).
import nest_asyncio
nest_asyncio.apply()

await run_client_demo()

if client_state["received_data"]:
    df_trades = pd.DataFrame(client_state["received_data"])
    df_trades['timestamp'] = pd.to_datetime(df_trades['timestamp'])
    df_trades.set_index('timestamp', inplace=True)

    print("\n--- Raw Trade Data (First 5 rows) ---")
    display(df_trades.head())

    print("\n--- Trade Data Summary Statistics ---")
    display(df_trades[['price', 'quantity', 'latency_ms']].describe())

    # Calculate total traded volume
    df_trades['value'] = df_trades['price'] * df_trades['quantity']
    total_traded_value = df_trades['value'].sum()
    print(f"\nTotal Traded Value: ${total_traded_value:,.2f}")

    # Resample to 1-second bars for price visualization
    df_resampled = df_trades['price'].resample('1s').ohlc()
    df_resampled = df_resampled.dropna()

    print("\n--- Resampled Price Data (1-Second OHLC) (First 5 rows) ---")
    display(df_resampled.head())

else:
    print("No trade data collected.")
    df_trades = pd.DataFrame()

if not df_trades.empty and 'latency_ms' in df_trades.columns:
    plt.figure(figsize=(12, 6))
    sns.histplot(df_trades['latency_ms'], bins=50, kde=True, color='skyblue')
    plt.title('Distribution of WebSocket Message Latency')
    plt.xlabel('Latency (ms)')
    plt.ylabel('Frequency')
    plt.axvline(df_trades['latency_ms'].mean(), color='red', linestyle='dashed', linewidth=1, label=f'Mean Latency: {df_trades["latency_ms"].mean():.2f} ms')
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.show()

    plt.figure(figsize=(12, 6))
    sns.lineplot(x=df_trades.index, y=df_trades['latency_ms'], color='orange', alpha=0.7)
    plt.title('Latency Over Time')
    plt.xlabel('Time')
    plt.ylabel('Latency (ms)')
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.show()
else:
    print("No latency data available for visualization.")