zc2 0.0.28

//! `zc bench mesh` — mesh warmup / calibration CLI (issue #47).
//!
//! Probes each worker URI and emits a per-worker report (latency mean + p95,
//! an optional bandwidth estimate) plus a recommended backpressure threshold,
//! mirroring the Python `zakuro.AdaptiveCompute.warmup()` primitive so operators
//! can calibrate a mesh without first spinning up training code. Output is JSON
//! on stdout (progress goes to stderr) so it pipes cleanly into `jq` or a config
//! store.
//!
//! ## Parity with `zakuro/adaptive.py::AdaptiveCompute.warmup`
//! All latency figures are in **seconds** to match the Python report fields.
//!   * `latency_p95`            = `sorted[min(len-1, int(0.95 * len))]`
//!   * `bandwidth_bps`          = `bytes / max(best - mean, 1e-3)`  (bytes/sec; None on probe failure)
//!   * `recommended_backpressure` = `max(healthy p95) * 1.5`         (seconds; None when no healthy worker)
//!
//! ## Transport (v1)
//! Latency is measured with the worker's unauthenticated `GET /health`, and the
//! bandwidth probe POSTs a sized body to `/execute` — the same HTTP worker
//! contract `zc bench` already speaks. The full zakuro-wire postcard `Envelope`
//! probe (HMAC-signed, over QUIC `OP_EXECUTE`) is a deliberate follow-up: it
//! needs the per-tenant key and a live Docker mesh to validate the documented
//! 2× parity, neither of which the CLI can exercise on its own.

use std::time::{Duration, Instant};

use serde::Serialize;

/// Successful round-trips to attempt per worker (warmup default `rounds=3`).
const DEFAULT_ROUNDS: usize = 3;
/// Per-probe timeout (warmup default `timeout=10.0`).
const DEFAULT_TIMEOUT_SECS: u64 = 10;
/// Bandwidth probe payload size (warmup default `bandwidth_probe_bytes=1 MiB`).
const DEFAULT_BANDWIDTH_BYTES: usize = 1024 * 1024;
/// Bandwidth samples per worker (warmup runs `for _ in range(2)`).
const BANDWIDTH_SAMPLES: usize = 2;

/// Mesh probe configuration parsed from CLI args.
#[derive(Debug, Clone)]
pub struct MeshProbeConfig {
    /// Worker URIs to probe (`zc://name`, `http://host:port`, or bare `host:port`).
    pub workers: Vec<String>,
    pub rounds: usize,
    pub timeout_secs: u64,
    pub bandwidth_bytes: usize,
}

/// Per-worker probe report. Field names mirror the Python warmup report so the
/// two outputs are directly comparable.
#[derive(Debug, Clone, Serialize, PartialEq)]
pub struct WorkerReport {
    pub uri: String,
    pub ok: bool,
    pub rounds_succeeded: usize,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub latency_mean: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub latency_p95: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bandwidth_bps: Option<f64>,
    pub observed: Vec<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reason: Option<String>,
}

impl WorkerReport {
    fn failed(uri: &str, reason: String) -> Self {
        WorkerReport {
            uri: uri.to_string(),
            ok: false,
            rounds_succeeded: 0,
            latency_mean: None,
            latency_p95: None,
            bandwidth_bps: None,
            observed: Vec::new(),
            reason: Some(reason),
        }
    }
}

/// Top-level mesh report (matches the Python warmup return schema).
#[derive(Debug, Clone, Serialize, PartialEq)]
pub struct MeshReport {
    pub rounds: usize,
    pub workers: Vec<WorkerReport>,
    pub recommended_backpressure: Option<f64>,
}

// ── pure math (byte-for-byte parity with zakuro/adaptive.py) ────────────────

/// Arithmetic mean of the samples (0.0 for an empty slice).
fn mean(samples: &[f64]) -> f64 {
    if samples.is_empty() {
        return 0.0;
    }
    samples.iter().sum::<f64>() / samples.len() as f64
}

/// p95 using the exact index the Python side uses:
/// `sorted[min(len - 1, int(0.95 * len))]`. Rust and Python both use IEEE-754
/// doubles and both truncate toward zero, so `(0.95 * len) as usize` reproduces
/// `int(0.95 * len)` bit-for-bit.
fn p95(samples: &[f64]) -> f64 {
    if samples.is_empty() {
        return 0.0;
    }
    let mut sorted = samples.to_vec();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let idx = ((0.95 * sorted.len() as f64) as usize).min(sorted.len() - 1);
    sorted[idx]
}

