circles/tests/end_to_end_integration.rs

//! End-to-end example demonstrating client_ws and server_ws working together
//!
//! This example:
//! 1. Starts a Redis server
//! 2. Starts a server_ws instance with authentication enabled
//! 3. Creates a client_ws instance with authentication
//! 4. Demonstrates the complete authentication flow
//! 5. Sends authenticated requests and receives responses
//! 6. Verifies in Redis that the public key was passed with the script
//! 7. Cleans up all processes
//!
//! To run this example:
//! ```bash
//! cargo run --example end_to_end_integration -p integration_tests
//! ```

use std::time::Duration;
use tokio::time::sleep;
use log::{info, error, warn};
use std::process::{Command, Child, Stdio};
use std::sync::Arc;


// Client-side imports
use circle_client_ws::{CircleWsClientBuilder, auth};
// Launcher imports
use launcher::{setup_and_spawn_circles, shutdown_circles, CircleConfig, RunningCircleInfo};
use redis::AsyncCommands;
use secp256k1::{Secp256k1, PublicKey, SecretKey};
use tokio::sync::Mutex;

struct ChildProcessGuard {
    child: Child,
    name: String,
}

impl ChildProcessGuard {
    fn new(child: Child, name: String) -> Self {
        Self { child, name }
    }
}

impl Drop for ChildProcessGuard {
    fn drop(&mut self) {
        info!("Cleaning up {} process (PID: {})...", self.name, self.child.id());
        match self.child.kill() {
            Ok(_) => {
                info!("Successfully sent kill signal to {} (PID: {}).", self.name, self.child.id());
                match self.child.wait() {
                    Ok(status) => info!("{} (PID: {}) exited with status: {}", self.name, self.child.id(), status),
                    Err(e) => warn!("Error waiting for {} (PID: {}): {}", self.name, self.child.id(), e),
                }
            }
            Err(e) => error!("Failed to kill {} (PID: {}): {}", self.name, self.child.id(), e),
        }
    }
}

