leaf/examples/ws_manager_demo.rs

//! WebSocket Manager Demo
//!
//! This example demonstrates how to use the WsManager to connect to multiple
//! WebSocket servers, authenticate, and execute scripts on connected clients.
//!
//! # Prerequisites
//!
//! 1. Start one or more Hero WebSocket servers:
//!    ```bash
//!    # Terminal 1
//!    cd hero/interfaces/websocket/server
//!    cargo run --example circle_auth_demo -- --port 3030
//!
//!    # Terminal 2 (optional)
//!    cargo run --example circle_auth_demo -- --port 3031
//!    ```
//!
//! 2. Run this example:
//!    ```bash
//!    cd framework
//!    cargo run --example ws_manager_demo
//!    ```

use framework::WsManager;
use log::{info, warn, error};
use std::time::Duration;
use tokio::time::sleep;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    info!("🚀 Starting WebSocket Manager Demo");

    // Example private key (in real usage, load from secure storage)
    let private_key = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";

    // Create a manager with multiple servers
    let manager = WsManager::builder()
        .private_key(private_key.to_string())
        .add_server_url("ws://localhost:3030".to_string())
        .add_server_url("ws://localhost:3031".to_string())
        .build();

    info!("📡 Connecting to WebSocket servers...");

    // Connect to all configured servers
    match manager.connect().await {
        Ok(()) => info!("✅ Successfully initiated connections"),
        Err(e) => {
            error!("❌ Failed to connect: {}", e);
            return Err(e.into());
        }
    }

    // Give connections time to establish
    sleep(Duration::from_secs(2)).await;

    // Check connection status
    info!("🔍 Checking connection status...");
    let statuses = manager.get_all_connection_statuses();
    for (url, status) in &statuses {
        info!("  {} -> {}", url, status);
    }

    // Demonstrate script execution on a specific server
    info!("🎯 Executing script on specific server...");
    let script = r#"
        console.log("Hello from WebSocket Manager!");
        return "Script executed successfully";
    "#;

    if let Some(future) = manager.with_client("ws://localhost:3030", |client| {
        info!("  📤 Sending script to ws://localhost:3030");
        client.play(script.to_string())
    }) {
        let result = future.await;
        match result {
            Ok(output) => info!("  ✅ Script output: {}", output),
            Err(e) => warn!("  ⚠️ Script error: {}", e),
        }
    } else {
        warn!("  ❌ Server ws://localhost:3030 is not connected");
    }

    // Demonstrate operations on all connected clients
    info!("🌐 Executing operations on all connected clients...");
    
    let ping_script = "return 'pong from ' + new Date().toISOString();";
    
    // Get list of connected URLs first
    let connected_urls = manager.get_connected_urls();
    
    for url in connected_urls {
        info!("  📤 Pinging {}", url);
        
        if let Some(future) = manager.with_client(&url, |client| {
            client.play(ping_script.to_string())
        }) {
            match future.await {
                Ok(output) => info!("  ✅ {} responded: {}", url, output),
                Err(e) => warn!("  ⚠️ {} error: {}", url, e),
            }
        } else {
            warn!("  ❌ {} is not connected", url);
        }
    }

    // Wait for async operations to complete
    sleep(Duration::from_secs(3)).await;

    // Demonstrate authentication status check
    info!("🔐 Checking authentication status...");
    // Note: For complex async operations, you may need to handle them differently
    // due to Rust's lifetime constraints with closures and futures
    info!("  💡 Authentication check would be done here in a real application");
    info!("  💡 Use manager.with_client() to access client methods like whoami()");

    // Demonstrate connection management
    info!("🔧 Testing connection management...");
    
    // Disconnect from a specific server
    info!("  🔌 Disconnecting from ws://localhost:3031...");
    if let Err(e) = manager.disconnect_from_server("ws://localhost:3031").await {
        warn!("  ⚠️ Disconnect error: {}", e);
    } else {
        info!("  ✅ Disconnected successfully");
    }

    // Check status after disconnect
    sleep(Duration::from_secs(1)).await;
    let new_statuses = manager.get_all_connection_statuses();
    for (url, status) in &new_statuses {
        info!("  {} -> {}", url, status);
    }

    // Reconnect
    info!("  🔌 Reconnecting to ws://localhost:3031...");
    if let Err(e) = manager.connect_to_server("ws://localhost:3031").await {
        warn!("  ⚠️ Reconnect error: {}", e);
    } else {
        info!("  ✅ Reconnected successfully");
        sleep(Duration::from_secs(1)).await;
    }

    // Final status check
    info!("📊 Final connection status:");
    let final_statuses = manager.get_all_connection_statuses();
    for (url, status) in &final_statuses {
        info!("  {} -> {}", url, status);
    }

    // Demonstrate adding a connection at runtime
    info!("➕ Adding new connection at runtime...");
    manager.add_connection(
        "ws://localhost:3032".to_string(),
        Some(private_key.to_string())
    );

    // Try to connect to the new server (will fail if server isn't running)
    if let Err(e) = manager.connect_to_server("ws://localhost:3032").await {
        warn!("  ⚠️ Could not connect to ws://localhost:3032: {}", e);
        warn!("  💡 Start a server on port 3032 to test this feature");
    } else {
        info!("  ✅ Connected to new server ws://localhost:3032");
    }

    info!("🎉 WebSocket Manager Demo completed!");
    info!("💡 Key takeaways:");
    info!("   - Use WsManager::builder() to configure multiple servers");
    info!("   - Call connect() to establish all connections");
    info!("   - Use with_client() for operations on specific servers");
    info!("   - Use with_all_clients() for bulk operations");
    info!("   - Connections are managed automatically with keep-alive");

    Ok(())
}