/// Bandwidth estimate in bytes/sec: `bytes / max(best - mean, 1e-3)`. Subtracts
/// the latency floor (mean RTT) so fixed overhead isn't counted against
/// throughput, flooring the transfer time at 1ms to avoid divide-by-zero.
fn bandwidth_bps(bytes: usize, best: f64, mean_latency: f64) -> f64 {
    let transfer = (best - mean_latency).max(1e-3);
    bytes as f64 / transfer
}

/// Recommended backpressure threshold = `max(healthy p95) * 1.5`; `None` when no
/// worker reported a healthy p95.
fn recommended_backpressure(p95s: &[f64]) -> Option<f64> {
    p95s.iter()
        .cloned()
        .fold(None, |acc: Option<f64>, v| {
            Some(acc.map_or(v, |m| m.max(v)))
        })
        .map(|m| m * 1.5)
}

// ── transport ───────────────────────────────────────────────────────────────

/// Normalize a worker entry to an `http://host:port` base URL for HTTP probing.
fn worker_base_url(entry: &str) -> Result<String, String> {
    let e = entry.trim();
    if e.is_empty() {
        return Err("empty worker URI".to_string());
    }
    if e.starts_with("zc://") {
        // Reuse the broker URI resolver for named `zc://` hosts.
        return super::uri::resolve(e);
    }
    if e.starts_with("http://") || e.starts_with("https://") {
        return Ok(e.trim_end_matches('/').to_string());
    }
    // Bare `host:port`.
    Ok(format!("http://{}", e.trim_end_matches('/')))
}

/// Probe one worker: `rounds` latency round-trips against `/health`, then a
/// best-effort sized bandwidth probe against `/execute`.
fn probe_worker(uri: &str, cfg: &MeshProbeConfig) -> WorkerReport {
    let base = match worker_base_url(uri) {
        Ok(b) => b,
        Err(e) => return WorkerReport::failed(uri, format!("invalid URI: {e}")),
    };
    let timeout = Duration::from_secs(cfg.timeout_secs);
    let health_url = format!("{base}/health");

    let mut observed = Vec::with_capacity(cfg.rounds);
    for _ in 0..cfg.rounds {
        let start = Instant::now();
        let ok = crate::vpn::mesh_agent(std::time::Duration::from_secs(30))
            .get(&health_url)
            .config()
            .timeout_global(Some(timeout))
            .build()
            .call()
            .is_ok();
        let elapsed = start.elapsed().as_secs_f64();
        if ok {
            observed.push(elapsed);
        }
    }

    if observed.is_empty() {
        return WorkerReport::failed(
            uri,
            format!("no successful /health probe in {} round(s)", cfg.rounds),
        );
    }

    let mean_latency = mean(&observed);
    let p95_latency = p95(&observed);
    let bandwidth = probe_bandwidth(&base, cfg, mean_latency);

    WorkerReport {
        uri: uri.to_string(),
        ok: true,
        rounds_succeeded: observed.len(),
        latency_mean: Some(mean_latency),
        latency_p95: Some(p95_latency),
        bandwidth_bps: bandwidth,
        observed,
        reason: None,
    }
}

/// Best-effort bandwidth probe: POST a sized body to `/execute` `BANDWIDTH_SAMPLES`
/// times and estimate throughput from the fastest sample. Returns `None` when the
/// worker rejects the sized probe — exactly the warmup's `else None` fallback.
fn probe_bandwidth(base: &str, cfg: &MeshProbeConfig, mean_latency: f64) -> Option<f64> {
    if cfg.bandwidth_bytes == 0 {
        return None;
    }
    let url = format!("{base}/execute");
    let payload = vec![0u8; cfg.bandwidth_bytes];
    let timeout = Duration::from_secs(cfg.timeout_secs);

    let mut samples = Vec::with_capacity(BANDWIDTH_SAMPLES);
    for _ in 0..BANDWIDTH_SAMPLES {
        let start = Instant::now();
        let res = crate::vpn::mesh_agent(std::time::Duration::from_secs(310))
            .post(&url)
            .config()
            .timeout_global(Some(timeout))
            .build()
            .header("Content-Type", "application/octet-stream")
            .header("X-Zakuro-User", "bench-mesh")
            .send(&payload);
        match res {
            Ok(_) => samples.push(start.elapsed().as_secs_f64()),
            Err(_) => break,
        }
    }

    if samples.is_empty() {
        return None;
    }
    let best = samples.iter().cloned().fold(f64::INFINITY, f64::min);
    Some(bandwidth_bps(cfg.bandwidth_bytes, best, mean_latency))
}

