//! Throughput benchmarks for OmniVox v2.1, answering: what does the
//! protocol cost versus raw loopback TCP on this machine?
//!
//! - `crypto`: criterion stats for frame seal/open at three sizes.
//! - `transport`: hand-rolled async harness — spray (small messages/s),
//!   dump (GiB/s at 1 MiB stream frames), and a raw-TCP echo baseline.
//!
//! Run: `cargo bench --bench transport`

use criterion::{BenchmarkId, Criterion, Throughput, criterion_group};
use omnivox::{Config, FrameType, frame::seal_into};
use serde::{Deserialize, Serialize};
use std::{
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
    time::{Duration, Instant},
};

#[derive(Serialize, Deserialize)]
struct M(#[allow(dead_code)] Vec<u8>);

fn crypto_benches(c: &mut Criterion) {
    let mut g = c.benchmark_group("crypto");
    let key = [1u8; 32];
    for size in [256usize, 4096, 1 << 20] {
        let payload = vec![7u8; size];
        g.throughput(Throughput::Bytes(size as u64));
        g.bench_function(BenchmarkId::new("seal", size), |b| {
            let mut buf = Vec::with_capacity(size + 64);
            b.iter(|| {
                seal_into(
                    &key,
                    1,
                    FrameType::Message,
                    std::hint::black_box(&payload),
                    &mut buf,
                )
            });
        });
        // open measured as copy + decrypt-in-place of a sealed frame
        let sealed = omnivox::frame::seal(&key, 1, FrameType::Message, &payload);
        g.bench_function(BenchmarkId::new("open", size), |b| {
            b.iter(|| {
                let mut body = std::hint::black_box(&sealed[4..]).to_vec();
                let (kind, _pt) = omnivox::frame::open_in_place(&key, 1, &mut body).unwrap();
                std::hint::black_box(kind)
            });
        });
    }
    g.finish();
}

fn test_config(stream_frame: u32) -> Config {
    Config {
        idle_timeout_ms: 60_000,
        body_stall_ms: 60_000,
        send_timeout_ms: 60_000,
        max_message_size: 2 * 1024 * 1024,
        max_stream_frame: stream_frame,
        allow_insecure_defaults: true,
        ..Config::default()
    }
}

fn transport_harness() {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap();
    rt.block_on(async {
        // ---- spray: small messages per second ----
        const SPRAY_PORT: u16 = 18701;
        let rx = Arc::new(AtomicU64::new(0));
        let r = Arc::clone(&rx);
        let server = omnivox::server::<M>()
            .key([7u8; 32])
            .config(test_config(1 << 20))
            .on_data(move |_: &M, _| {
                r.fetch_add(1, Ordering::Relaxed);
            })
            .serve(SPRAY_PORT)
            .await
            .unwrap();
        let client = omnivox::client::<M>()
            .key([7u8; 32])
            .config(test_config(1 << 20))
            .on_data(|_: &M, _| {})
            .connect("127.0.0.1", SPRAY_PORT)
            .await
            .unwrap();

        let msg = M(vec![9u8; 256]);
        let n: u64 = 200_000;
        let start = Instant::now();
        let mut sent = 0;
        while sent < n {
            if client.send(&msg).await {
                sent += 1;
            } else {
                break;
            }
        }
        // drain
        while rx.load(Ordering::Relaxed) < sent {
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
        let secs = start.elapsed().as_secs_f64();
        println!(
            "spray:        {:>10.0} msg/s  ({} x 256 B, encrypted round trip loopback)",
            sent as f64 / secs,
            sent
        );
        client.stop().await;
        server.stop().await;

        // ---- dump: stream GiB/s ----
        const DUMP_PORT: u16 = 18702;
        let rx = Arc::new(AtomicU64::new(0));
        let r = Arc::clone(&rx);
        let server = omnivox::server::<M>()
            .key([7u8; 32])
            .config(test_config(1 << 20))
            .on_stream(move |b, _| {
                r.fetch_add(b.len() as u64, Ordering::Relaxed);
            })
            .serve(DUMP_PORT)
            .await
            .unwrap();
        let client = omnivox::client::<M>()
            .key([7u8; 32])
            .config(test_config(1 << 20))
            .on_data(|_: &M, _| {})
            .connect("127.0.0.1", DUMP_PORT)
            .await
            .unwrap();

        let page = vec![3u8; 1 << 20]; // 1 MiB
        let pages: u64 = 256; // 256 MiB
        let start = Instant::now();
        for _ in 0..pages {
            client.send_stream(&page).await;
        }
        let want = pages << 20;
        while rx.load(Ordering::Relaxed) < want {
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
        let secs = start.elapsed().as_secs_f64();
        println!(
            "dump:         {:>10.2} GiB/s  ({} MiB over encrypted stream lane, loopback)",
            want as f64 / secs / (1u64 << 30) as f64,
            want >> 20
        );
        client.stop().await;
        server.stop().await;

        // ---- raw loopback baseline (no protocol) ----
        let listener = tokio::net::TcpListener::bind(("127.0.0.1", 18703))
            .await
            .unwrap();
        let echo = tokio::spawn(async move {
            let (s, _) = listener.accept().await.unwrap();
            let (mut r, mut w) = s.into_split();
            let mut buf = vec![0u8; 1 << 20];
            loop {
                use tokio::io::{AsyncReadExt, AsyncWriteExt};
                if r.read_exact(&mut buf).await.is_err() {
                    break;
                }
                if w.write_all(&buf).await.is_err() {
                    break;
                }
            }
        });
        let raw = tokio::net::TcpStream::connect(("127.0.0.1", 18703))
            .await
            .unwrap();
        raw.set_nodelay(true).unwrap();
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        let (mut rr, mut rw) = raw.into_split();
        let chunk: Vec<u8> = vec![5u8; 1 << 20];
        let total: u64 = 64 << 20;
        let reader = tokio::spawn(async move {
            let mut got = 0u64;
            let mut sink = vec![0u8; 1 << 20];
            while got < total {
                let n = rr.read(&mut sink).await.unwrap();
                if n == 0 {
                    break;
                }
                got += n as u64;
            }
            got
        });
        let start = Instant::now();
        let mut written = 0u64;
        while written < total {
            rw.write_all(&chunk).await.unwrap();
            written += chunk.len() as u64;
        }
        let want = total;
        while !reader.is_finished() {
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
        let got = reader.await.unwrap();
        assert!(got >= want, "raw echo short: {got}");
        let secs = start.elapsed().as_secs_f64();
        println!(
            "raw baseline: {:>10.2} GiB/s  (unencrypted loopback write+echo, same box)",
            want as f64 / secs / (1u64 << 30) as f64
        );
        echo.abort();
    });
}

criterion_group! {
    name = crypto;
    config = Criterion::default().sample_size(20).measurement_time(Duration::from_secs(2));
    targets = crypto_benches
}

fn main() {
    crypto();
    transport_harness();
}