//! A synthetic copy of the Internet routing table, as the UPDATE messages a
//! peer would send to announce it.
//!
//! It is made up, not captured, so that the benchmarks need no data files
//! and give the same input every run. It is shaped like the real thing where
//! that matters for cost:
//!
//! - about a million IPv4 prefixes, most of them /24s, and about a quarter
//!   of a million IPv6 prefixes, most of them /48s;
//! - prefixes that share their attributes are announced together, a handful
//!   to an UPDATE, as routers do;
//! - AS paths a few hops long, and communities on most routes;
//! - several peers announce the same prefixes, each with attributes of its
//!   own.

pub mod memory;
pub mod mrt;

use std::collections::HashSet;
use std::net::{Ipv4Addr, Ipv6Addr};

use bgp_packet::message::{
    AsPathAttribute, AsPathSegment, CommunitiesPathAttribute, LargeCommunitiesPathAttribute,
    LargeCommunity, Message, MpReachNlriPathAttribute, MultiExitDiscPathAttribute,
    NextHopPathAttribute, NlriNextHop, OriginPathAttribute, PathAttribute, UpdateMessage,
};
use bgp_primitives::constants::{AfiSafi, IpVersion};
use bgp_primitives::ip_prefix::{IpPrefix, Nlri};
use bgp_primitives::parser::ParserContext;
use bytes::BytesMut;

/// How big a table to make.
#[derive(Debug, Clone, Copy)]
pub struct Shape {
    pub ipv4_prefixes: usize,
    pub ipv6_prefixes: usize,
    /// How many prefixes share one set of attributes, on average.
    pub prefixes_per_attribute_set: usize,
}

impl Shape {
    /// Roughly the default-free zone in 2026.
    pub const INTERNET: Self = Self {
        ipv4_prefixes: 1_000_000,
        ipv6_prefixes: 230_000,
        prefixes_per_attribute_set: 6,
    };

    pub fn prefixes(&self) -> usize {
        self.ipv4_prefixes + self.ipv6_prefixes
    }

    /// The same shape, `factor` times smaller.
    pub fn scaled_down(&self, factor: usize) -> Self {
        Self {
            ipv4_prefixes: self.ipv4_prefixes / factor,
            ipv6_prefixes: self.ipv6_prefixes / factor,
            ..*self
        }
    }
}

/// A small deterministic generator (SplitMix64); no need for a dependency.
pub struct Rng(u64);

impl Rng {
    pub fn new(seed: u64) -> Self {
        Self(seed)
    }

    pub fn next_u64(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
        z ^ (z >> 31)
    }

    /// A number below `bound`.
    pub fn below(&mut self, bound: u64) -> u64 {
        self.next_u64() % bound
    }

    /// An entry of `weighted`, picked with probability proportional to its
    /// weight.
    fn pick<T: Copy>(&mut self, weighted: &[(T, u64)]) -> T {
        let total: u64 = weighted.iter().map(|(_, w)| w).sum();
        let mut at = self.below(total);
        for (value, weight) in weighted {
            if at < *weight {
                return *value;
            }
            at -= weight;
        }
        weighted[weighted.len() - 1].0
    }
}

/// The prefixes of the table. Every peer announces these same prefixes.
pub struct Prefixes {
    pub ipv4: Vec<IpPrefix>,
    pub ipv6: Vec<IpPrefix>,
}

impl Prefixes {
    pub fn generate(shape: Shape) -> Self {
        let mut rng = Rng::new(0x1234_5678);
        // Shares of prefix lengths, close to what the table looks like.
        let v4_lengths: [(u8, u64); 9] = [
            (24, 60),
            (23, 10),
            (22, 11),
            (21, 5),
            (20, 4),
            (19, 3),
            (18, 2),
            (16, 4),
            (15, 1),
        ];
        let v6_lengths: [(u8, u64); 7] = [
            (48, 48),
            (32, 14),
            (40, 10),
            (44, 10),
            (36, 6),
            (29, 6),
            (47, 6),
        ];

        let mut seen = HashSet::new();
        let mut ipv4 = Vec::with_capacity(shape.ipv4_prefixes);
        while ipv4.len() < shape.ipv4_prefixes {
            let length = rng.pick(&v4_lengths);
            // Unicast space, 1.0.0.0 to 223.255.255.255.
            let address = 0x0100_0000 + (rng.next_u64() as u32) % 0xdf00_0000;
            let prefix = IpPrefix::new(IpVersion::Ipv4, address.to_be_bytes(), length)
                .expect("four bytes cover any IPv4 prefix");
            if seen.insert(prefix) {
                ipv4.push(prefix);
            }
        }
        let mut ipv6 = Vec::with_capacity(shape.ipv6_prefixes);
        while ipv6.len() < shape.ipv6_prefixes {
            let length = rng.pick(&v6_lengths);
            // Global unicast, 2000::/3.
            let address = (0x2000_u128 << 112) | (u128::from(rng.next_u64()) << 61 >> 3);
            let prefix = IpPrefix::new(IpVersion::Ipv6, address.to_be_bytes(), length)
                .expect("sixteen bytes cover any IPv6 prefix");
            if seen.insert(prefix) {
                ipv6.push(prefix);
            }
        }
        Self { ipv4, ipv6 }
    }
}