/// Probe every worker and assemble the mesh report.
pub fn run_mesh(cfg: MeshProbeConfig) -> MeshReport {
    let mut workers = Vec::with_capacity(cfg.workers.len());
    for uri in &cfg.workers {
        eprintln!("  probing {uri} ...");
        workers.push(probe_worker(uri, &cfg));
    }

    let healthy_p95: Vec<f64> = workers
        .iter()
        .filter(|w| w.ok)
        .filter_map(|w| w.latency_p95)
        .collect();
    let recommended = recommended_backpressure(&healthy_p95);

    MeshReport {
        rounds: cfg.rounds,
        workers,
        recommended_backpressure: recommended,
    }
}

/// Entry point for `zc bench mesh ...` (args are everything after `mesh`).
pub fn run_from_args(args: &[String]) {
    let mut workers_raw: Option<String> = None;
    let mut rounds = DEFAULT_ROUNDS;
    let mut timeout_secs = DEFAULT_TIMEOUT_SECS;
    let mut bandwidth_bytes = DEFAULT_BANDWIDTH_BYTES;

    let mut i = 0;
    while i < args.len() {
        match args[i].as_str() {
            "-w" | "--workers" => {
                if i + 1 < args.len() {
                    workers_raw = Some(args[i + 1].clone());
                    i += 1;
                }
            }
            "--rounds" => {
                if i + 1 < args.len() {
                    rounds = args[i + 1].parse().unwrap_or(DEFAULT_ROUNDS);
                    i += 1;
                }
            }
            "--timeout" => {
                if i + 1 < args.len() {
                    timeout_secs = args[i + 1].parse().unwrap_or(DEFAULT_TIMEOUT_SECS);
                    i += 1;
                }
            }
            "--bandwidth-bytes" => {
                if i + 1 < args.len() {
                    bandwidth_bytes = args[i + 1].parse().unwrap_or(DEFAULT_BANDWIDTH_BYTES);
                    i += 1;
                }
            }
            "-h" | "--help" => {
                print_mesh_help();
                return;
            }
            other => {
                eprintln!("Unknown option: {other}");
                print_mesh_help();
                return;
            }
        }
        i += 1;
    }

    let workers: Vec<String> = match workers_raw {
        Some(s) => s
            .split(',')
            .map(|w| w.trim().to_string())
            .filter(|w| !w.is_empty())
            .collect(),
        None => {
            eprintln!("Error: --workers <uri,uri,...> is required");
            print_mesh_help();
            std::process::exit(1);
        }
    };
    if workers.is_empty() {
        eprintln!("Error: no worker URIs given");
        std::process::exit(1);
    }

    let cfg = MeshProbeConfig {
        workers,
        rounds: rounds.max(1),
        timeout_secs,
        bandwidth_bytes,
    };

    let report = run_mesh(cfg);
    match serde_json::to_string_pretty(&report) {
        Ok(json) => println!("{json}"),
        Err(e) => {
            eprintln!("Error serializing report: {e}");
            std::process::exit(1);
        }
    }
}

