//! A simple end-to-end attachments benchmark. It sets up a pair of Ditto instances,
//! and measures the latency between a peer creating and advertising an attachment through a new
//! document, and the opposing peer completing the fetching of said attachment after discovering it
//! through the advertising document.
//!
//! The measured section thus includes:
//! - Peer A creating the attachment from a random slice of bytes. This includes hashing the data
//!   and writing it to the blob-store (i.e., disk).
//! - Peer A creating and writing the advertising document to its Ditto store.
//! - Peer B receiving the document through its subscription observer.
//! - Peer B fetching the entire attachment.
//!
//! Note that this benchmark *does not* work on `in-memory-storage` backends, as attachments are
//! currently noops on these.
//!
//! The `DITTO_BENCH_TRANSPORT` and `DITTO_REPLICATION_OVER_NGN` environment variables can be used
//! to configure the transport used between the peers and whether to use NGN or legacy for doc _and_
//! attachment sync.

use std::{
    cell::UnsafeCell,
    sync::{Arc, Mutex},
};

use criterion::*;
use dittolive_ditto::prelude::{CborValueGetters, DittoAttachmentFetchEvent};

mod common;
use common::{generate_payload, setup};
use serde_json::json;

/// Only two-peers is supported by this benchmark for now, as attachments don't chain-sync like
/// documents do.
const PEERS: usize = 2;
const KB: usize = 1024;

fn simple(c: &mut Criterion) {
    let transport = std::env::var("DITTO_BENCH_TRANSPORT")
        .map(|s| {
            s.parse()
                .expect("Invalid transport type in DITTO_BENCH_TRANSPORT")
        })
        .unwrap_or_default();

    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap();

    let mut group = c.benchmark_group("attachments_e2e");

    let (_temp_root, dittos) = runtime.block_on(setup(PEERS, transport)).unwrap();
    let ditto = dittos.into_iter().map(Arc::new).collect::<Vec<_>>();

    let subs = ditto
        .iter()
        .map(|d| {
            d.sync().start().unwrap();
            d.sync()
                .register_subscription("SELECT * FROM COLLECTION att")
                .unwrap()
        })
        .collect::<Vec<_>>();

    let chunk_sizes_in_kb = [1usize, 64, 512, 1024, 4096, 8192, 16 * 1024];

    let mode = std::env::var("DITTO_REPLICATION_OVER_NGN").is_ok_and(|v| v == "true" || v == "1");

    let name = format!(
        "{}{}",
        transport.scheme(),
        if mode { "_ngn" } else { "_legacy" }
    );

    for payload_size in chunk_sizes_in_kb {
        let dst = Arc::clone(ditto.last().unwrap());
        let src = Arc::clone(ditto.first().unwrap());
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        let (ftx, frx) = tokio::sync::mpsc::unbounded_channel();
        let _observer = dst
            .store()
            .register_observer("SELECT * FROM COLLECTION att (at ATTACHMENT)", {
                let dst = Arc::clone(&dst);
                move |result| {
                    for item in result.iter() {
                        let att = item.value()["at"]
                            .as_object()
                            .expect("Couldn't deserialize Att")
                            .clone();
                        struct ProgressIndicator(
                            Mutex<Option<tokio::sync::mpsc::UnboundedSender<Result<(), ()>>>>,
                        );
                        impl ProgressIndicator {
                            fn send_complete(&self) {
                                if let Some(tx) = self.0.lock().unwrap().take() {
                                    let _ = tx.send(Ok(()));
                                }
                            }
                        }
                        impl Drop for ProgressIndicator {
                            fn drop(&mut self) {
                                if let Some(tx) = self.0.lock().unwrap().take() {
                                    let _ = tx.send(Err(()));
                                }
                            }
                        }
                        let tx = ProgressIndicator(Mutex::new(Some(tx.clone())));
                        let _fetcher = dst
                            .store()
                            .fetch_attachment(att, move |event| {
                                if let DittoAttachmentFetchEvent::Completed { attachment: _ } =
                                    event
                                {
                                    tx.send_complete();
                                }
                            })
                            .expect("Couldn't start attachment fetcher");
                        // println!("Started fetcher for attachment");
                        _ = ftx.send(_fetcher);
                    }
                }
            })
            .unwrap();

        // this isn't quite accurate but it's close enough (we don't account for id or document
        // wrapper overhead)
        group.throughput(Throughput::Bytes((payload_size * KB) as _));

        // To fit the signature of iter_batched, we need to wrap `rx` and `frx` in a something that
        // appears `'static`, despite the guarantee that all clones of these Arcs will be dropped
        // by the time `bench_function` returns.
        // We use `UnsafeCell` to allow each iteration of the `routine` to recover a mutable
        // reference to the receivers, as `routine` is never run concurrently (it _is_
        // `FnMut`).
        // All of this is needed because `routine` isn't a proper async closure (those didn't exist
        // back when criterion introduced this API), but a closure returning a new future.
        #[allow(clippy::arc_with_non_send_sync)]
        let rx = Arc::new(UnsafeCell::new(rx));
        #[allow(clippy::arc_with_non_send_sync)]
        let frx = Arc::new(UnsafeCell::new(frx));
        group.bench_function(BenchmarkId::new(&name, format!("{payload_size}KiB")), |b| {
            b.to_async(&runtime).iter_batched(
                {
                    || {
                        (
                            Arc::clone(&src),
                            generate_payload(payload_size * KB),
                            Arc::clone(&rx),
                            Arc::clone(&frx),
                        )
                    }
                },
                |(src, data, rx, frx)| async move {
                    let att = src
                        .store()
                        .new_attachment_from_bytes(&data, Default::default())
                        .await
                        .expect("Couldn't create attachment");

                    src.store()
                        .execute((
                            "INSERT INTO COLLECTION att (at ATTACHMENT) VALUES (:att)",
                            json!({"att": {"at": att}}),
                        ))
                        .await
                        .unwrap();

                    let _fetcher = unsafe { &mut *frx.get() }
                        .recv()
                        .await
                        .expect("Failed to receive completion signal");

                    // println!("Attachment fetcher started, waiting for completion");

                    unsafe { &mut *rx.get() }
                        .recv()
                        .await
                        .expect("Failed to receive completion signal")
                        .expect("Attachment fetcher was closed");
                },
                criterion::BatchSize::LargeInput,
            );
        });
    }

    println!("Finished benchmarks, should return now");
    group.finish();
    println!("Finished group");
    for sub in subs {
        sub.cancel();
    }
    for d in &ditto {
        println!("Stopping ditto");
        d.sync().stop();
        println!("Stopped ditto");
    }
}

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