//! The same questions as `full_table`, asked of a real table dump: every
//! route that the peers of a RIPE RIS route collector had announced to it.
//!
//! 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, so that
//! `cargo bench` works on a fresh checkout. `BENCH_PEERS` limits the run to
//! the peers with the most routes (default: all of them).

use std::collections::HashSet;
use std::net::Ipv4Addr;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Instant;

use bgp_bench::memory::Region;
use bgp_bench::mrt::{Dump, parse_attributes};
use bgp_bench::report::{mib, per_second, row};
use bgp_packet::message::{NlriNextHop, PathAttribute};
use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rib::table::{Event, PeerIndex, PeerInfo, RibTable};

/// Bytes allocated and not freed since `region` began; zero when memory
/// is not being counted (see `bgp_bench::memory`), and zero too in the
/// odd case of more freed than allocated, which is nothing held.
fn retained(region: &Region) -> usize {
    usize::try_from(region.held().unwrap_or(0)).unwrap_or(0)
}

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 = (Arc<[PathAttribute]>, Option<NlriNextHop>);

fn main() {
    if Region::begin().held().is_none() {
        println!(
            "Memory is not counted in this run: sizes print as 0. For sizes (and slower\ntimings) add --features count-memory.\n"
        );
    }
    let Some(path) = dump_path() else {
        println!("No table dump found; see crates/bench/data/README.md. Nothing to do.");
        return;
    };
    println!("== Reading {}", path.display());
    let started = Instant::now();
    let dump = Dump::read(&path).expect("a readable dump");
    println!(
        "{} peers, {} routes, {} distinct attribute sets, read in {:.0} s",
        dump.peers.len(),
        dump.route_count(),
        dump.attribute_sets.len(),
        started.elapsed().as_secs_f64()
    );

    let limit = std::env::var("BENCH_PEERS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(dump.peers.len());
    let peers = dump.largest_peers(limit);
    let peers: Vec<usize> = peers
        .into_iter()
        .filter(|&peer| !dump.routes[peer].is_empty())
        .collect();
    describe(&dump, &peers);

    // Parsed once, outside the timings: reading messages is measured by
    // `full_table`, and here the attributes arrive already parsed.
    let region = Region::begin();
    let parsed: Vec<Option<Attributes>> = dump
        .attribute_sets
        .iter()
        .map(|wire| parse_attributes(wire).map(|(a, next_hop)| (Arc::from(a), next_hop)))
        .collect();
    let parsed_bytes = retained(&region);
    println!(
        "\nattribute sets, parsed (Arc<[PathAttribute]>)    {:.0} MiB = {:.0} bytes a set",
        mib(parsed_bytes),
        parsed_bytes as f64 / parsed.len() as f64
    );
    let unparsed = parsed.iter().filter(|a| a.is_none()).count();
    if unparsed > 0 {
        println!("{unparsed} attribute sets could not be parsed; their routes are left out");
    }

    let threads = std::thread::available_parallelism().map_or(1, |n| n.get());
    sharded_table(&dump, &peers, &parsed, 1);
    sharded_table(&dump, &peers, &parsed, threads);
    followed_table(&dump, &peers, &parsed, threads);
}

/// What the table looks like, as far as it decides what a RIB costs.
fn describe(dump: &Dump, peers: &[usize]) {
    let routes: usize = peers.iter().map(|&p| dump.routes[p].len()).sum();
    let full = peers
        .iter()
        .filter(|&&p| dump.routes[p].len() > 800_000)
        .count();
    let prefixes: HashSet<IpPrefix> = peers
        .iter()
        .flat_map(|&p| dump.routes[p].iter().map(|(prefix, _)| *prefix))
        .collect();
    let mut sets_used = HashSet::new();
    let mut sets_per_peer = 0usize;
    for &peer in peers {
        let own: HashSet<u32> = dump.routes[peer].iter().map(|(_, set)| *set).collect();
        sets_per_peer += own.len();
        sets_used.extend(own);
    }
    let attribute_bytes: usize = sets_used
        .iter()
        .map(|&set| dump.attribute_sets[set as usize].len())
        .sum();
    println!("\n== The table ({} peers used)", peers.len());
    println!("routes                                       {routes:>12}");
    println!("peers with more than 800,000 routes          {full:>12}");
    println!(
        "distinct prefixes                            {:>12}",
        prefixes.len()
    );
    println!(
        "routes per prefix                            {:>12.1}",
        routes as f64 / prefixes.len() as f64
    );
    println!(
        "distinct attribute sets                      {:>12}",
        sets_used.len()
    );
    println!(
        "  if each peer's were kept apart             {sets_per_peer:>12}   ({:.2}x: sets shared between peers)",
        sets_per_peer as f64 / sets_used.len() as f64
    );
    println!(
        "routes per attribute set                     {:>12.1}",
        routes as f64 / sets_used.len() as f64
    );
    println!(
        "attribute set on the wire, mean              {:>12.0} bytes",
        attribute_bytes as f64 / sets_used.len() as f64
    );
}

/// The sharded table, filled by `threads` threads. Each takes whole peers,
/// as sessions would, and for every UPDATE interns the attributes and then
/// announces the prefixes: everything between a decoded message and a
/// chosen best route.
fn sharded_table(dump: &Dump, peers: &[usize], parsed: &[Option<Attributes>], threads: usize) {
    println!("\n== The sharded table, {threads} thread(s)");
    let updates = grouped(dump, peers, parsed);
    let routes: usize = updates
        .iter()
        .flatten()
        .map(|(_, prefixes)| prefixes.len())
        .sum();

    let region = Region::begin();
    let table = Arc::new(RibTable::new(threads * 16));
    let indices = add_peers(&table, dump, peers);
    let started = Instant::now();
    let best_changes = load(&table, dump, &updates, &indices, threads);
    let elapsed = started.elapsed();
    assert_eq!(table.len(), routes);
    let held = retained(&region);
    row(
        "intern attributes, insert, choose best",
        elapsed,
        format!(
            "{} routes; {:.0} MiB = {:.0} bytes a route, attribute sets included",
            per_second(routes, elapsed),
            mib(held),
            held as f64 / routes as f64
        ),
    );
    println!(
        "    {} prefixes, {} attribute sets, best route changed {} times",
        table.prefix_count(),
        table.attrs().len(),
        best_changes
    );

    let started = Instant::now();
    let mut best = 0usize;
    table.for_each_best(|prefix, _, _| {
        best += usize::from(std::hint::black_box(prefix).length() > 0)
    });
    row(
        "read every best route",
        started.elapsed(),
        format!("{best} of them, in place, table not stopped"),
    );

    // Reading has to look at each route for the timing to mean anything.
    let started = Instant::now();
    let mut all = 0usize;
    let mut hops = 0u64;
    table.for_each_route(|route| {
        all += 1;
        hops += u64::from(route.attrs.key().as_path_len);
    });
    row(
        "read every route and its attributes",
        started.elapsed(),
        format!(
            "{all} of them, mean AS path {:.1}",
            hops as f64 / all as f64
        ),
    );

    let started = Instant::now();
    let outcome = table.remove_peer(indices[0]);
    let elapsed = started.elapsed();
    row(
        "remove one peer (session down)",
        elapsed,
        format!(
            "{} routes; best route changed for {} prefixes",
            per_second(outcome.removed, elapsed),
            outcome.best_changed
        ),
    );
    let started = Instant::now();
    let swept = table.attrs().sweep();
    row(
        "drop attribute sets no longer used",
        started.elapsed(),
        format!("{swept} dropped"),
    );
}

/// A peer's routes grouped as UPDATEs: the attribute set and the prefixes
/// announced with it.
type Updates = Vec<Vec<(u32, Vec<IpPrefix>)>>;

fn grouped(dump: &Dump, peers: &[usize], parsed: &[Option<Attributes>]) -> Updates {
    peers
        .iter()
        .map(|&peer| {
            let mut routes = dump.routes[peer].clone();
            routes.sort_by_key(|(_, set)| *set);
            routes
                .chunk_by(|a, b| a.1 == b.1)
                .filter(|group| parsed[group[0].1 as usize].is_some())
                .map(|group| {
                    (
                        group[0].1,
                        group.iter().map(|(prefix, _)| *prefix).collect(),
                    )
                })
                .collect()
        })
        .collect()
}

fn add_peers(table: &RibTable, dump: &Dump, peers: &[usize]) -> Vec<PeerIndex> {
    peers
        .iter()
        .map(|&peer| {
            table
                .add_peer(PeerInfo {
                    asn: dump.peers[peer].asn,
                    router_id: Ipv4Addr::from(peer as u32 + 1),
                    address: dump.peers[peer].address,
                    interface: 0,
                    ibgp: false,
                    reflector_client: false,
                })
                .expect("room for the peer")
        })
        .collect()
}

/// Announces everything in `updates` on `threads` threads, each taking
/// whole peers. Returns how many times a best route changed.
fn load(
    table: &RibTable,
    dump: &Dump,
    updates: &Updates,
    indices: &[PeerIndex],
    threads: usize,
) -> usize {
    let next = AtomicUsize::new(0);
    let best_changes = AtomicUsize::new(0);
    std::thread::scope(|scope| {
        for _ in 0..threads {
            scope.spawn(|| {
                let mut changes = 0;
                loop {
                    let at = next.fetch_add(1, Ordering::Relaxed);
                    let Some(updates) = updates.get(at) else {
                        break;
                    };
                    for (set, prefixes) in updates {
                        // From the bytes, as a dump has them and as a
                        // session could hand them over without parsing.
                        let Some(attrs) = table
                            .attrs()
                            .intern_wire(&dump.attribute_sets[*set as usize])
                        else {
                            continue;
                        };
                        changes += table.announce(indices[at], &attrs, prefixes).best_changed;
                    }
                }
                best_changes.fetch_add(changes, Ordering::Relaxed);
            });
        }
    });
    best_changes.load(Ordering::Relaxed)
}

/// The same load with someone following the best routes through the
/// journal: once reading as fast as it can while the table fills, once not
/// reading at all until the table is full.
fn followed_table(dump: &Dump, peers: &[usize], parsed: &[Option<Attributes>], threads: usize) {
    println!("\n== The sharded table with a follower, {threads} thread(s)");
    let updates = grouped(dump, peers, parsed);
    let routes: usize = updates
        .iter()
        .flatten()
        .map(|(_, prefixes)| prefixes.len())
        .sum();
    let count = |counts: &mut (usize, usize, usize), event: Event<'_>| match event {
        Event::Best { attrs, .. } => {
            // Look at the route, as a follower would.
            counts.0 += usize::from(attrs.key().as_path_len > 0);
        }
        Event::Withdrawn { .. } => counts.1 += 1,
        Event::Reset | Event::Synced => counts.2 += 1,
        // Not given to a follower of the best routes.
        Event::Routes { .. } => {}
    };

    // Reading while the table fills.
    let table = Arc::new(RibTable::new(threads * 16));
    let indices = add_peers(&table, dump, peers);
    let done = AtomicBool::new(false);
    let started = Instant::now();
    let (best_changes, counts, longest, after_load) = std::thread::scope(|scope| {
        let follower = scope.spawn(|| {
            let mut reader = table.reader();
            let mut counts = (0, 0, 0);
            let mut longest = 0;
            while !done.load(Ordering::Acquire) {
                reader.poll(4096, |event| count(&mut counts, event));
                longest = longest.max(table.journal_len());
            }
            let finishing = Instant::now();
            while !reader
                .poll(4096, |event| count(&mut counts, event))
                .caught_up
            {}
            (counts, longest, finishing.elapsed())
        });
        let best_changes = load(&table, dump, &updates, &indices, threads);
        done.store(true, Ordering::Release);
        let (counts, longest, after_load) = follower.join().expect("the follower runs");
        (best_changes, counts, longest, after_load)
    });
    let elapsed = started.elapsed();
    assert_eq!(table.len(), routes);
    row(
        "load, follower reading throughout",
        elapsed,
        format!(
            "{} routes; caught up {:.3} s after the last route",
            per_second(routes, elapsed),
            after_load.as_secs_f64()
        ),
    );
    println!(
        "    best route changed {best_changes} times; the follower was given {} events for it; \n    journal at most {longest} entries for {} prefixes",
        counts.0 + counts.1,
        table.prefix_count()
    );
    drop(table);

    // Not reading: the journal has to stay small on its own account.
    let table = Arc::new(RibTable::new(threads * 16));
    let indices = add_peers(&table, dump, peers);
    let mut reader = table.reader();
    let mut counts = (0, 0, 0);
    reader.poll(usize::MAX, |event| count(&mut counts, event));
    let started = Instant::now();
    load(&table, dump, &updates, &indices, threads);
    let elapsed = started.elapsed();
    row(
        "load, follower not reading",
        elapsed,
        format!(
            "{} routes; journal holds {} entries for {} prefixes",
            per_second(routes, elapsed),
            table.journal_len(),
            table.prefix_count()
        ),
    );
    let started = Instant::now();
    let mut counts = (0, 0, 0);
    while !reader
        .poll(4096, |event| count(&mut counts, event))
        .caught_up
    {}
    let elapsed = started.elapsed();
    row(
        "the follower then catches up",
        elapsed,
        format!(
            "{} events, {} of them a reset",
            counts.0 + counts.1,
            counts.2
        ),
    );
}