//! What it costs to take in the Internet routing table, from the bytes on
//! the wire to a chosen best route.
//!
//! Run with `cargo bench -p bgp-bench --bench full_table`. Set
//! `BENCH_PEERS` (default 4) and `BENCH_SCALE_DOWN` (default 1; 10 makes the
//! table ten times smaller) to change the size. The table is synthetic (see
//! the crate documentation); `real_table` runs on a real one.
//!
//! This is a stopwatch around whole operations at full size, not a
//! statistical micro-benchmark: the questions are "how long until a full
//! table is in" and "how much memory does it take", and their answers are
//! seconds and gigabytes.

use std::hint::black_box;
use std::net::Ipv4Addr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};

use bgp_bench::memory::{Region, describe};
use bgp_bench::{Prefixes, Shape, announcements, encoded, wire_context};
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::message::{Message, PathAttribute};
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::table::{PeerIndex, PeerInfo, RibTable};
use bytes::BytesMut;

fn env(name: &str, default: usize) -> usize {
    std::env::var(name)
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(default)
}

fn mib(bytes: usize) -> f64 {
    bytes as f64 / (1024.0 * 1024.0)
}

fn per_second(count: usize, elapsed: Duration) -> String {
    let rate = count as f64 / elapsed.as_secs_f64();
    if rate >= 1e6 {
        format!("{:.2} M/s", rate / 1e6)
    } else {
        format!("{:.0} k/s", rate / 1e3)
    }
}

fn row(what: &str, elapsed: Duration, detail: String) {
    println!("{what:<44} {:>9.3} s   {detail}", elapsed.as_secs_f64());
}

fn main() {
    let peers = env("BENCH_PEERS", 4) as u32;
    let shape = Shape::INTERNET.scaled_down(env("BENCH_SCALE_DOWN", 1));
    let threads = std::thread::available_parallelism().map_or(1, |n| n.get());

    println!(
        "== Table: {} IPv4 + {} IPv6 prefixes, from each of {peers} peers",
        shape.ipv4_prefixes, shape.ipv6_prefixes
    );
    println!(
        "(IpPrefix is {} bytes, PathAttribute {})",
        size_of::<IpPrefix>(),
        size_of::<PathAttribute>()
    );
    let prefixes = Prefixes::generate(shape);
    let wire: Vec<Vec<BytesMut>> = (1..=peers)
        .map(|peer| encoded(&announcements(&prefixes, shape, peer)))
        .collect();
    let updates: usize = wire.iter().map(Vec::len).sum();
    let routes = shape.prefixes() * peers as usize;
    let wire_bytes: usize = wire.iter().flatten().map(BytesMut::len).sum();
    println!(
        "{updates} UPDATE messages, {:.1} prefixes each, {:.1} MiB on the wire, {routes} routes",
        routes as f64 / updates as f64,
        mib(wire_bytes),
    );

    println!("\n== Reading the messages (one thread)");
    let ctx = wire_context();
    let started = Instant::now();
    for message in wire.iter().flatten() {
        black_box(Message::decode(&ctx, message).expect("what we wrote is well formed"));
    }
    let elapsed = started.elapsed();
    row(
        "decode UPDATEs",
        elapsed,
        format!(
            "{} routes, {} messages, {:.0} MiB/s",
            per_second(routes, elapsed),
            per_second(updates, elapsed),
            mib(wire_bytes) / elapsed.as_secs_f64()
        ),
    );

    println!("\n== From bytes to best routes, without the network");
    for threads in [1, threads] {
        into_table(&wire, routes, threads);
    }

    println!("\n== End to end: {peers} peers send the table over TCP at once");
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("a runtime");
    runtime.block_on(async {
        let server = TestServer::start(65000).await.expect("a server");
        let mut senders = Vec::new();
        for (index, messages) in wire.iter().enumerate() {
            let peer = index as u32 + 1;
            let mut session = TestPeer::connect(server.addr).await.expect("a connection");
            session
                .handshake(64512 + peer, Ipv4Addr::new(10, 0, 0, peer as u8), 90)
                .await
                .expect("a session");
            let stream: Vec<u8> = messages.iter().flat_map(|m| m.iter().copied()).collect();
            senders.push((session, stream));
        }
        let started = Instant::now();
        let sending: Vec<_> = senders
            .into_iter()
            .map(|(mut session, stream)| {
                tokio::spawn(async move {
                    session.send_raw(&stream).await.expect("the server reads");
                    // Held open until the table is in.
                    session
                })
            })
            .collect();
        while server.rib.len() < routes {
            tokio::time::sleep(Duration::from_millis(2)).await;
        }
        let elapsed = started.elapsed();
        row(
            "sessions, state machine, table",
            elapsed,
            format!(
                "{} routes; {} prefixes with a best route",
                per_second(routes, elapsed),
                server.rib.prefix_count()
            ),
        );
        for session in sending {
            drop(session.await);
        }
    });
}

/// What a session does with each message, without the session: decode it,
/// find its attribute set, put its routes in the table, choose the best.
/// `threads` threads each take whole peers.
fn into_table(wire: &[Vec<BytesMut>], routes: usize, threads: usize) {
    let region = Region::begin();
    let table = RibTable::new(threads * 16);
    let peers: Vec<PeerIndex> = (0..wire.len())
        .map(|peer| {
            table
                .add_peer(PeerInfo {
                    asn: 64512 + peer as u32 + 1,
                    router_id: Ipv4Addr::from(peer as u32 + 1),
                    address: Ipv4Addr::new(192, 0, 2, peer as u8 + 1).into(),
                    ibgp: false,
                })
                .expect("room for the peer")
        })
        .collect();
    let ctx = wire_context();
    let next = AtomicUsize::new(0);
    let started = Instant::now();
    std::thread::scope(|scope| {
        for _ in 0..threads {
            scope.spawn(|| {
                loop {
                    let at = next.fetch_add(1, Ordering::Relaxed);
                    let Some(messages) = wire.get(at) else {
                        break;
                    };
                    for message in messages {
                        let decoded =
                            Message::decode_with_report(&ctx, message).expect("well formed");
                        let Message::Update(update) = decoded.message else {
                            continue;
                        };
                        let wire = decoded.attributes_wire.as_deref().expect("kept");
                        if !update.announced_routes.is_empty() {
                            let attrs = table.attrs().intern_received(wire, None).expect("valid");
                            table.announce(peers[at], &attrs, &update.announced_routes);
                        }
                        for attribute in &update.path_attributes {
                            if let PathAttribute::MpReachNlri(mp) = attribute {
                                let attrs = table
                                    .attrs()
                                    .intern_received(wire, Some(&mp.next_hop))
                                    .expect("valid");
                                table.announce(peers[at], &attrs, &mp.prefixes);
                            }
                        }
                    }
                }
            });
        }
    });
    let elapsed = started.elapsed();
    assert_eq!(table.len(), routes);
    row(
        &format!("decode, intern, insert, choose: {threads} thread(s)"),
        elapsed,
        format!(
            "{} routes; {}",
            per_second(routes, elapsed),
            describe(region.held(), routes, "route")
        ),
    );
}