use criterion::{black_box, criterion_group, criterion_main, Criterion, BenchmarkId};
use hero_supervisor_openrpc_client::SupervisorClientBuilder;
use hero_job::Job;
use tokio::runtime::Runtime;
use std::time::Duration;
use std::collections::HashMap;
use uuid::Uuid;
use chrono::Utc;

/// Benchmark configuration
const SUPERVISOR_URL: &str = "http://127.0.0.1:3030";
const ADMIN_SECRET: &str = "SECRET";

/// Helper to create a tokio runtime for benchmarks
fn create_runtime() -> Runtime {
    Runtime::new().unwrap()
}

/// Helper to create a test job
fn create_test_job(runner: &str, command: &str, args: Vec<String>) -> Job {
    Job {
        id: Uuid::new_v4().to_string(),
        caller_id: "benchmark".to_string(),
        context_id: "test".to_string(),
        payload: serde_json::json!({
            "command": command,
            "args": args
        }).to_string(),
        runner: runner.to_string(),
        timeout: 30,
        env_vars: HashMap::new(),
        created_at: Utc::now(),
        updated_at: Utc::now(),
        signatures: vec![],
    }
}

/// Stress test: High-frequency job submissions
fn stress_high_frequency_jobs(c: &mut Criterion) {
    let rt = create_runtime();
    
    let client = rt.block_on(async {
        SupervisorClientBuilder::new()
            .url(SUPERVISOR_URL)
            .secret(ADMIN_SECRET)
            .timeout(Duration::from_secs(120))
            .build()
            .expect("Failed to create supervisor client")
    });

    // Ensure runner is registered
    rt.block_on(async {
        let _ = client.runner_create("hero").await;
    });

    let mut group = c.benchmark_group("stress_high_frequency");
    group.sample_size(10); // Fewer samples for stress tests
    group.measurement_time(Duration::from_secs(20));
    
    for num_jobs in [50, 100, 200].iter() {
        group.bench_with_input(
            BenchmarkId::from_parameter(num_jobs),
            num_jobs,
            |b, &num_jobs| {
                b.to_async(&rt).iter(|| async {
                    let mut handles = vec![];
                    
                    for i in 0..num_jobs {
                        let client = client.clone();
                        let handle = tokio::spawn(async move {
                            let job = create_test_job("hero", "echo", vec![format!("stress_{}", i)]);
                            client.job_create(job).await
                        });
                        handles.push(handle);
                    }
                    
                    // Wait for all jobs to be submitted
                    for handle in handles {
                        let _ = black_box(handle.await);
                    }
                });
            },
        );
    }
    
    group.finish();
}

/// Stress test: Sustained load over time
fn stress_sustained_load(c: &mut Criterion) {
    let rt = create_runtime();
    
    let client = rt.block_on(async {
        SupervisorClientBuilder::new()
            .url(SUPERVISOR_URL)
            .secret(ADMIN_SECRET)
            .timeout(Duration::from_secs(120))
            .build()
            .expect("Failed to create supervisor client")
    });

    // Ensure runner is registered
    rt.block_on(async {
        let _ = client.runner_create("hero").await;
    });

    let mut group = c.benchmark_group("stress_sustained_load");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(30));
    
    group.bench_function("continuous_submissions", |b| {
        b.to_async(&rt).iter(|| async {
            // Submit jobs continuously for the measurement period
            for i in 0..20 {
                let job = create_test_job("hero", "echo", vec![format!("sustained_{}", i)]);
                let _ = black_box(client.job_create(job).await);
            }
        });
    });
    
    group.finish();
}

