use hero_supervisor::{
    Supervisor, SupervisorBuilder, WorkerConfig, WorkerLifecycleManager, 
    WorkerLifecycleManagerBuilder, ScriptType
};
use log::{info, warn, error};
use std::collections::HashMap;
use std::path::PathBuf;
use std::time::Duration;
use tokio::time::sleep;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    env_logger::init();
    
    info!("Starting Worker Lifecycle Management Demo");
    
    // Configuration
    let redis_url = "redis://localhost:6379";
    let zinit_socket = "/var/run/zinit.sock";
    
    // Create supervisor
    let supervisor = SupervisorBuilder::new()
        .redis_url(redis_url)
        .caller_id("lifecycle_demo")
        .context_id("demo_context")
        .build()?;
    
    // Configure workers for different script types
    let mut worker_configs = Vec::new();
    
    // OSIS workers (Rhai/HeroScript)
    for i in 0..2 {
        let config = WorkerConfig::new(
            format!("osis_worker_{}", i),
            PathBuf::from("/usr/local/bin/osis_worker"),
            ScriptType::OSIS,
        )
        .with_args(vec![
            "--redis-url".to_string(),
            redis_url.to_string(),
            "--worker-id".to_string(),
            format!("osis_worker_{}", i),
        ])
        .with_env({
            let mut env = HashMap::new();
            env.insert("RUST_LOG".to_string(), "info".to_string());
            env.insert("WORKER_TYPE".to_string(), "osis".to_string());
            env
        })
        .with_health_check("/usr/local/bin/osis_worker --health-check".to_string())
        .with_dependencies(vec!["redis".to_string()]);
        
        worker_configs.push(config);
    }
    
    // SAL workers (System Abstraction Layer)
    for i in 0..3 {
        let config = WorkerConfig::new(
            format!("sal_worker_{}", i),
            PathBuf::from("/usr/local/bin/sal_worker"),
            ScriptType::SAL,
        )
        .with_args(vec![
            "--redis-url".to_string(),
            redis_url.to_string(),
            "--worker-id".to_string(),
            format!("sal_worker_{}", i),
        ])
        .with_env({
            let mut env = HashMap::new();
            env.insert("RUST_LOG".to_string(), "info".to_string());
            env.insert("WORKER_TYPE".to_string(), "sal".to_string());
            env
        })
        .with_health_check("/usr/local/bin/sal_worker --health-check".to_string())
        .with_dependencies(vec!["redis".to_string()]);
        
        worker_configs.push(config);
    }
    
    // V workers (HeroScript in V language)
    for i in 0..2 {
        let config = WorkerConfig::new(
            format!("v_worker_{}", i),
            PathBuf::from("/usr/local/bin/v_worker"),
            ScriptType::V,
        )
        .with_args(vec![
            "--redis-url".to_string(),
            redis_url.to_string(),
            "--worker-id".to_string(),
            format!("v_worker_{}", i),
        ])
        .with_env({
            let mut env = HashMap::new();
            env.insert("RUST_LOG".to_string(), "info".to_string());
            env.insert("WORKER_TYPE".to_string(), "v".to_string());
            env
        })
        .with_health_check("/usr/local/bin/v_worker --health-check".to_string())
        .with_dependencies(vec!["redis".to_string()]);
        
        worker_configs.push(config);
    }
    
    // Create lifecycle manager
    let mut lifecycle_manager = WorkerLifecycleManagerBuilder::new(zinit_socket.to_string())
        .with_supervisor(supervisor.clone());
    
    // Add all worker configurations
    for config in worker_configs {
        lifecycle_manager = lifecycle_manager.add_worker(config);
    }
    
    let mut lifecycle_manager = lifecycle_manager.build();
    
    // Demonstrate lifecycle operations
    info!("=== Starting Worker Lifecycle Demo ===");
    
    // 1. Start all workers
    info!("1. Starting all workers...");
    match lifecycle_manager.start_all_workers().await {
        Ok(_) => info!("✅ All workers started successfully"),
        Err(e) => {
            error!("❌ Failed to start workers: {}", e);
            return Err(e.into());
        }
    }
    
    // Wait for workers to initialize
    sleep(Duration::from_secs(5)).await;
    
    // 2. Check worker status
    info!("2. Checking worker status...");
    match lifecycle_manager.get_all_worker_status().await {
        Ok(status_map) => {
            for (worker_name, status) in status_map {
                info!("   Worker '{}': State={:?}, PID={}", worker_name, status.state, status.pid);
            }
        }
        Err(e) => warn!("Failed to get worker status: {}", e),
    }
    
    // 3. Demonstrate scaling
    info!("3. Demonstrating worker scaling...");
    
    // Scale up OSIS workers
    info!("   Scaling up OSIS workers to 3...");
    if let Err(e) = lifecycle_manager.scale_workers(&ScriptType::OSIS, 3).await {
        warn!("Failed to scale OSIS workers: {}", e);
    }
    
    sleep(Duration::from_secs(3)).await;
    
    // Scale down SAL workers
    info!("   Scaling down SAL workers to 1...");
    if let Err(e) = lifecycle_manager.scale_workers(&ScriptType::SAL, 1).await {
        warn!("Failed to scale SAL workers: {}", e);
    }
    
    sleep(Duration::from_secs(3)).await;
    
    // 4. Check running worker counts
    info!("4. Checking running worker counts after scaling...");
    for script_type in [ScriptType::OSIS, ScriptType::SAL, ScriptType::V] {
        let count = lifecycle_manager.get_running_worker_count(&script_type).await;
        info!("   {:?}: {} workers running", script_type, count);
    }
    
    // 5. Demonstrate restart functionality
    info!("5. Demonstrating worker restart...");
    if let Err(e) = lifecycle_manager.restart_worker("osis_worker_0").await {
        warn!("Failed to restart worker: {}", e);
    } else {
        info!("   ✅ Successfully restarted osis_worker_0");
    }
    
    sleep(Duration::from_secs(3)).await;
    
    // 6. Simulate job dispatch and health monitoring
    info!("6. Simulating job dispatch and health monitoring...");
    
    // Update job time for a worker (simulating job dispatch)
    lifecycle_manager.update_worker_job_time("sal_worker_0");
    info!("   Updated job time for sal_worker_0");
    
    // Perform health monitoring check
    if let Err(e) = lifecycle_manager.monitor_worker_health().await {
        warn!("Health monitoring failed: {}", e);
    } else {
        info!("   ✅ Health monitoring completed");
    }
    
    // 7. Create and execute a test job
    info!("7. Creating and executing a test job...");
    let test_job = supervisor
        .new_job()
        .script_type(ScriptType::OSIS)
        .script_content("println!(\"Hello from worker!\");".to_string())
        .timeout(Duration::from_secs(30))
        .build()?;
    
    match supervisor.run_job_and_await_result(&test_job).await {
        Ok(result) => info!("   ✅ Job executed successfully: {}", result),
        Err(e) => warn!("   ❌ Job execution failed: {}", e),
    }
    
    // 8. Demonstrate graceful shutdown
    info!("8. Demonstrating graceful shutdown...");
    
    // Stop specific workers
    info!("   Stopping specific workers...");
    for worker_name in ["osis_worker_1", "v_worker_0"] {
        if let Err(e) = lifecycle_manager.stop_worker(worker_name).await {
            warn!("Failed to stop worker {}: {}", worker_name, e);
        } else {
            info!("   ✅ Stopped worker: {}", worker_name);
        }
    }
    
    sleep(Duration::from_secs(2)).await;
    
    // Stop all remaining workers
    info!("   Stopping all remaining workers...");
    if let Err(e) = lifecycle_manager.stop_all_workers().await {
        error!("Failed to stop all workers: {}", e);
    } else {
        info!("   ✅ All workers stopped successfully");
    }
    
    info!("=== Worker Lifecycle Demo Completed ===");
    
    // Optional: Start health monitoring loop (commented out for demo)
    // info!("Starting health monitoring loop (Ctrl+C to stop)...");
    // lifecycle_manager.start_health_monitoring().await;
    
    Ok(())
}