# Install the websocket-client library
!pip install websocket-client
!pip install websockets # Install websockets for the mock server

# Import necessary libraries
import websocket
import json
import time
import threading
from collections import deque
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from datetime import datetime

# Imports for mock WebSocket server
import asyncio
import websockets
import random

### Global Data Storage and WebSocket Client Callbacks
# A deque to store the last N prices for visualization
price_data_queue = deque(maxlen=100)  # Store up to 100 data points

def on_message(ws, message):
    """
    Callback function to handle incoming WebSocket messages.
    Parses the JSON message and stores relevant price data.

    Args:
        ws: The WebSocketApp instance.
        message: The raw message string received from the WebSocket.
    """
    try:
        json_message = json.loads(message)
        # For our mock server, we will use 'p' for price and 'E' for event time
        if 'p' in json_message and 'E' in json_message:
            price = float(json_message['p'])
            event_time_ms = int(json_message['E'])
            # Convert milliseconds to seconds for datetime object
            event_time = datetime.fromtimestamp(event_time_ms / 1000)

            price_data_queue.append({'time': event_time, 'price': price})
            # print(f"Time: {event_time.strftime('%H:%M:%S')}, Price: {price}") # Uncomment to see live prints

    except Exception as e:
        print(f"Error parsing message: {e}")

def on_error(ws, error):
    """
    Callback function to handle WebSocket errors.

    Args:
        ws: The WebSocketApp instance.
        error: The error object.
    """
    print(f"WebSocket Error: {error}")

def on_close(ws, close_status_code, close_msg):
    """
    Callback function to handle WebSocket closure.

    Args:
        ws: The WebSocketApp instance.
        close_status_code: The status code of the closure.
        close_msg: The close message.
    """
    print("### closed ###")

def on_open(ws):
    """
    Callback function when the WebSocket connection is opened.

    Args:
        ws: The WebSocketApp instance.
    """
    print("### opened ###")

### WebSocket Client Execution Function
def run_websocket_client(url, duration_seconds=30):
    """
    Runs a WebSocket client for a specified duration.

    Args:
        url (str): The WebSocket URL to connect to.
        duration_seconds (int): The duration in seconds to keep the connection open.
    """
    ws = websocket.WebSocketApp(url,
                                on_open=on_open,
                                on_message=on_message,
                                on_error=on_error,
                                on_close=on_close)

    # Run in a separate thread to allow control flow in the main thread
    wst = threading.Thread(target=ws.run_forever)
    wst.daemon = True  # Allow the main program to exit even if server thread is running
    wst.start()

    print(f"Listening for price updates for {duration_seconds} seconds...")
    time.sleep(duration_seconds)  # Keep the connection open for the specified duration
    ws.close()
    print("WebSocket client stopped.")

### Mock WebSocket Server Implementation

import asyncio
import websockets
import random
import threading
import time

# Global variables
_mock_server = None
_mock_server_thread = None
_server_loop = None

async def time_sync_mock_server(websocket):  # REMOVED the 'path' parameter
    """Mock WebSocket server that generates simulated price data.
    Sends price updates every 100ms for a duration.
    """
    current_price = 100.0  # Starting price for simulation
    print("Mock WebSocket client connected.")
    try:
        for i in range(200):  # Emit data points for a duration (e.g., 200 data points)
            # Simulate price fluctuation
            current_price += random.uniform(-0.5, 0.5)
            # Ensure price doesn't go too low for realism
            if current_price < 90.0:
                current_price = 90.0 + random.uniform(0, 1)

            message_data = {
                "E": int(time.time() * 1000),  # Event time in milliseconds
                "p": f"{current_price:.2f}"  # Price as string with 2 decimal places
            }
            await websocket.send(json.dumps(message_data))
            await asyncio.sleep(0.1)  # Simulate real-time updates every 100ms
    except websockets.exceptions.ConnectionClosed:
        print("Mock WebSocket server connection closed by client.")
    except Exception as e:
        print(f"Mock WebSocket server error: {e}")
    finally:
        print("Mock WebSocket client disconnected.")

async def start_server(host, port):
    """Start the WebSocket server."""
    return await websockets.serve(time_sync_mock_server, host, port)

