//! A load test of the relay's fan-out: one publisher and N subscribers of one track, over native
//! QUIC on loopback, the relay a process of its own, so what it spends is its own.
//!
//! ```text
//! fanout_load --subscribers N [--payload B] [--rate R] [--seconds S] [--relay-threads T] [--relay-cpus 0-7]
//!             [--keyframe-every K --keyframe-bytes B] [--per-endpoint E] [--relay-sockets N]
//!             [--mode subscribe|subscribe-tracks]
//! ```
//!
//! It prints one JSON object: the relay's resident memory before the subscribers, with them idle
//! and at its peak; its CPU over the run; how many objects arrived of how many sent; and how late.
//! A subscriber's lateness is from the publisher's clock, which is the subscriber's: one host. Run
//! the driver on other CPUs than the relay's (`taskset -c 8-31 fanout_load ...`), so that the clients
//! never take the relay's time. Every wait is bounded, and the relay exits when the driver's end of
//! its standard input closes, so a killed run leaves nothing behind. Linux, for `/proc`.
//!
//! A **bystander** runs beside the audience: a track of its own with one subscriber, at the same
//! rate and with small objects. It shares nothing with the audience but the relay, so its lateness
//! is what the audience's fan-out costs everyone else on the relay.
//!
//! `--mode` is how a subscriber asks for the track: `subscribe` (the default), a SUBSCRIBE of the
//! track by name; `subscribe-tracks`, a SUBSCRIBE_TRACKS of its namespace, which the relay answers
//! by starting the subscription itself with a PUBLISH that the subscriber accepts (-22 §3.6, §7.6),
//! so the subscription the objects go out on is the relay's. The bystander's track is then under a
//! namespace of its own, so that no subscriber is offered it. The figures are the same in both, and
//! `cpu_us_per_delivery` is each path's CPU a delivered object.

use joule_moq::draft22::data::{SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::kvp::Properties;
use joule_moq::draft22::{ControlMessage, NamespaceRequest, Parameters, RequestOk, UniEvent};
use joule_moq::{TrackName, TrackNamespace};
use joule_moq_relay::client::{Client, ClientOptions};
use joule_moq_relay::transport::quic;
use joule_moq_relay::{RelayConfig, RelayServer};
use std::io::{BufRead, Read, Write};
use std::process::{Child, Command, Stdio};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::time::timeout;

/// Bytes each client socket asks of the kernel.
const CLIENT_SOCKET_BUFFER: usize = 4 << 20;
/// Latency buckets of 50 µs, up to 10 s; the last holds everything later.
const BUCKET_US: u64 = 50;
const BUCKETS: usize = 200_000;
const STEP: Duration = Duration::from_secs(30);

fn arg<T: std::str::FromStr>(args: &[String], flag: &str, default: T) -> T {
    args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).and_then(|v| v.parse().ok()).unwrap_or(default)
}

fn relay_main(args: &[String]) {
    let threads: usize = arg(args, "--relay-threads", 8);
    let sockets: usize = arg(args, "--relay-sockets", 0);
    let rt = tokio::runtime::Builder::new_multi_thread().worker_threads(threads).enable_all().build().expect("a runtime");
    rt.block_on(async {
        let mut c = RelayConfig::loopback();
        c.websocket_listen = None;
        c.limits.max_sessions = 1 << 20;
        c.limits.quic_sockets = sockets;
        let server = RelayServer::start(c).await.expect("the relay starts");
        let cert: String = server.identity.chain[0].as_ref().iter().map(|b| format!("{b:02x}")).collect();
        println!("{} {cert}", server.quic_addr.expect("a UDP endpoint"));
        std::io::stdout().flush().expect("stdout");
        // Until the driver goes.
        tokio::task::spawn_blocking(|| {
            let mut sink = Vec::new();
            let _ = std::io::stdin().read_to_end(&mut sink);
        })
        .await
        .expect("the wait");
        std::process::exit(0);
    });
}

fn proc_kb(pid: u32, field: &str) -> u64 {
    std::fs::read_to_string(format!("/proc/{pid}/status")).ok().and_then(|s| s.lines().find(|l| l.starts_with(field)).and_then(|l| l.split_whitespace().nth(1)?.parse().ok())).unwrap_or(0)
}

