use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
use hirun::runtime;

fn hun_setup() {
    let _ = runtime::Builder::new()
        //.nth(4)
        .build();
}

static mut RT: Option<tokio::runtime::Runtime> = None;

fn tokio_setup<'a>() -> &'a tokio::runtime::Runtime {
    unsafe {
        //RT = Some(tokio::runtime::Runtime::new().unwrap());
        RT = Some(
            tokio::runtime::Builder::new_multi_thread()
                .enable_all()
                //.worker_threads(4)
                .build()
                .unwrap(),
        );
        RT.as_ref().unwrap()
    }
}

async fn tokio_value() -> i32 {
    let rt = unsafe { RT.as_ref().unwrap() };
    let mut handles = vec![];
    for _ in 0..10 {
        handles.push(rt.spawn(value()));
    }
    let mut val = 0;
    for handle in handles {
        val += handle.await.unwrap();
    }
    val
}

fn tokio_compute_sync(cnt: usize, rt: &tokio::runtime::Runtime) {
    let mut handles = vec![];
    for _ in 0..cnt {
        handles.push(rt.spawn(tokio_value()));
    }
    rt.block_on(async move {
        for handle in handles {
            let _ = handle.await;
        }
    });
}

fn tokio_compute_async(cnt: usize, rt: &tokio::runtime::Runtime) {
    rt.block_on(async {
        let mut set = tokio::task::JoinSet::new(); 
        for _ in 0..cnt {
            let _ = set.spawn(tokio_value());
        }

        while let Some(_) = set.join_next().await {
        }
    });
}

async fn value() -> i32 {
    1
}

async fn hun_value() -> i32 {
    let mut handles = vec![];
    for _ in 0..10 {
        handles.push(runtime::spawn(value()));
    }
    let mut val = 0;
    for handle in handles {
        val += handle.await.unwrap();
    }
    val
}

fn hun_compute_sync(cnt: usize) {
    let mut handles = vec![];
    for _ in 0..cnt {
        handles.push(runtime::spawn(hun_value()));
    }
    runtime::block_on(async move {
        for handle in handles {
            let _ = handle.await;
        }
    })
    .unwrap();
}

fn hun_compute_async(cnt: usize) {
    runtime::block_on(async move {
        let mut set = runtime::JoinSet::new();
        for _ in 0..cnt {
            let _ = set.spawn(hun_value());
        }
        let _ = set.wait_all().await;
    })
    .unwrap();
}

pub fn hun_benchmark(c: &mut Criterion) {
    let cnts = [100, 1000, 10000, 100000];
    //let cnts = [100, 1000];
    //let cnts = [100000];
    //let cnts = [100, 1000, 10000];
    let mut grp = c.benchmark_group("compute");
    for n in cnts {
        grp.bench_function(format!("hun_sync_{}", n), |b| {
            b.iter_batched(
                || hun_setup(),
                |_| hun_compute_sync(black_box(n)),
                BatchSize::SmallInput,
            )
        });
        grp.bench_function(format!("hirun_{}", n), |b| {
            b.iter_batched(
                || hun_setup(),
                |_| hun_compute_async(black_box(n)),
                BatchSize::SmallInput,
            )
        });
        grp.bench_function(format!("tokio_sync_{}", n), |b| {
            b.iter_batched(
                || tokio_setup(),
                |rt| tokio_compute_sync(black_box(n), rt),
                BatchSize::SmallInput,
            )
        });
        grp.bench_function(format!("tokio_async_{}", n), |b| {
            b.iter_batched(
                || tokio_setup(),
                |rt| tokio_compute_async(black_box(n), rt),
                BatchSize::SmallInput,
            )
        });
    }
    grp.finish();
}

criterion_group!(benches, hun_benchmark);
criterion_main!(benches);