//! What the operator's policy script costs, on a real table: the routes
//! one peer of a RIPE RIS route collector had announced to it, each with
//! the attributes it really has.
//!
//! Needs a dump in `crates/bench/data/` (see the README there) or at the
//! path in `BENCH_MRT`; without one it says so and does nothing.
//!
//! A function is asked about a route once for every UPDATE if it does not
//! look at the prefix, and once for every prefix if it does. How many
//! UPDATEs a table arrives in is the sender's business; the two ends are
//! one UPDATE for every distinct set of attributes, and one for every
//! route. Both are measured, with what the same work costs without a
//! script.

use std::collections::HashSet;
use std::hint::black_box;
use std::path::PathBuf;
use std::time::{Duration, Instant};

use std::net::Ipv4Addr;

use bgp_bench::mrt::{Dump, parse_attributes};
use bgp_bench::{Prefixes, Shape, announcements, encoded};
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::message::{NlriNextHop, PathAttribute};
use bgp_policy::{Direction, Hook, Peer, Policy, Subject};
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_server::config::PeerConfig;
use bgp_server::family::FamilySet;
use bgp_server::policy::ImportPolicy;
use bytes::BytesMut;

const SCRIPT: &str = r#"
    const BLACKHOLE = "65535:666";
    const LEARNED_FROM = #{ customer: "64601:1000", peer: "64601:2000", provider: "64601:3000" };
    const PREFERENCE = #{ customer: 300, peer: 200, provider: 100 };

    // Nothing at all: what asking costs.
    fn nothing(route, peer) { true }

    // Looks, and changes nothing.
    fn filter(route, peer) {
        if route.has_community(BLACKHOLE) { return false; }
        if route.as_path_len > 30 { return false; }
        if route.path_has(64512) { return false; }
        true
    }

    // What a network would have for its transit: refuse some, mark and
    // rank the rest.
    fn typical(route, peer) {
        if route.has_community(BLACKHOLE) { return false; }
        if route.as_path_len > 30 { return false; }
        route.remove_community("64601:*");
        route.add_community(LEARNED_FROM[peer.role]);
        route.local_pref = PREFERENCE[peer.role];
        true
    }

    // For any peer, whatever it is to us.
    fn marks(route, peer) {
        if route.has_community(BLACKHOLE) { return false; }
        route.remove_community("64601:*");
        route.add_community("64601:3000");
        route.local_pref = 100;
        true
    }

    fn marks_by_prefix(route, peer) {
        if route.prefix_len > 24 && route.family == "ipv4" { return false; }
        marks(route, peer)
    }

    // The same, and it looks at the prefix: asked for every route.
    fn by_prefix(route, peer) {
        if route.prefix_len > 24 && route.family == "ipv4" { return false; }
        if route.has_community(BLACKHOLE) { return false; }
        route.remove_community("64601:*");
        route.add_community(LEARNED_FROM[peer.role]);
        route.local_pref = PREFERENCE[peer.role];
        true
    }
"#;

fn dump_path() -> Option<PathBuf> {
    if let Ok(path) = std::env::var("BENCH_MRT") {
        return Some(path.into());
    }
    let data = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("data");
    let mut dumps: Vec<PathBuf> = std::fs::read_dir(data)
        .ok()?
        .filter_map(|entry| Some(entry.ok()?.path()))
        .filter(|path| path.extension().is_some_and(|e| e == "gz"))
        .collect();
    dumps.sort();
    dumps.pop()
}

type Attributes = (Vec<PathAttribute>, Option<NlriNextHop>);

struct Outcome {
    elapsed: Duration,
    asked: usize,
    taken: usize,
    changed: usize,
}

/// Asks `hook` about each of `routes`.
fn ask<'a>(
    hook: &Hook,
    peer: &Peer,
    routes: impl Iterator<Item = (&'a IpPrefix, &'a Attributes)>,
) -> Outcome {
    let started = Instant::now();
    let (mut asked, mut taken, mut changed) = (0, 0, 0);
    for (prefix, (attributes, next_hop)) in routes {
        let verdict = hook
            .run(
                peer,
                &Subject {
                    prefix,
                    attributes,
                    next_hop: *next_hop,
                    direction: Direction::Import,
                },
            )
            .expect("the function runs");
        asked += 1;
        taken += usize::from(verdict.accept);
        changed += usize::from(verdict.attributes.is_some());
        black_box(&verdict);
    }
    Outcome {
        elapsed: started.elapsed(),
        asked,
        taken,
        changed,
    }
}

fn row(what: &str, outcome: &Outcome) {
    println!(
        "{what:<52} {:>8.3} s {:>7.0} ns each   {} asked, {} taken, {} changed",
        outcome.elapsed.as_secs_f64(),
        outcome.elapsed.as_nanos() as f64 / outcome.asked.max(1) as f64,
        outcome.asked,
        outcome.taken,
        outcome.changed
    );
}