fn print_mesh_help() {
    println!();
    println!("Usage: zc bench mesh --workers <uri,uri,...> [OPTIONS]");
    println!();
    println!("Probe each worker and print a JSON warmup/calibration report");
    println!("(per-worker latency p95 + bandwidth, recommended backpressure).");
    println!();
    println!("Options:");
    println!("  -w, --workers <list>     Comma-separated worker URIs (required)");
    println!("                           zc://name | http://host:port | host:port");
    println!("      --rounds <N>         Latency round-trips per worker (default: 3)");
    println!("      --timeout <SECS>     Per-probe timeout (default: 10)");
    println!("      --bandwidth-bytes <N> Bandwidth probe size, 0 to skip (default: 1048576)");
    println!("  -h, --help               Show this help");
    println!();
    println!("Example:");
    println!("  zc bench mesh --workers a:3960,b:3960,c:3960 | jq .");
    println!();
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── p95 parity: sorted[min(len-1, int(0.95*len))] ──────────────────────
    #[test]
    fn p95_single_sample() {
        assert_eq!(p95(&[0.42]), 0.42);
    }

    #[test]
    fn p95_three_samples_is_max() {
        // int(0.95*3)=int(2.85)=2, min(2,2)=2 → the max
        assert_eq!(p95(&[0.1, 0.3, 0.2]), 0.3);
    }

    #[test]
    fn p95_ten_samples() {
        // int(0.95*10)=int(9.5)=9 → sorted[9]
        let s: Vec<f64> = (0..10).map(|x| x as f64 / 10.0).collect();
        assert!((p95(&s) - 0.9).abs() < 1e-12);
    }

    #[test]
    fn p95_twenty_samples() {
        // int(0.95*20)=19 → sorted[19] (the max)
        let s: Vec<f64> = (0..20).map(|x| x as f64 / 20.0).collect();
        assert!((p95(&s) - 0.95).abs() < 1e-12);
    }

    #[test]
    fn p95_sorts_unsorted_input() {
        assert_eq!(p95(&[0.3, 0.1, 0.2]), 0.3);
    }

    #[test]
    fn p95_empty_is_zero() {
        assert_eq!(p95(&[]), 0.0);
    }

    // ── mean ───────────────────────────────────────────────────────────────
    #[test]
    fn mean_basic() {
        assert!((mean(&[0.1, 0.2, 0.3]) - 0.2).abs() < 1e-12);
    }

    #[test]
    fn mean_empty_is_zero() {
        assert_eq!(mean(&[]), 0.0);
    }

    // ── bandwidth_bps = bytes / max(best - mean, 1e-3) ─────────────────────
    #[test]
    fn bandwidth_basic() {
        // 1 MiB, best 0.05s, mean 0.04s → transfer ~0.01s
        let bw = bandwidth_bps(1024 * 1024, 0.05, 0.04);
        assert!((bw - (1024.0 * 1024.0 / 0.01)).abs() < 100.0);
    }

    #[test]
    fn bandwidth_floors_transfer_at_1ms() {
        // best <= mean → transfer clamped to 1e-3, never divide-by-zero/negative
        let bw = bandwidth_bps(1000, 0.04, 0.041);
        assert!((bw - (1000.0 / 1e-3)).abs() < 1e-6);
    }

    // ── recommended_backpressure = max(healthy p95) * 1.5 ──────────────────
    #[test]
    fn recommended_is_max_p95_times_1_5() {
        // max(0.1, 0.2, 0.05) * 1.5 = 0.3 (within IEEE-754 tolerance)
        let got = recommended_backpressure(&[0.1, 0.2, 0.05]).unwrap();
        assert!((got - 0.3).abs() < 1e-12, "got {got}");
    }

    #[test]
    fn recommended_none_when_no_healthy() {
        assert_eq!(recommended_backpressure(&[]), None);
    }

    // ── URI normalization ──────────────────────────────────────────────────
    #[test]
    fn worker_base_url_forms() {
        assert_eq!(
            worker_base_url("1.2.3.4:9000").unwrap(),
            "http://1.2.3.4:9000"
        );
        assert_eq!(worker_base_url("http://h:1/").unwrap(), "http://h:1");
        assert!(worker_base_url("  ").is_err());
    }

    // ── JSON schema matches the Python warmup report ───────────────────────
    #[test]
    fn report_serializes_expected_schema() {
        let report = MeshReport {
            rounds: 3,
            workers: vec![WorkerReport {
                uri: "w1".into(),
                ok: true,
                rounds_succeeded: 3,
                latency_mean: Some(0.10),
                latency_p95: Some(0.12),
                bandwidth_bps: None,
                observed: vec![0.10, 0.11, 0.12],
                reason: None,
            }],
            recommended_backpressure: Some(0.18),
        };
        let v: serde_json::Value = serde_json::to_value(&report).unwrap();
        assert_eq!(v["rounds"], 3);
        assert_eq!(v["recommended_backpressure"], 0.18);
        assert_eq!(v["workers"][0]["uri"], "w1");
        assert_eq!(v["workers"][0]["latency_p95"], 0.12);
        // None fields are omitted, not serialized as null.
        assert!(v["workers"][0].get("bandwidth_bps").is_none());
        assert!(v["workers"][0].get("reason").is_none());
    }

    #[test]
    fn failed_report_carries_reason_and_not_ok() {
        let r = WorkerReport::failed("dead:1", "boom".into());
        assert!(!r.ok);
        assert_eq!(r.rounds_succeeded, 0);
        let v: serde_json::Value = serde_json::to_value(&r).unwrap();
        assert_eq!(v["ok"], false);
        assert_eq!(v["reason"], "boom");
        assert!(v.get("latency_p95").is_none());
    }
}