//! Reproducible loopback comparison. Not a claim about production networks.
use anyhow::Result;
use mirror_tool::{
    config::{NetworkConfig, ReplicaPolicy, ServiceConfig, Team},
    engine::Engine,
};
use serde_json::json;
use std::{net::SocketAddr, time::Instant};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
};

async fn transfer(addr: SocketAddr, megabytes: usize) -> Result<f64> {
    let stream = TcpStream::connect(addr).await?;
    stream.set_nodelay(true)?;
    let (mut read, mut write) = stream.into_split();
    let start = Instant::now();
    let sender = async {
        let data = vec![0x5au8; 65536];
        for _ in 0..megabytes * 16 {
            write.write_all(&data).await?;
        }
        write.shutdown().await
    };
    let receiver = async {
        let mut data = vec![0; 65536];
        let mut total = 0;
        loop {
            let n = read.read(&mut data).await?;
            if n == 0 {
                break;
            }
            anyhow::ensure!(data[..n].iter().all(|b| *b == 0x5a), "byte corruption");
            total += n;
        }
        anyhow::ensure!(total == megabytes * 1048576, "truncated stream");
        Ok::<_, anyhow::Error>(())
    };
    tokio::try_join!(
        async {
            sender.await?;
            Ok::<_, anyhow::Error>(())
        },
        receiver
    )?;
    Ok((megabytes as f64 * 2.0) / start.elapsed().as_secs_f64())
}
async fn latency(addr: SocketAddr) -> Result<(f64, f64)> {
    let mut stream = TcpStream::connect(addr).await?;
    stream.set_nodelay(true)?;
    let mut times = vec![];
    let mut reply = [0u8; 32];
    for _ in 0..500 {
        let t = Instant::now();
        stream.write_all(&[42u8; 32]).await?;
        stream.read_exact(&mut reply).await?;
        anyhow::ensure!(reply == [42; 32], "latency reply corruption");
        times.push(t.elapsed().as_secs_f64() * 1e6);
    }
    stream.shutdown().await?;
    times.sort_by(f64::total_cmp);
    Ok((times[250], times[475]))
}
#[tokio::main]
async fn main() -> Result<()> {
    let megabytes = std::env::args()
        .nth(1)
        .and_then(|v| v.parse().ok())
        .unwrap_or(64);
    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let direct = listener.local_addr()?;
    let upstream = tokio::spawn(async move {
        loop {
            let (stream, _) = listener.accept().await.unwrap();
            stream.set_nodelay(true).unwrap();
            tokio::spawn(async move {
                let (mut read, mut write) = stream.into_split();
                let _ = tokio::io::copy(&mut read, &mut write).await;
                let _ = write.shutdown().await;
            });
        }
    });
    let reserve = TcpListener::bind("127.0.0.1:0").await?;
    let proxy = reserve.local_addr()?;
    drop(reserve);
    let engine = Engine::new(NetworkConfig::default());
    let service = ServiceConfig {
        id: "bench".into(),
        name: "Benchmark".into(),
        listen: proxy,
        upstream_port: direct.port(),
        primary: Some("local".into()),
        replicas: vec![],
        enabled: true,
        auto_switch: false,
        replica_policy: ReplicaPolicy::Isolate,
    };
    engine
        .upsert(
            &service,
            &[Team {
                id: "local".into(),
                name: "Loopback".into(),
                address: direct.ip(),
            }],
        )
        .await?;
    transfer(direct, 4).await?;
    transfer(proxy, 4).await?;
    let mut direct_runs = vec![];
    let mut proxy_runs = vec![];
    for round in 0..3 {
        if round % 2 == 0 {
            direct_runs.push(transfer(direct, megabytes).await?);
            proxy_runs.push(transfer(proxy, megabytes).await?);
        } else {
            proxy_runs.push(transfer(proxy, megabytes).await?);
            direct_runs.push(transfer(direct, megabytes).await?);
        }
    }
    let direct_latency = latency(direct).await?;
    let proxy_latency = latency(proxy).await?;
    engine.shutdown().await;
    upstream.abort();
    let _ = upstream.await;
    println!(
        "{}",
        serde_json::to_string_pretty(
            &json!({"platform":std::env::consts::OS,"arch":std::env::consts::ARCH,"payload_mib_per_run":megabytes,"throughput_unit":"MiB/s aggregate upload+echo download","direct":direct_runs,"proxy":proxy_runs,"direct_latency_us":{"p50":direct_latency.0,"p95":direct_latency.1},"proxy_latency_us":{"p50":proxy_latency.0,"p95":proxy_latency.1},"byte_integrity":"verified every byte", "stats":engine.snapshots()})
        )?
    );
    Ok(())
}