/// One peer sends the whole (synthetic) table over TCP to a server that
/// gives its routes to `function`, or to none. Returns how long it took
/// for all of them to be in the table.
fn end_to_end(wire: &[BytesMut], routes: usize, function: Option<&str>) -> Duration {
    const PEER_ASN: u32 = 64513;
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("a runtime");
    runtime.block_on(async {
        // The configured peer is whoever connects from this machine; the
        // listener is where the server's own attempts to reach it go.
        let listener = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0))
            .await
            .expect("a port");
        let endpoint = listener.local_addr().expect("its address");
        let policy = Policy::compile(SCRIPT).expect("the script");
        let import_hook = function.map(|name| policy.hook(name).expect("the function"));
        let server = TestServer::start_with(4_200_000_000, |config| {
            config.accept_unconfigured_peers = false;
            config.peers = vec![PeerConfig {
                name: "transit".into(),
                endpoint: endpoint.into(),
                remote_asn: PEER_ASN,
                hold_time_secs: 90,
                local_address: None,
                families: FamilySet::supported(),
                add_path: None,
                import: ImportPolicy::permissive(),
                role: None,
                require_role: false,
                import_hook,
                export_hook: None,
            }];
        })
        .await
        .expect("a server");
        let mut session = TestPeer::connect(server.addr).await.expect("a connection");
        session
            .handshake(PEER_ASN, 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 {
            session.send_raw(&stream).await.expect("the server reads");
            session
        });
        while server.rib.len() < routes {
            tokio::time::sleep(Duration::from_millis(2)).await;
        }
        let elapsed = started.elapsed();
        drop(sending.await);
        elapsed
    })
}

fn main() {
    // The whole way, first: a session, its state machine, the general
    // checks, the function, the table. On a table that is made up, so
    // that it runs anywhere.
    let shape = Shape::INTERNET;
    let prefixes = Prefixes::generate(shape);
    let wire = encoded(&announcements(&prefixes, shape, 1));
    let routes = shape.prefixes();
    println!(
        "== End to end: one peer sends {routes} routes in {} UPDATEs ({:.1} routes each) over TCP",
        wire.len(),
        routes as f64 / wire.len() as f64
    );
    let mut without = Duration::ZERO;
    for (what, function) in [
        ("no policy function", None),
        ("a function that does nothing", Some("nothing")),
        ("marks and ranks every route", Some("marks")),
        ("the same, and looks at the prefix", Some("marks_by_prefix")),
    ] {
        // The better of two: the first run of all also warms the machine.
        let elapsed = end_to_end(&wire, routes, function).min(end_to_end(&wire, routes, function));
        if function.is_none() {
            without = elapsed;
        }
        println!(
            "{what:<52} {:>8.3} s   {:.2} M routes/s   {:+.0}%",
            elapsed.as_secs_f64(),
            routes as f64 / elapsed.as_secs_f64() / 1e6,
            100.0 * (elapsed.as_secs_f64() / without.as_secs_f64() - 1.0)
        );
    }
    println!();

    let Some(path) = dump_path() else {
        println!("No table dump found; see crates/bench/data/README.md. Nothing more to do.");
        return;
    };
    println!("== Reading {}", path.display());
    let dump = Dump::read(&path).expect("a readable dump");
    let peer_index = dump.largest_peers(1)[0];
    let routes = &dump.routes[peer_index];
    let parsed: Vec<Option<Attributes>> = dump
        .attribute_sets
        .iter()
        .map(|wire| parse_attributes(wire))
        .collect();
    let routes: Vec<(&IpPrefix, &Attributes)> = routes
        .iter()
        .filter_map(|(prefix, set)| Some((prefix, parsed[*set as usize].as_ref()?)))
        .collect();
    let distinct: HashSet<*const Attributes> = routes
        .iter()
        .map(|(_, attributes)| std::ptr::from_ref(*attributes))
        .collect();
    let mut seen = HashSet::new();
    let once_per_set: Vec<(&IpPrefix, &Attributes)> = routes
        .iter()
        .filter(|(_, attributes)| seen.insert(std::ptr::from_ref(*attributes)))
        .copied()
        .collect();
    let communities: usize = routes
        .iter()
        .map(|(_, (attributes, _))| {
            attributes
                .iter()
                .map(|attribute| match attribute {
                    PathAttribute::Communities(communities) => communities.0.len(),
                    _ => 0,
                })
                .sum::<usize>()
        })
        .sum();
    println!(
        "The peer with the most routes, AS{}: {} routes with {} distinct sets of attributes \
         ({:.1} routes a set), {:.1} communities a route\n",
        dump.peers[peer_index].asn,
        routes.len(),
        distinct.len(),
        routes.len() as f64 / distinct.len() as f64,
        communities as f64 / routes.len() as f64
    );

    let policy = Policy::compile(SCRIPT).expect("the script");
    let peer = Peer::new(
        "transit",
        dump.peers[peer_index].address,
        dump.peers[peer_index].asn,
        64601,
        "provider",
    );

    // What it costs to go over the routes and do the work of the typical
    // function in Rust.
    let started = Instant::now();
    let mut taken = 0usize;
    for (_, (attributes, _)) in &routes {
        let mut communities: Vec<(u16, u16)> = attributes
            .iter()
            .find_map(|attribute| match attribute {
                PathAttribute::Communities(communities) => Some(communities.0.clone()),
                _ => None,
            })
            .unwrap_or_default();
        if communities.contains(&(65535, 666)) {
            continue;
        }
        communities.retain(|community| community.0 != 64601);
        communities.push((64601, 3000));
        let mut changed = attributes.clone();
        changed.push(PathAttribute::Communities(
            bgp_packet::message::CommunitiesPathAttribute(communities),
        ));
        black_box(&changed);
        taken += 1;
    }
    row(
        "the typical one in Rust, for every route",
        &Outcome {
            elapsed: started.elapsed(),
            asked: routes.len(),
            taken,
            changed: taken,
        },
    );
    println!();

    for function in ["nothing", "filter", "typical", "by_prefix"] {
        let hook = policy.hook(function).expect("the function");
        row(
            &format!("{function}, asked for every route"),
            &ask(&hook, &peer, routes.iter().copied()),
        );
        if function != "by_prefix" {
            row(
                &format!("{function}, asked once for each set of attributes"),
                &ask(&hook, &peer, once_per_set.iter().copied()),
            );
        }
    }
}