#![allow(deprecated)]
use std::sync::{Arc, Mutex};

use criterion::*;
use dittolive_ditto::experimental::bus::*;
use rand::Rng;
use safer_ffi::bytes::Bytes;

mod common;
use common::{pubkey, setup};

/// Generate a random payload of `len` bytes. It's important to use random data
/// so that the compression middleware in the transports doesn't skew the results.
fn generate_payload(len: usize) -> Bytes<'static> {
    let mut rng = rand::thread_rng();
    let mut payload = vec![0u8; len];
    rng.fill(&mut payload[..]);

    payload.into()
}

fn stream(c: &mut Criterion) {
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap();

    let mut group = c.benchmark_group("bus_throughput");
    const N: usize = 3;
    let ditto = runtime.block_on(setup(N)).unwrap();
    let pk2 = pubkey(&ditto[N - 1]);

    let src = ditto[0].bus();
    let dst = ditto[N - 1].bus();
    let acceptor = dst
        .bind_topic("throughput")
        .on_receive_factory(tokio::sync::mpsc::unbounded_channel)
        .finish(std::sync::mpsc::channel())
        .unwrap();
    let mut tx = None;
    runtime.block_on(async {
        for i in (0..100).rev() {
            match src
                .connect(pk2.clone(), "throughput")
                .on_receive_factory(WriteOnly)
                .finish_async()
                .await
            {
                Ok(t) => {
                    tx = Some(t);
                    break;
                }
                Err(e) => {
                    println!(
                        "Failed to establish stream (reattempting in 100ms, {i} attempts \
                         remaining): {:?}",
                        e
                    );
                    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
                    continue;
                }
            }
        }
    });
    let tx = tx.expect("Failed to establish stream after 100 attempts");
    let rx = acceptor.recv().unwrap();
    let rx = Arc::new(Mutex::new(rx));

    let chunk_sizes = [
        1,  // Mostly sending overhead
        64, // Small user payloads
        512,
        1024,
        4096,
        8192, // Peak performance at time of this comment: 16KiB chunks halve throughput
        16 * 1024, // 16KiB chunks already start degrading performance, dividing it by 4 compared to 8KiB chunks at time of writing
                   // 1024 * 1024, // 1MiB, benching this takes 20 minutes at time of writing, likely due to compression scaling poorly
    ];
    for payload_size in chunk_sizes {
        group.throughput(Throughput::Bytes(payload_size as _));

        group.bench_function(format!("stream_{}_bytes", payload_size), |b| {
            b.to_async(&runtime).iter_batched(
                || (tx.clone(), rx.clone(), generate_payload(payload_size)),
                |(tx, rx, message)| {
                    async move {
                        let send = tx.message(message).send().changed();

                        let mut guard = rx.lock().unwrap();
                        let recv = guard.recv();

                        // await both in parallel - we're not finished until both are done.
                        let (_sent, _rect) = tokio::join!(send, recv);
                    }
                },
                criterion::BatchSize::LargeInput,
            );
        });
    }

    group.finish();
}

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