/// The process's user and system CPU, in seconds.
fn proc_cpu(pid: u32) -> f64 {
    let Ok(s) = std::fs::read_to_string(format!("/proc/{pid}/stat")) else { return 0.0 };
    let after = s.rsplit(')').next().unwrap_or("");
    let f: Vec<&str> = after.split_whitespace().collect();
    let t = |i: usize| f.get(i).and_then(|v| v.parse::<u64>().ok()).unwrap_or(0);
    (t(11) + t(12)) as f64 / 100.0
}

/// The packets the kernel dropped on UDP sockets on loopback: the relay's, at `port`, and the rest.
fn udp_drops(port: u16) -> (u64, u64) {
    let Ok(s) = std::fs::read_to_string("/proc/net/udp") else { return (0, 0) };
    let (mut relay, mut others) = (0, 0);
    for l in s.lines().skip(1) {
        let f: Vec<&str> = l.split_whitespace().collect();
        let (Some(local), Some(drops)) = (f.get(1), f.last().and_then(|d| d.parse::<u64>().ok())) else { continue };
        let Some((host, p)) = local.split_once(':') else { continue };
        if host != "0100007F" {
            continue;
        }
        if u16::from_str_radix(p, 16) == Ok(port) {
            relay += drops;
        } else {
            others += drops;
        }
    }
    (relay, others)
}

fn now_us() -> u64 {
    SystemTime::now().duration_since(UNIX_EPOCH).expect("after 1970").as_micros() as u64
}

fn spawn_relay(args: &[String]) -> (Child, std::net::SocketAddr, Vec<u8>) {
    let exe = std::env::current_exe().expect("this binary");
    let threads: usize = arg(args, "--relay-threads", 8);
    let cpus: String = arg(args, "--relay-cpus", "0-7".to_string());
    let mut child = Command::new("taskset")
        .args(["-c", &cpus])
        .arg(&exe)
        .args(["relay", "--relay-threads", &threads.to_string(), "--relay-sockets", &arg::<usize>(args, "--relay-sockets", 0).to_string()])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()
        .expect("taskset runs");
    let mut line = String::new();
    std::io::BufReader::new(child.stdout.take().expect("its stdout")).read_line(&mut line).expect("its address");
    let (addr, cert) = line.trim().split_once(' ').expect("an address and a certificate");
    let cert = (0..cert.len() / 2).map(|i| u8::from_str_radix(&cert[2 * i..2 * i + 2], 16).expect("hex")).collect();
    (child, addr.parse().expect("an address"), cert)
}

struct Seen {
    objects: AtomicU64,
    buckets: Vec<AtomicU64>,
}

impl Seen {
    fn note(&self, late_us: u64) {
        self.objects.fetch_add(1, Ordering::Relaxed);
        self.buckets[((late_us / BUCKET_US) as usize).min(BUCKETS - 1)].fetch_add(1, Ordering::Relaxed);
    }

    /// The latency, in milliseconds, at each fraction.
    fn at(&self, fractions: &[f64]) -> Vec<f64> {
        let counts: Vec<u64> = self.buckets.iter().map(|b| b.load(Ordering::Relaxed)).collect();
        let total: u64 = counts.iter().sum();
        fractions
            .iter()
            .map(|f| {
                let want = ((total as f64) * f).ceil() as u64;
                let mut sum = 0;
                for (i, c) in counts.iter().enumerate() {
                    sum += c;
                    if sum >= want.max(1) {
                        return (i as u64 * BUCKET_US) as f64 / 1000.0;
                    }
                }
                0.0
            })
            .collect()
    }
}