/// The UPDATE messages with which the peer numbered `peer` announces the
/// whole table. The peer's AS is `64512 + peer`.
pub fn announcements(prefixes: &Prefixes, shape: Shape, peer: u32) -> Vec<UpdateMessage> {
    let mut rng = Rng::new(0x9999 + u64::from(peer));
    let peer_asn = 64512 + peer;
    let attributes = |rng: &mut Rng| {
        let hops = 2 + rng.below(4) as usize;
        let mut path = vec![peer_asn];
        path.extend((0..hops).map(|_| 1 + rng.below(400_000) as u32));
        let mut attributes = vec![
            PathAttribute::Origin(OriginPathAttribute::IGP),
            PathAttribute::AsPath(AsPathAttribute {
                segments: vec![AsPathSegment {
                    ordered: true,
                    path,
                }],
            }),
        ];
        let mut later = Vec::new();
        if rng.below(4) == 0 {
            later.push(PathAttribute::MultiExitDisc(MultiExitDiscPathAttribute(
                rng.below(1000) as u32,
            )));
        }
        let communities = rng.below(7) as usize;
        if communities > 0 {
            later.push(PathAttribute::Communities(CommunitiesPathAttribute(
                (0..communities)
                    .map(|_| (rng.below(65000) as u16, rng.below(1000) as u16))
                    .collect(),
            )));
        }
        let large = if rng.below(5) == 0 {
            Some(PathAttribute::LargeCommunities(
                LargeCommunitiesPathAttribute {
                    communities: vec![LargeCommunity {
                        global_admin: peer_asn,
                        data_1: rng.below(100) as u32,
                        data_2: rng.below(100) as u32,
                    }],
                },
            ))
        } else {
            None
        };
        (std::mem::take(&mut attributes), later, large)
    };

    let mut updates = Vec::new();
    // Prefixes that share attributes travel together. A set's prefixes are a
    // run of the (randomly ordered) list, a few long.
    let mut grouped = |list: &[IpPrefix], rng: &mut Rng, ipv6: bool| {
        let mut rest = list;
        while !rest.is_empty() {
            let size = (1 + rng.below(2 * shape.prefixes_per_attribute_set as u64 - 1) as usize)
                .min(rest.len());
            let (group, after) = rest.split_at(size);
            rest = after;

            let (mut path_attributes, later, large) = attributes(rng);
            let update = if ipv6 {
                // In order of type code: MED and COMMUNITIES come before
                // MP_REACH_NLRI.
                path_attributes.extend(later);
                path_attributes.push(PathAttribute::MpReachNlri(MpReachNlriPathAttribute {
                    family: AfiSafi::IPV6_UNICAST,
                    next_hop: NlriNextHop::Ipv6(Ipv6Addr::new(
                        0x2001,
                        0xdb8,
                        0,
                        0,
                        0,
                        0,
                        0,
                        peer as u16,
                    )),
                    prefixes: Nlri::plain(group.iter().copied()),
                }));
                path_attributes.extend(large);
                UpdateMessage {
                    withdrawn_routes: vec![],
                    path_attributes,
                    announced_routes: vec![],
                }
            } else {
                path_attributes.push(PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::new(
                    192, 0, 2, peer as u8,
                ))));
                path_attributes.extend(later);
                path_attributes.extend(large);
                UpdateMessage {
                    withdrawn_routes: vec![],
                    path_attributes,
                    announced_routes: Nlri::plain(group.iter().copied()),
                }
            };
            updates.push(update);
        }
    };
    grouped(&prefixes.ipv4, &mut rng, false);
    grouped(&prefixes.ipv6, &mut rng, true);
    updates
}

/// `updates` as they are on the wire, on a session with four-octet AS
/// numbers.
pub fn encoded(updates: &[UpdateMessage]) -> Vec<BytesMut> {
    let ctx = wire_context();
    updates
        .iter()
        .map(|update| {
            let mut wire = BytesMut::new();
            Message::Update(update.clone())
                .to_wire(&ctx, &mut wire)
                .expect("a few prefixes fit one message");
            wire
        })
        .collect()
}

/// The context of the sessions the benchmarks pretend to have: an external
/// peer with four-octet AS numbers.
pub fn wire_context() -> ParserContext {
    ParserContext {
        four_octet_asn: true,
        ibgp: false,
        add_path: Vec::new(),
    }
}