use std::{sync::{atomic::{AtomicBool, Ordering}, Arc}, thread};

use qanat::broadcast;
use criterion::{criterion_group, criterion_main, Criterion};
use futures_lite::future::block_on;

#[allow(unused_must_use)]
pub fn broadcast_read_write(c: &mut Criterion) {
    let (mut qs, mut qr1) = broadcast::channel(1.try_into().unwrap());

    let mut n = 0;

    let mut qr2 = qr1.clone();
    let mut qr3 = qr1.clone();
    let mut qr4 = qr1.clone();
    let mut qr5 = qr1.clone();
    let mut qr6 = qr1.clone();
    let mut qr7 = qr1.clone();
    let mut qr8 = qr1.clone();

    c.bench_function("qanat::broadcast read write", |b| {
        b.iter(|| {
            block_on(async {
                qs.send(n);
                assert_eq!(qr1.recv().await.unwrap(), n);
                assert_eq!(qr2.recv().await.unwrap(), n);
                assert_eq!(qr3.recv().await.unwrap(), n);
                assert_eq!(qr4.recv().await.unwrap(), n);
                assert_eq!(qr5.recv().await.unwrap(), n);
                assert_eq!(qr6.recv().await.unwrap(), n);
                assert_eq!(qr7.recv().await.unwrap(), n);
                assert_eq!(qr8.recv().await.unwrap(), n);
                n += 1;
            })
        })
    });


    let (s, mut r1) = async_broadcast::broadcast(1);

    let mut n = 0;

    let mut r2 = r1.clone();
    let mut r3 = r1.clone();
    let mut r4 = r1.clone();
    let mut r5 = r1.clone();
    let mut r6 = r1.clone();
    let mut r7 = r1.clone();
    let mut r8 = r1.clone();

    c.bench_function("async_broadcast read write", |b| {
        b.iter(|| {
            block_on(async {
                s.broadcast(n).await;
                assert_eq!(r1.recv().await.unwrap(), n);
                assert_eq!(r2.recv().await.unwrap(), n);
                assert_eq!(r3.recv().await.unwrap(), n);
                assert_eq!(r4.recv().await.unwrap(), n);
                assert_eq!(r5.recv().await.unwrap(), n);
                assert_eq!(r6.recv().await.unwrap(), n);
                assert_eq!(r7.recv().await.unwrap(), n);
                assert_eq!(r8.recv().await.unwrap(), n);
                n += 1;
            })
        })
    });
}

#[allow(unused_must_use)]
pub fn broadcast_write(c: &mut Criterion) {
    let (mut qs, _qr1) = broadcast::channel(1.try_into().unwrap());
    let mut n = 0;

    c.bench_function("qanat::broadcast write", |b| {
        b.iter(|| {
            block_on(async {
                qs.send(n);
                n += 1;
            })
        })
    });

    let (mut s, _r1) = async_broadcast::broadcast(1);
    s.set_overflow(true);
    let mut n = 0;

    c.bench_function("async_broadcast write", |b| {
        b.iter(|| {
            block_on(async {
                s.broadcast(n).await;
                n += 1;
            })
        })
    });
}

#[allow(unused_must_use)]
pub fn broadcast_contended_write(c: &mut Criterion) {
    thread::scope(|s| {
        let (mut qs, mut _qr1) = broadcast::channel(20.try_into().unwrap());
        let stop = Arc::new(AtomicBool::new(false));

        // Start some background noise threads
        for _ in 0..1 {
            let mut cqs = qs.clone();
            let st = stop.clone();
            s.spawn(move || {
                let mut n = 0;
                while !st.load(Ordering::Relaxed) {
                    cqs.send(n);
                    n += 1;
                }
            });
        }

        let mut n = 0;
        c.bench_function("qanat::broadcast contended write", |b| {
            b.iter(|| {
                qs.send(n);
                n += 1;
            })
        });

        stop.store(true, Ordering::Relaxed);
    });

    thread::scope(|ts| {
        let (mut s, mut _r1) = async_broadcast::broadcast(20);
        s.set_overflow(true);
        let stop = Arc::new(AtomicBool::new(false));

        // Start some background noise threads
        for _ in 0..1 {
            let cs = s.clone();
            let st = stop.clone();
            ts.spawn(move || block_on(async {
                let mut n = 0;
                while !st.load(Ordering::Relaxed) {
                    cs.broadcast(n).await;
                    n += 1;
                }
            }));
        }

        let mut n = 0;
        c.bench_function("async_broadcast contended write", |b| {
            b.iter(|| block_on(async {
                s.broadcast(n).await;
                n += 1;
            }))
        });

        stop.store(true, Ordering::Relaxed);
    });
}

criterion_group!(benches, broadcast_write, broadcast_contended_write, broadcast_read_write);
criterion_main!(benches);
