use std::hint::black_box;

use criterion::{criterion_group, criterion_main, Criterion};
use vapora_swarm::{AgentProfile, SwarmCoordinator};

fn setup_swarm_with_agents(count: usize) -> SwarmCoordinator {
    let coordinator = SwarmCoordinator::new();

    for i in 0..count {
        let profile = AgentProfile {
            id: format!("agent-{}", i),
            roles: vec![match i % 3 {
                0 => "developer".to_string(),
                1 => "reviewer".to_string(),
                _ => "architect".to_string(),
            }],
            capabilities: vec![
                "coding".to_string(),
                "analysis".to_string(),
                format!("domain-{}", i % 5),
            ],
            current_load: 0.3 + (i as f64 * 0.05) % 0.7,
            success_rate: 0.7 + (i as f64 * 0.02),
            availability: i % 5 != 0,
        };
        coordinator.register_agent(profile).ok();
    }

    coordinator
}

fn coordinator_register_agent(c: &mut Criterion) {
    c.bench_function("register_single_agent", |b| {
        b.iter(|| {
            let coordinator = SwarmCoordinator::new();
            let profile = AgentProfile {
                id: black_box("agent-1".to_string()),
                roles: vec!["developer".to_string()],
                capabilities: vec!["coding".to_string()],
                current_load: 0.5,
                success_rate: 0.9,
                availability: true,
            };
            black_box(coordinator.register_agent(black_box(profile)))
        });
    });
}

fn coordinator_task_assignment(c: &mut Criterion) {
    c.bench_function("submit_task_for_bidding_50_agents", |b| {
        b.to_async(tokio::runtime::Runtime::new().unwrap())
            .iter_batched(
                || setup_swarm_with_agents(50),
                |coordinator| async move {
                    black_box(
                        coordinator
                            .submit_task_for_bidding(
                                black_box("task-1".to_string()),
                                black_box("Develop feature X".to_string()),
                                black_box(vec!["coding".to_string()]),
                            )
                            .await,
                    )
                },
                criterion::BatchSize::SmallInput,
            );
    });
}

fn coordinator_coalition_formation(c: &mut Criterion) {
    c.bench_function("create_coalition_100_agents", |b| {
        b.to_async(tokio::runtime::Runtime::new().unwrap())
            .iter_batched(
                || setup_swarm_with_agents(100),
                |coordinator| async move {
                    black_box(
                        coordinator
                            .create_coalition(
                                black_box("coordinator-1".to_string()),
                                black_box(vec![
                                    "developer".to_string(),
                                    "reviewer".to_string(),
                                    "architect".to_string(),
                                ]),
                            )
                            .await,
                    )
                },
                criterion::BatchSize::SmallInput,
            );
    });
}

fn coordinator_update_status(c: &mut Criterion) {
    c.bench_function("update_agent_status_200_agents", |b| {
        b.to_async(tokio::runtime::Runtime::new().unwrap())
            .iter_batched(
                || setup_swarm_with_agents(200),
                |coordinator| async move {
                    for i in 0..50 {
                        black_box(coordinator.update_agent_status(
                            black_box(&format!("agent-{}", i)),
                            black_box(0.5 + (i as f64 * 0.01) % 0.4),
                            black_box(i % 10 != 0),
                        ))
                        .ok();
                    }
                    coordinator
                },
                criterion::BatchSize::SmallInput,
            );
    });
}

fn coordinator_consensus_vote(c: &mut Criterion) {
    c.bench_function("consensus_vote_50_voters", |b| {
        b.to_async(tokio::runtime::Runtime::new().unwrap())
            .iter(|| async {
                let coordinator = SwarmCoordinator::new();
                let mut votes = Vec::new();
                for i in 0..50 {
                    let vote = if i < 30 {
                        vapora_swarm::Vote::Agree
                    } else if i < 40 {
                        vapora_swarm::Vote::Disagree
                    } else {
                        vapora_swarm::Vote::Abstain
                    };
                    votes.push((format!("voter-{}", i), vote));
                }
                black_box(
                    coordinator
                        .consensus_vote(black_box("proposal-1"), black_box(votes))
                        .await,
                )
            });
    });
}

fn coordinator_get_stats(c: &mut Criterion) {
    c.bench_function("get_swarm_stats_500_agents", |b| {
        b.iter_batched(
            || setup_swarm_with_agents(500),
            |coordinator| black_box(coordinator.get_swarm_stats()),
            criterion::BatchSize::SmallInput,
        );
    });
}

criterion_group!(
    benches,
    coordinator_register_agent,
    coordinator_task_assignment,
    coordinator_coalition_formation,
    coordinator_update_status,
    coordinator_consensus_vote,
    coordinator_get_stats
);
criterion_main!(benches);
