//! What passing routes on costs: one peer sends the table over a session,
//! and another is told of every route, over its own. The time is from the
//! first byte sent to the last route heard; alongside, the time for the
//! same table to be in our own table and no further.

use std::net::Ipv4Addr;
use std::sync::Arc;
use std::time::{Duration, Instant};

use bgp_bench::{Prefixes, Shape, announcements, encoded, env};
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::message::PathAttribute;
use bgp_server::marks::{Control, Marks};
use bgp_server::policy::ExportPolicy;
use bytes::BytesMut;

/// How a server passes routes on in one run.
#[derive(Clone, Copy)]
enum Passing {
    Not,
    Plain,
    /// With the communities that say who is told, and a mark put on every
    /// route.
    Dressed,
}

/// Sends the table from one peer and, where it is passed on, hears all of
/// it at another. Returns when the table was in, and when it was all
/// heard.
fn run(
    wire: &[BytesMut],
    routes: usize,
    heard_routes: usize,
    passing: Passing,
) -> (Duration, Option<Duration>) {
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("a runtime");
    runtime.block_on(async {
        let server = TestServer::start_with(4_200_000_000, |config| {
            config.export = Arc::new(ExportPolicy {
                learned_routes: !matches!(passing, Passing::Not),
                all_paths: false,
                marks: match passing {
                    Passing::Dressed => Marks::parse(&[], &["64601:7".to_string()], &[]).unwrap(),
                    _ => Marks::default(),
                },
                prepend: 0,
            });
            config.control = matches!(passing, Passing::Dressed).then_some(Control {
                local_asn: 4_200_000_000,
            });
        })
        .await
        .expect("a server");
        let mut hearer = TestPeer::connect(server.addr).await.expect("a connection");
        hearer
            // An AS number that is on none of the generated paths: a route
            // is not passed to a peer whose AS it has been through.
            .handshake(4_100_000_000, Ipv4Addr::new(10, 0, 0, 100), 90)
            .await
            .expect("a session");
        let mut sender = TestPeer::connect(server.addr).await.expect("a connection");
        sender
            .handshake(64513, Ipv4Addr::new(10, 0, 0, 1), 90)
            .await
            .expect("a session");
        let stream: Vec<u8> = wire.iter().flat_map(|m| m.iter().copied()).collect();
        let started = Instant::now();
        let sending = tokio::spawn(async move {
            sender.send_raw(&stream).await.expect("the server reads");
            sender
        });
        let hearing = tokio::spawn(async move {
            if matches!(passing, Passing::Not) {
                return (hearer, None, 0);
            }
            let mut heard = 0usize;
            let mut messages = 0usize;
            while heard < heard_routes {
                let update = hearer.recv_update().await.expect("an UPDATE");
                messages += 1;
                heard += update.announced_routes.len();
                heard += update
                    .path_attributes
                    .iter()
                    .map(|attribute| match attribute {
                        PathAttribute::MpReachNlri(mp) => mp.prefixes.len(),
                        _ => 0,
                    })
                    .sum::<usize>();
            }
            (hearer, Some(started.elapsed()), messages)
        });
        while server.rib.len() < routes {
            tokio::time::sleep(Duration::from_millis(2)).await;
        }
        let in_table = started.elapsed();
        let (hearer, heard, messages) = hearing.await.expect("the hearer");
        if messages > 0 {
            println!(
                "    ({messages} UPDATEs, {:.1} routes each)",
                heard_routes as f64 / messages as f64
            );
        }
        drop(sending.await);
        drop(hearer);
        (in_table, heard)
    })
}

fn main() {
    let shape = Shape::INTERNET.scaled_down(env("BENCH_SCALE_DOWN", 1));
    let prefixes = Prefixes::generate(shape);
    let wire = encoded(&announcements(&prefixes, shape, 1));
    let routes = shape.prefixes();
    // Over a session on IPv4 only the IPv4 routes go.
    let heard_routes = shape.ipv4_prefixes;
    println!(
        "== One peer sends {routes} routes in {} UPDATEs over TCP; another, over IPv4, is told of the {heard_routes} IPv4 ones",
        wire.len()
    );
    for (what, passing) in [
        ("not passed on", Passing::Not),
        ("passed on", Passing::Plain),
        (
            "passed on, with control communities and a mark",
            Passing::Dressed,
        ),
    ] {
        let (in_table, heard) = run(&wire, routes, heard_routes, passing);
        let heard = heard.map_or(String::new(), |heard| {
            format!(
                "   all heard by the other peer after {:>6.3} s ({:.2} M routes/s)",
                heard.as_secs_f64(),
                heard_routes as f64 / heard.as_secs_f64() / 1e6
            )
        });
        println!(
            "{what:<48} in the table after {:>6.3} s{heard}",
            in_table.as_secs_f64()
        );
    }
}