#[tokio::test]
async fn test_full_end_to_end_example() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging with default level if not set
    if std::env::var("RUST_LOG").is_err() {
        std::env::set_var("RUST_LOG", "info");
    }
    let _ = env_logger::try_init();
    
    info!("🚀 Starting self-contained end-to-end authentication example");
    info!("🔗 Running full end-to-end example with server");

    // Step 1: Start Redis server
    info!("Starting Redis server...");
    let redis_server_process = Command::new("redis-server")
        .arg("--port")
        .arg("6379")
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()?;
    let _redis_server_guard = ChildProcessGuard::new(redis_server_process, "redis-server".to_string());
    info!("Redis server started with PID {}", _redis_server_guard.child.id());
    sleep(Duration::from_millis(500)).await;
    
    // Step 2 & 3: Setup and spawn circle using launcher
    info!("🚀 Setting up and spawning circle via launcher...");
    let circle_name = "e2e_test_circle";
    let circle_configs = vec![CircleConfig {
        name: circle_name.to_string(),
        script_path: Some("./tests/test_data/e2e_script.rhai".to_string()),
        port: 8080, // Correct field name is port
    }];

    let (running_circles_store, _circle_outputs) = setup_and_spawn_circles(circle_configs).await?;
    
    info!("Circles spawned by launcher:");
    for circle_info_arc_loop in &running_circles_store {
        let circle_info_locked_loop = circle_info_arc_loop.lock().expect("Failed to lock circle info for logging");
        info!("  ✅ Name: {}, WS Port: {}, Public Key: {}...",
              circle_info_locked_loop.config.name, circle_info_locked_loop.config.port, &circle_info_locked_loop.public_key[..10]);
    }

    let target_circle_name = "e2e_test_circle"; // This was 'circle_name'
    let mut found_circle_arc_opt: Option<Arc<Mutex<RunningCircleInfo>>> = None; // std::sync::Mutex
    for info_arc_find in &running_circles_store {
        if info_arc_find.lock().expect("Failed to lock circle info for finding target").config.name == target_circle_name {
            found_circle_arc_opt = Some(info_arc_find.clone());
            break;
        }
    }
    let circle_info_arc = found_circle_arc_opt
        .ok_or_else(|| Into::<Box<dyn std::error::Error>>::into(format!("Circle '{}' not found in running_circles_store", target_circle_name)))?;
    
    let circle_info_locked = circle_info_arc.lock().expect("Failed to lock target circle info"); // Lock it for use
    let server_address = format!("127.0.0.1:{}", circle_info_locked.config.port); // Access port via config
    // The main info log for the specific test circle is covered by the loop.
    // If a specific log for the *target* circle is still desired here, it can be added, e.g.:
    // info!("Target circle for test: '{}' at ws://{}/ws, Public Key: {}...",
    //       circle_info_locked.name, server_address, &circle_info_locked.public_key[..10]);
    // The circle_public_key_hex for the server is now circle_info.public_key
    // Client generates its own keypair (Step 4, formerly Step 2 for client keys)
    sleep(Duration::from_millis(1000)).await; // Allow services to fully start
    
    // Step 4: Generate a keypair for the client
    info!("🔑 Generating a new keypair for the client...");
    let client_private_key = auth::generate_private_key()?;
    info!("🔑 Generated client private key: {}...", &client_private_key[..10]);

    let shared_secret = auth::generate_shared_secret(
        &client_private_key,
        &auth::pubkey_from_hex(&circle_info_locked.public_key).expect("Failed to get pubkey from hex")?, // Use public key from the locked RunningCircleInfo
    )?;

    // Step 5: Create authenticated client
    info!("🔌 Creating authenticated WebSocket client...");
    let mut client = CircleWsClientBuilder::new(format!("ws://{}/ws", server_address))
        .with_keypair(client_private_key.clone())
        .build();
    
    // Step 5: Connect to WebSocket
    info!("🔗 Connecting to WebSocket server...");
    client.connect().await?;

    // Step 6: Authenticate the client
    info!("🔐 Authenticating client...");
    match client.authenticate().await {
        Ok(true) => {
            info!("✅ Authentication successful!");
        }
        Ok(false) => {
            error!("❌ Authentication failed!");
            return Err("Authentication failed".into());
        }
        Err(e) => {
            error!("❌ Authentication failed: {}", e);
            return Err(e.into());
        }
    }
    
    // Step 7: Send authenticated requests
    info!("📤 Sending authenticated Rhai script requests...");

    let secp = Secp256k1::new();
    let secret_key_bytes = &hex::decode(&client_private_key).unwrap();
    let secret_key = SecretKey::from_slice(secret_key_bytes).unwrap();
    let expected_public_key = PublicKey::from_secret_key(&secp, &secret_key);
    let expected_public_key_hex = hex::encode(expected_public_key.serialize_uncompressed());
    
    let test_scripts = vec![
        "print(\"Hello from authenticated client!\"); 42",
        "let x = 10; let y = 20; x + y",
        "print(\"Testing authentication...\"); \"success\"",
        "CALLER_PUBLIC_KEY",
    ];

    for (i, script) in test_scripts.iter().enumerate() {
        info!("📝 Executing script {}: {}", i + 1, script);

        match client.play(script.to_string()).await {
            Ok(result) => {
                info!("✅ Script {} result: {}", i + 1, result.output);
                if script == &"CALLER_PUBLIC_KEY" {
                    assert_eq!(result.output, expected_public_key_hex);
                    info!("✅ CALLER_PUBLIC_KEY verification successful!");
                }
            }
            Err(e) => {
                panic!("client.play() failed with error: {:#?}", e);
            }
        }
        
        // Small delay between requests
        sleep(Duration::from_millis(500)).await;
    }
    
    // Step 8: Verify public key in Redis
    info!("🔍 Verifying public key in Redis...");
    let redis_client = redis::Client::open("redis://127.0.0.1:6379/")?;
    let mut redis_conn = redis_client.get_multiplexed_async_connection().await?;

    let mut found_task = false;
    let task_keys: Vec<String> = redis_conn.keys("rhai_task_details:*").await?;
    for key in task_keys {
        let script_content: String = redis_conn.hget(&key, "script").await?;
        if script_content.contains("Testing authentication...") {
            let stored_public_key: String = redis_conn.hget(&key, "publicKey").await?;
            assert_eq!(stored_public_key, expected_public_key_hex);
            info!("✅ Public key verified in Redis for task: {}", key);
            found_task = true;
            break;
        }
    }
    if !found_task {
        return Err("Could not find the test task in Redis to verify public key.".into());
    }


    // Step 9: Disconnect client
    info!("🔌 Disconnecting client...");
    client.disconnect().await;
    info!("✅ Client disconnected");
    
    // Step 10: Shutdown circles via launcher
    info!("🔌 Shutting down circles via launcher...");
    shutdown_circles(running_circles_store).await; // Pass the Vec<Arc<Mutex<RunningCircleInfo>>>
    info!("✅ Circles shut down.");

    info!("🎉 Full end-to-end authentication example completed successfully!");
    
    Ok(())
}