//! WAN pipelining latency simulation benchmark — Track 8.
//!
//! Measures throughput of the pipelined push protocol (sliding window W) under
//! simulated WAN latency.  RTT delay is injected into the server's Ack sending
//! path via `tokio::time::sleep`, keeping the benchmark fully self-contained
//! (no OS-level traffic shaping required).
//!
//! # What it demonstrates
//!
//! The stop-and-wait baseline (W=1) sends one packet, waits one RTT for an
//! Ack, then sends the next.  For 100 packets at 50 ms RTT, that is 5 seconds.
//! A W=16 pipeline keeps 16 packets in-flight simultaneously, reducing total
//! time to roughly ceil(100/16) × RTT ≈ 7 × 50 ms = 350 ms (15× faster).
//!
//! # Groups
//!
//! `wan_pipeline/rtt_{R}ms_{label}` — 100-revision push, RTT=R ms, window=W
//!
//! Run:
//!   cargo bench -p clawsync-agent -- wan_pipeline_bench

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use clawhdf5_onion::format::NO_PARENT;
use clawhdf5_onion::writer::OnionFile;
use clawsync_onion::differ::diff_revisions;
use clawsync_onion::manifest::ClawSyncManifest;
use clawsync_transport::protocol::SyncMessage;
use clawsync_transport::tcp::{TcpConnection, TcpServer};
use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main};
use tempfile::TempDir;
use tokio::sync::{Semaphore, mpsc};

// ─────────────────────────────────────────────────────────────────────────────

const H5_MAGIC: &[u8] = b"\x89HDF\r\n\x1a\n";
/// Number of revisions used in every bench iteration.
const N: usize = 100;

fn any_addr() -> SocketAddr {
    SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)
}

/// Build an OnionFile with `n` committed revisions in `dir`.
fn make_onion_n(dir: &TempDir, n: usize) -> (PathBuf, OnionFile) {
    let h5 = dir.path().join("src.h5");
    std::fs::write(&h5, H5_MAGIC).unwrap();
    let mut onion = OnionFile::create(&h5, 4096).unwrap();
    for i in 0..n {
        let mut s = onion.begin_session(None).unwrap();
        s.record_page(0, &vec![i as u8; 4096]);
        onion.commit_session(s, Some(&format!("rev {i}"))).unwrap();
    }
    onion.flush().unwrap();
    (h5, onion)
}

// ─────────────────────────────────────────────────────────────────────────────
// Latency-injecting server
//
// Receives exactly N LayerPacket messages, sleeping `rtt_ms` before each Ack
// to simulate the round-trip time of a WAN link.  Uses the ManifestRequest
// protocol (no IBLT) for simplicity.
// ─────────────────────────────────────────────────────────────────────────────

async fn serve_with_latency(server: TcpServer, rtt_ms: f64) {
    let (mut conn, _) = server.accept().await.unwrap();

    // Client sends ManifestRequest; server declares it has 0 revisions so
    // the client will push all N of its own.
    let client_rev_count = match conn.recv().await.unwrap() {
        SyncMessage::ManifestRequest { revision_count, .. } => revision_count,
        other => panic!("server expected ManifestRequest, got {other:?}"),
    };

    conn.send(&SyncMessage::ManifestResponse {
        manifest: ClawSyncManifest {
            agent_id: "bench-server".to_string(),
            file_blake3: [0u8; 32],
            revision_count: 0,
            head_revision: 0,
            head_blake3: [0u8; 32],
            revisions: vec![],
            last_write: 0.0,
        },
    })
    .await
    .unwrap();

    let delay = Duration::from_secs_f64(rtt_ms / 1000.0);
    let mut received = 0u64;
    loop {
        match conn.recv().await.unwrap() {
            SyncMessage::LayerPacket { packet } => {
                let rev = packet.revision;
                // Inject simulated WAN RTT before sending Ack.
                tokio::time::sleep(delay).await;
                conn.send(&SyncMessage::Ack { revision: rev })
                    .await
                    .unwrap();
                received += 1;
            }
            SyncMessage::SyncComplete { .. } => break,
            other => panic!("server: unexpected {other:?}"),
        }
    }
    assert_eq!(received, client_rev_count);
}