fn driver(args: &[String]) {
    let subscribers: usize = arg(args, "--subscribers", 100);
    let payload: usize = arg(args, "--payload", 8192);
    let rate: u64 = arg(args, "--rate", 30);
    let seconds: u64 = arg(args, "--seconds", 20);
    let keyframe_every: u64 = arg(args, "--keyframe-every", 0);
    let keyframe_bytes: usize = arg(args, "--keyframe-bytes", 0);
    let per_endpoint: usize = arg(args, "--per-endpoint", 8);
    let mode: String = arg(args, "--mode", "subscribe".to_string());
    let by_tracks = match mode.as_str() {
        "subscribe" => false,
        "subscribe-tracks" => true,
        other => panic!("--mode is subscribe or subscribe-tracks, not {other}"),
    };
    let (child, addr, cert) = spawn_relay(args);
    let pid = child.id();
    let rt = tokio::runtime::Builder::new_multi_thread().enable_all().build().expect("a runtime");
    rt.block_on(async move {
        tokio::time::sleep(Duration::from_millis(300)).await;
        let cert = rustls::pki_types::CertificateDer::from(cert);
        let alpn = joule_moq::draft22::ALPN.to_vec();
        let bind: std::net::SocketAddr = "127.0.0.1:0".parse().expect("an address");
        let (space, track) = (TrackNamespace::from_strs(["load"]), TrackName::from_str("video"));
        let base_kb = proc_kb(pid, "VmRSS");

        let pub_ep = quic::client_endpoint(bind, vec![cert.clone()], vec![alpn.clone()], None).expect("an endpoint");
        let conn = timeout(STEP, quic::connect(&pub_ep, addr, "localhost")).await.expect("in time").expect("a session");
        let publisher = Arc::new(Client::connect(Arc::new(conn), ClientOptions::default()).await.expect("SETUP"));
        let _published = publisher.publish(space.clone(), track.clone(), 1, Parameters::new(), Default::default()).await.expect("PUBLISH");

        let seen = Arc::new(Seen { objects: AtomicU64::new(0), buckets: (0..BUCKETS).map(|_| AtomicU64::new(0)).collect() });
        let endpoints: Vec<quinn::Endpoint> = (0..subscribers.div_ceil(per_endpoint))
            .map(|_| {
                let ep = quic::client_endpoint(bind, vec![cert.clone()], vec![alpn.clone()], None).expect("an endpoint");
                // A viewer has a socket of its own, and these share one: a small buffer here would
                // drop what no real viewer's would, and measure this program, not the relay.
                ep.rebind(quic::bind_udp(bind, CLIENT_SOCKET_BUFFER).expect("a socket")).expect("rebinding");
                ep
            })
            .collect();
        let connecting = Instant::now();
        let gate = Arc::new(tokio::sync::Semaphore::new(128));
        let mut joins = Vec::with_capacity(subscribers);
        for i in 0..subscribers {
            let ep = endpoints[i / per_endpoint].clone();
            let (gate, space, track, seen) = (gate.clone(), space.clone(), track.clone(), seen.clone());
            joins.push(tokio::spawn(async move {
                let permit = gate.acquire_owned().await.expect("the gate");
                let conn = timeout(STEP, quic::connect(&ep, addr, "localhost")).await.expect("in time").expect("a session");
                let c = Client::connect(Arc::new(conn), ClientOptions::default()).await.expect("SETUP");
                // What the subscriber holds open for the run: its SUBSCRIBE's answer and request
                // stream, or its SUBSCRIBE_TRACKS's request stream and the relay's PUBLISH's.
                let held: Box<dyn std::any::Any + Send> = if by_tracks {
                    // SUBSCRIBE_TRACKS of the namespace; the relay's PUBLISH of the one track under it,
                    // accepted (PUBLISH_OK): its subscription is the one the objects arrive on.
                    let id = c.next_request_id();
                    let (ttx, mut trx) = c.request(&ControlMessage::SubscribeTracks(NamespaceRequest { request_id: id, prefix: space, params: Parameters::new() })).await.expect("SUBSCRIBE_TRACKS goes");
                    let answer = timeout(STEP, Client::response(&mut trx)).await.expect("an answer in time").expect("an answer");
                    assert!(matches!(answer, ControlMessage::RequestOk(_)), "SUBSCRIBE_TRACKS answered {}", answer.name());
                    let (ptx, prx, first) = c.next_request(STEP).await.expect("the relay's PUBLISH in time");
                    assert!(matches!(&first, ControlMessage::Publish(p) if p.name == track), "the relay opened {} where a PUBLISH of the track was expected", first.name());
                    ptx.send(&ControlMessage::RequestOk(RequestOk { params: Parameters::new(), track_properties: Properties::new() })).await.expect("PUBLISH_OK goes");
                    Box::new((ttx, trx, ptx, prx))
                } else {
                    Box::new(c.subscribe(space, track, Parameters::new()).await.expect("SUBSCRIBE"))
                };
                drop(permit);
                tokio::spawn(async move {
                    let _held = held;
                    while let Some(mut r) = c.next_data(Duration::from_secs(3600)).await {
                        while let Ok(Some(e)) = r.next().await {
                            if let UniEvent::SubgroupObject(o) = e {
                                let sent = u64::from_le_bytes(o.payload[..8].try_into().expect("a timestamp"));
                                seen.note(now_us().saturating_sub(sent));
                            }
                        }
                    }
                });
            }));
        }
        for j in joins {
            j.await.expect("a subscriber");
        }
        let connect_s = connecting.elapsed().as_secs_f64();

        // The bystander: a publisher and a subscriber of their own track.
        let aside = Arc::new(Seen { objects: AtomicU64::new(0), buckets: (0..BUCKETS).map(|_| AtomicU64::new(0)).collect() });
        let by_ep = quic::client_endpoint(bind, vec![cert.clone()], vec![alpn.clone()], None).expect("an endpoint");
        let by_name = TrackName::from_str("bystander");
        // Under SUBSCRIBE_TRACKS the audience asks for every track of `space`: the bystander's is
        // under another, so it is not offered to them.
        let by_space = if by_tracks { TrackNamespace::from_strs(["aside"]) } else { space.clone() };
        let conn = timeout(STEP, quic::connect(&by_ep, addr, "localhost")).await.expect("in time").expect("a session");
        let by_pub = Arc::new(Client::connect(Arc::new(conn), ClientOptions::default()).await.expect("SETUP"));
        let _by_published = by_pub.publish(by_space.clone(), by_name.clone(), 2, Parameters::new(), Default::default()).await.expect("PUBLISH");
        let conn = timeout(STEP, quic::connect(&by_ep, addr, "localhost")).await.expect("in time").expect("a session");
        let by_sub = Client::connect(Arc::new(conn), ClientOptions::default()).await.expect("SETUP");
        let by_held = by_sub.subscribe(by_space, by_name, Parameters::new()).await.expect("SUBSCRIBE");
        let by_seen = aside.clone();
        tokio::spawn(async move {
            let _held = by_held;
            while let Some(mut r) = by_sub.next_data(Duration::from_secs(3600)).await {
                while let Ok(Some(e)) = r.next().await {
                    if let UniEvent::SubgroupObject(o) = e {
                        by_seen.note(now_us().saturating_sub(u64::from_le_bytes(o.payload[..8].try_into().expect("a timestamp"))));
                    }
                }
            }
        });
        tokio::time::sleep(Duration::from_secs(3)).await;
        let idle_kb = proc_kb(pid, "VmRSS");
        let idle_cpu_start = proc_cpu(pid);
        tokio::time::sleep(Duration::from_secs(3)).await;
        let idle_cores = (proc_cpu(pid) - idle_cpu_start) / 3.0;

        // The run: one object to a group, at the rate, each stamped with when it left.
        let drops_start = udp_drops(addr.port());
        let cpu_start = proc_cpu(pid);
        let driver_start = proc_cpu(std::process::id());
        let started = Instant::now();
        let by_tick = Duration::from_nanos(1_000_000_000 / rate);
        let by_run = Duration::from_secs(seconds);
        let by_sender = tokio::spawn(async move {
            let mut tick = tokio::time::interval(by_tick);
            let (t0, mut g) = (Instant::now(), 0u64);
            while t0.elapsed() < by_run {
                tick.tick().await;
                let mut b = vec![0u8; 1024];
                b[..8].copy_from_slice(&now_us().to_le_bytes());
                let h = SubgroupHeader { track_alias: 2, group_id: g, subgroup_id: SubgroupId::Zero, publisher_priority: Some(128), properties: false, end_of_group: true, first_object: true };
                let _ = timeout(STEP, by_pub.send_subgroup(h, &[SubgroupObject::normal(0, b)])).await;
                g += 1;
            }
        });
        let mut tick = tokio::time::interval(Duration::from_nanos(1_000_000_000 / rate));
        let body = vec![0x5au8; payload.max(8)];
        let mut sent = 0u64;
        let mut lagging = 0u64;
        let mut peak_kb = 0;
        while started.elapsed() < Duration::from_secs(seconds) {
            tick.tick().await;
            let mut b = if keyframe_every != 0 && sent.is_multiple_of(keyframe_every) { vec![0x5au8; keyframe_bytes.max(8)] } else { body.clone() };
            b[..8].copy_from_slice(&now_us().to_le_bytes());
            let h = SubgroupHeader { track_alias: 1, group_id: sent, subgroup_id: SubgroupId::Zero, publisher_priority: Some(128), properties: false, end_of_group: true, first_object: true };
            let t0 = Instant::now();
            timeout(STEP, publisher.send_subgroup(h, &[SubgroupObject::normal(0, b)])).await.expect("in time").expect("a send");
            if t0.elapsed() > Duration::from_nanos(1_000_000_000 / rate) {
                lagging += 1;
            }
            sent += 1;
            if sent.is_multiple_of(rate) {
                peak_kb = peak_kb.max(proc_kb(pid, "VmRSS"));
            }
        }
        let run_s = started.elapsed().as_secs_f64();
        let run_cpu = proc_cpu(pid) - cpu_start;
        let driver_cores = (proc_cpu(std::process::id()) - driver_start) / run_s;
        let drops_end = udp_drops(addr.port());
        let _ = by_sender.await;
        tokio::time::sleep(Duration::from_secs(3)).await;
        let got = seen.objects.load(Ordering::Relaxed);
        let want = sent * subscribers as u64;
        let l = seen.at(&[0.5, 0.9, 0.99, 0.999]);
        let by = aside.at(&[0.5, 0.99, 0.999]);
        println!(
            "{{\"mode\":\"{mode}\",\"subscribers\":{subscribers},\"payload\":{payload},\"rate\":{rate},\"seconds\":{seconds},\"connect_s\":{connect_s:.1},\"rss_base_mb\":{:.1},\"rss_idle_mb\":{:.1},\"kb_per_idle_subscriber\":{:.1},\"idle_cores\":{idle_cores:.3},\"rss_peak_mb\":{:.1},\"relay_cores\":{:.2},\"cpu_us_per_delivery\":{:.1},\"sent\":{sent},\"send_lagged\":{lagging},\"delivered\":{got},\"delivered_ratio\":{:.4},\"gbps\":{:.2},\"p50_ms\":{:.2},\"p90_ms\":{:.2},\"p99_ms\":{:.2},\"p999_ms\":{:.2},\"driver_cores\":{driver_cores:.2},\"udp_drops_relay\":{},\"udp_drops_clients\":{},\"bystander_p50_ms\":{:.2},\"bystander_p99_ms\":{:.2},\"bystander_p999_ms\":{:.2}}}",
            base_kb as f64 / 1024.0,
            idle_kb as f64 / 1024.0,
            idle_kb.saturating_sub(base_kb) as f64 / subscribers.max(1) as f64,
            peak_kb.max(proc_kb(pid, "VmHWM")) as f64 / 1024.0,
            run_cpu / run_s,
            run_cpu * 1e6 / got.max(1) as f64,
            got as f64 / want.max(1) as f64,
            got as f64 * payload as f64 * 8.0 / run_s / 1e9,
            l[0],
            l[1],
            l[2],
            l[3],
            drops_end.0.saturating_sub(drops_start.0),
            drops_end.1.saturating_sub(drops_start.1),
            by[0],
            by[1],
            by[2],
        );
        std::process::exit(0);
    });
}

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.first().map(String::as_str) == Some("relay") {
        relay_main(&args);
    } else {
        driver(&args);
    }
}