/// Stress test: Large payload handling
fn stress_large_payloads(c: &mut Criterion) {
    let rt = create_runtime();
    
    let client = rt.block_on(async {
        SupervisorClientBuilder::new()
            .url(SUPERVISOR_URL)
            .secret(ADMIN_SECRET)
            .timeout(Duration::from_secs(120))
            .build()
            .expect("Failed to create supervisor client")
    });

    // Ensure runner is registered
    rt.block_on(async {
        let _ = client.runner_create("hero").await;
    });

    let mut group = c.benchmark_group("stress_large_payloads");
    group.sample_size(10);
    
    for size_kb in [1, 10, 100].iter() {
        group.bench_with_input(
            BenchmarkId::from_parameter(format!("{}KB", size_kb)),
            size_kb,
            |b, &size_kb| {
                b.to_async(&rt).iter(|| async {
                    // Create a large payload
                    let large_data = "x".repeat(size_kb * 1024);
                    let job = create_test_job("hero", "echo", vec![large_data]);
                    black_box(client.job_create(job).await.expect("Job create failed"))
                });
            },
        );
    }
    
    group.finish();
}

/// Stress test: Rapid API calls
fn stress_rapid_api_calls(c: &mut Criterion) {
    let rt = create_runtime();
    
    let client = rt.block_on(async {
        SupervisorClientBuilder::new()
            .url(SUPERVISOR_URL)
            .secret(ADMIN_SECRET)
            .build()
            .expect("Failed to create supervisor client")
    });

    let mut group = c.benchmark_group("stress_rapid_api");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(15));
    
    group.bench_function("rapid_info_calls", |b| {
        b.to_async(&rt).iter(|| async {
            // Make 100 rapid API calls
            for _ in 0..100 {
                let _ = black_box(client.get_supervisor_info().await);
            }
        });
    });
    
    group.bench_function("rapid_list_calls", |b| {
        b.to_async(&rt).iter(|| async {
            // Make 100 rapid list calls
            for _ in 0..100 {
                let _ = black_box(client.runner_list().await);
            }
        });
    });
    
    group.finish();
}

/// Stress test: Mixed workload
fn stress_mixed_workload(c: &mut Criterion) {
    let rt = create_runtime();
    
    let client = rt.block_on(async {
        SupervisorClientBuilder::new()
            .url(SUPERVISOR_URL)
            .secret(ADMIN_SECRET)
            .timeout(Duration::from_secs(120))
            .build()
            .expect("Failed to create supervisor client")
    });

    // Ensure runner is registered
    rt.block_on(async {
        let _ = client.runner_create("hero").await;
    });

    let mut group = c.benchmark_group("stress_mixed_workload");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(25));
    
    group.bench_function("mixed_operations", |b| {
        b.to_async(&rt).iter(|| async {
            let mut handles = vec![];
            
            // Mix of different operations
            for i in 0..10 {
                let client = client.clone();
                
                // Job submission
                let handle1 = tokio::spawn(async move {
                    let job = create_test_job("hero", "echo", vec![format!("mixed_{}", i)]);
                    client.job_create(job).await.map(|_| ())
                });
                handles.push(handle1);
            }
            
            // Wait for all operations
            for handle in handles {
                let _ = black_box(handle.await);
            }
        });
    });
    
    group.finish();
}

/// Stress test: Connection pool exhaustion
fn stress_connection_pool(c: &mut Criterion) {
    let rt = create_runtime();

    let mut group = c.benchmark_group("stress_connection_pool");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(20));
    
    for num_clients in [10, 50, 100].iter() {
        group.bench_with_input(
            BenchmarkId::from_parameter(num_clients),
            num_clients,
            |b, &num_clients| {
                b.to_async(&rt).iter(|| async {
                    let mut handles = vec![];
                    
                    // Create many clients and make concurrent requests
                    for _ in 0..num_clients {
                        let handle = tokio::spawn(async move {
                            let client = SupervisorClientBuilder::new()
                                .url(SUPERVISOR_URL)
                                .secret(ADMIN_SECRET)
                                .build()
                                .expect("Failed to create client");
                            
                            client.get_supervisor_info().await
                        });
                        handles.push(handle);
                    }
                    
                    // Wait for all requests
                    for handle in handles {
                        let _ = black_box(handle.await);
                    }
                });
            },
        );
    }
    
    group.finish();
}

criterion_group!(
    stress_tests,
    stress_high_frequency_jobs,
    stress_sustained_load,
    stress_large_payloads,
    stress_rapid_api_calls,
    stress_mixed_workload,
    stress_connection_pool,
);

criterion_main!(stress_tests);