def start_mock_server_in_thread(host="127.0.0.1", port=8765):
    """
    Starts the mock WebSocket server in a separate thread with its own event loop.
    """
    global _mock_server, _mock_server_thread, _server_loop

    if _mock_server:
        print("Mock server already running.")
        return

    print(f"Attempting to start mock WebSocket server on ws://{host}:{port}")

    def run_server():
        global _mock_server, _server_loop
        # Create a new event loop for this thread
        _server_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(_server_loop)

        # Start the server
        _mock_server = _server_loop.run_until_complete(start_server(host, port))
        print(f"Mock WebSocket server listening on ws://{host}:{port}")

        # Run the event loop forever
        _server_loop.run_forever()

    # Start server in a daemon thread
    _mock_server_thread = threading.Thread(target=run_server, daemon=True)
    _mock_server_thread.start()

    # Give the server time to start
    time.sleep(1)

def stop_mock_server_in_thread():
    """
    Stops the mock WebSocket server running in a separate thread.
    """
    global _mock_server, _mock_server_thread, _server_loop

    if _mock_server:
        print("Stopping mock WebSocket server...")

        def shutdown():
            _mock_server.close()
            _server_loop.stop()

        # Schedule shutdown in the server's event loop
        if _server_loop and _server_loop.is_running():
            _server_loop.call_soon_threadsafe(shutdown)

        # Wait a bit for cleanup
        time.sleep(0.5)
        _mock_server = None
        _mock_server_thread = None
        _server_loop = None
        print("Mock WebSocket server stopped.")
    else:
        print("Mock WebSocket server not running.")

### Orchestrate Mock Server and Client Execution

# Define the local WebSocket server URL
MOCK_WS_SERVER_URL = "ws://127.0.0.1:8765"

# Clear previous data before starting new collection
price_data_queue.clear()

# Start the mock server in a separate thread
print("Starting mock WebSocket server...")
start_mock_server_in_thread(host="127.0.0.1", port=8765)

# Give the server time to fully initialize
time.sleep(2)

# Run the WebSocket client to collect data from the mock server
print("Running WebSocket client to connect to mock server...")
run_websocket_client(MOCK_WS_SERVER_URL, duration_seconds=15)

# Stop the mock server
stop_mock_server_in_thread()

print("Mock server and client execution complete.")

# Display collected data count
print(f"\nCollected {len(price_data_queue)} price points.")

print(f"Collected {len(price_data_queue)} price points.")
if price_data_queue:
    print("First 5 collected data points (Mock Data):")
    for i, data in enumerate(list(price_data_queue)[:5]):
        print(f"  {i+1}. Time: {data['time'].strftime('%H:%M:%S.%f')[:-3]}, Price: {data['price']}")
    print(f"\nLast 5 collected data points (Mock Data):")
    for i, data in enumerate(list(price_data_queue)[-5:]):
        print(f"  {i+1}. Time: {data['time'].strftime('%H:%M:%S.%f')[:-3]}, Price: {data['price']}")
else:
    print("No data collected. Please ensure the mock server is running and the client connected successfully.")

if price_data_queue:
    # Extract times and prices
    times = [data['time'] for data in price_data_queue]
    prices = [data['price'] for data in price_data_queue]

    # Create the plot
    fig, ax = plt.subplots(figsize=(12, 6))
    ax.plot(times, prices, marker='.', linestyle='-', color='skyblue',
            markersize=4, linewidth=1.5, label='Price')

    # Add styling
    ax.set_title('Real-time Price Feed Simulation', fontsize=14, fontweight='bold')
    ax.set_xlabel('Time', fontsize=12)
    ax.set_ylabel('Price (USD)', fontsize=12)
    ax.grid(True, linestyle='--', alpha=0.3)
    ax.legend()

    # Format x-axis for dates
    fig.autofmt_xdate()

    # Add some statistics
    avg_price = sum(prices) / len(prices)
    min_price = min(prices)
    max_price = max(prices)

    # Add text box with stats
    stats_text = f'Avg: ${avg_price:.2f}\nMin: ${min_price:.2f}\nMax: ${max_price:.2f}'
    ax.text(0.02, 0.98, stats_text, transform=ax.transAxes,
            verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

    plt.tight_layout()
    plt.show()

    # Print statistics
    print(f"\nPrice Statistics:")
    print(f"  Average Price: ${avg_price:.2f}")
    print(f"  Minimum Price: ${min_price:.2f}")
    print(f"  Maximum Price: ${max_price:.2f}")
    print(f"  Price Range: ${max_price - min_price:.2f}")
else:
    print("Cannot create plot: No price data available.")