// ─────────────────────────────────────────────────────────────────────────────
// Pipelined push client (configurable window)
//
// Uses the ManifestRequest protocol and the W=`window` sliding-window pipeline
// identical to the real cmd_push implementation.
// ─────────────────────────────────────────────────────────────────────────────

async fn pipelined_push(addr: SocketAddr, onion: &OnionFile, window: usize) {
    let mut conn = TcpConnection::connect(addr).await.unwrap();

    let rev_count = onion.revision_count();
    conn.send(&SyncMessage::ManifestRequest {
        agent_id: "bench-client".to_string(),
        head_revision: rev_count.saturating_sub(1),
        revision_count: rev_count,
    })
    .await
    .unwrap();

    match conn.recv().await.unwrap() {
        SyncMessage::ManifestResponse { .. } => {} // server has 0 revisions — we push all
        other => panic!("client expected ManifestResponse, got {other:?}"),
    }

    let packets = diff_revisions(onion, NO_PARENT).unwrap();
    let total = packets.len();

    let (mut read_half, mut write_half) = conn.into_split();
    let semaphore = Arc::new(Semaphore::new(window));
    let (meta_tx, mut meta_rx) = mpsc::unbounded_channel::<u64>(); // bytes per packet
    let sem_writer = semaphore.clone();

    let writer = tokio::spawn(async move {
        for packet in packets {
            let bytes = packet.page_data_size() as u64;
            sem_writer.acquire().await.unwrap().forget();
            meta_tx.send(bytes).unwrap();
            write_half
                .send(&SyncMessage::LayerPacket { packet })
                .await
                .unwrap();
        }
        write_half // return so we can send SyncComplete
    });

    let mut sent = 0u64;
    let mut bytes_sent = 0u64;
    while sent < total as u64 {
        match read_half.recv().await.unwrap() {
            SyncMessage::Ack { .. } => {
                semaphore.add_permits(1);
                bytes_sent += meta_rx.recv().await.unwrap();
                sent += 1;
            }
            other => panic!("client: unexpected {other:?}"),
        }
    }

    let mut write_half = writer.await.unwrap();
    write_half
        .send(&SyncMessage::SyncComplete {
            revisions_transferred: sent,
            bytes_transferred: bytes_sent,
        })
        .await
        .unwrap();
    write_half.shutdown().await.unwrap();
}

// ─────────────────────────────────────────────────────────────────────────────
// Criterion benchmark
// ─────────────────────────────────────────────────────────────────────────────

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

    let mut group = c.benchmark_group("wan_pipeline");
    // Each iteration sleeps (rtt_ms × N/window) ms, so use fewer samples.
    group.sample_size(10);
    group.throughput(Throughput::Elements(N as u64));

    // (rtt_ms, window, label)
    // At 1 ms RTT: W=1 ≈ 100 ms/iter, W=16 ≈ 7 ms/iter
    // At 5 ms RTT: W=1 ≈ 500 ms/iter, W=16 ≈ 32 ms/iter
    let params: &[(f64, usize, &str)] = &[
        (1.0, 1, "rtt_1ms_W1"),
        (1.0, 4, "rtt_1ms_W4"),
        (1.0, 16, "rtt_1ms_W16"),
        (5.0, 1, "rtt_5ms_W1"),
        (5.0, 4, "rtt_5ms_W4"),
        (5.0, 16, "rtt_5ms_W16"),
    ];

    for &(rtt_ms, window, label) in params {
        group.bench_function(label, |b| {
            b.iter_batched(
                || {
                    // SETUP (not timed): build source onion and bind server
                    let dir = TempDir::new().unwrap();
                    let (path, onion) = make_onion_n(&dir, N);
                    let (server, addr) = rt.block_on(async {
                        let s = TcpServer::bind(any_addr()).await.unwrap();
                        let a = s.local_addr;
                        (s, a)
                    });
                    (dir, path, onion, server, addr)
                },
                |(dir, path, onion, server, addr)| {
                    // TIMED: full push with simulated WAN latency
                    let _ = (dir, path); // keep TempDir alive
                    rt.block_on(async {
                        let server_task = tokio::spawn(serve_with_latency(server, rtt_ms));
                        pipelined_push(addr, &onion, window).await;
                        server_task.await.unwrap();
                    });
                },
                BatchSize::SmallInput,
            );
        });
    }

    group.finish();
}

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