//! Putting the routes of a received UPDATE into the routing table, as far
//! as the peer's import policy lets them in.

use std::borrow::Cow;
use std::cell::OnceCell;
use std::collections::HashMap;
use std::sync::Arc;

use bgp_packet::message::{NlriNextHop, PathAttribute, UpdateMessage, only_to_customer_wire};
use bgp_primitives::constants::IpVersion;
use bgp_primitives::ip_prefix::{IpPrefix, Nlri};
use bgp_primitives::parser::EncodeError;
use bgp_rib::table::{AttrSet, Outcome, PeerIndex, RibTable};
use bgp_rpki::{PathValidity, Validity};

use crate::policy::{self, ImportPolicy, Rejection, RouteSource};
use crate::rejections::{Came, PeerRejections};
use crate::rpki::Checks;

/// The most routes of one peer that are held aside as invalid. Far more
/// than a peer sending the whole Internet has; a bound all the same.
const MAX_HELD: usize = 200_000;

/// What applying an UPDATE did.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub(crate) struct Applied {
    /// What changed in the table.
    pub(crate) outcome: Outcome,
    /// How many routes were announced, or found on a second look, that
    /// the policy refused.
    pub(crate) rejected: usize,
}

/// Whose UPDATE is being applied, and by what rules.
pub(crate) struct Import<'a> {
    pub(crate) peer: PeerIndex,
    pub(crate) from: RouteSource,
    pub(crate) policy: &'a ImportPolicy,
    /// What the peer's routes are validated against, with the data of the
    /// RPKI.
    pub(crate) checks: Checks<'a>,
    /// Where what is refused is noted.
    pub(crate) rejections: &'a PeerRejections,
}

/// A peer's routes that are not in the routing table because the RPKI's
/// data found them invalid: their origin (RFC 6811) or their path (ASPA).
///
/// Every other reason for refusing a route is a matter of the route and of
/// our configuration, and stays true for as long as the session lasts.
/// Validity is not: the data it rests on changes all the time, and a route
/// that is invalid now may be valid in ten minutes. So these routes are
/// kept, out of use, to be taken when that happens. The peer will not send
/// them again by itself.
#[derive(Default)]
pub(crate) struct Held {
    routes: HashMap<Nlri, Arc<AttrSet>>,
    /// Whether a route was not kept for want of room: then looking through
    /// `routes` again is not enough, and the peer has to be asked.
    overflowed: bool,
}

impl Held {
    pub(crate) fn len(&self) -> usize {
        self.routes.len()
    }

    /// Whether some invalid route was dropped instead of held; resets it.
    pub(crate) fn take_overflowed(&mut self) -> bool {
        std::mem::take(&mut self.overflowed)
    }

    fn hold(&mut self, route: Nlri, attrs: &Arc<AttrSet>) {
        if self.routes.len() >= MAX_HELD && !self.routes.contains_key(&route) {
            self.overflowed = true;
            return;
        }
        self.routes.insert(route, Arc::clone(attrs));
    }

    /// The peer has replaced or withdrawn these routes: what was held of
    /// them is no longer what the peer has.
    fn forget(&mut self, routes: &[Nlri]) {
        if self.routes.is_empty() {
            return;
        }
        for route in routes {
            self.routes.remove(route);
        }
    }
}

/// Applies `update` to `table`: its withdrawals first and then its
/// announcements, so that a prefix named in both ends up announced
/// (RFC 4271 §4.3).
///
/// An UPDATE names prefixes in four places: withdrawn routes and NLRI in the
/// message body, which are IPv4, and the MP_UNREACH_NLRI and MP_REACH_NLRI
/// attributes for every other family (RFC 4760). Where the peer sends
/// several routes for a prefix (ADD-PATH, RFC 7911), each is named by its
/// prefix and path identifier together, and is a route of its own.
///
/// A route the policy refuses is withdrawn instead: whatever the peer had
/// announced for it before is replaced by this, and this is not taken. One
/// refused as invalid by the RPKI's data is kept in `held`.
///
/// `attributes_wire` is the attributes as the decoder kept them, if it did
/// (see `DecodedUpdate::attributes_wire`). With it the table finds the
/// attribute set from those bytes, without writing the attributes out again.
pub(crate) fn apply(
    table: &RibTable,
    import: &Import<'_>,
    held: &mut Held,
    update: &UpdateMessage,
    attributes_wire: Option<&[u8]>,
) -> Result<Applied, EncodeError> {
    let mut applying = Applying {
        table,
        import,
        held,
        update,
        attributes_wire,
        applied: Applied::default(),
        // What the routes have in common is judged once.
        path_verdict: import
            .policy
            .check_path(&update.path_attributes, &import.from),
        path_invalid: OnceCell::new(),
        as_path: OnceCell::new(),
        refused_attributes: OnceCell::new(),
        mark: only_to_customer_mark(import, update),
        changed: OnceCell::new(),
        rules_holding: OnceCell::new(),
    };
    applying.withdraw();
    // Routes in the body have their next hop in the NEXT_HOP attribute,
    // which is one of the set's attributes; routes in MP_REACH_NLRI bring
    // theirs along.
    applying.announce(&update.announced_routes, None)?;
    for attribute in &update.path_attributes {
        if let PathAttribute::MpReachNlri(mp) = attribute {
            applying.announce(&mp.prefixes, Some(&mp.next_hop))?;
        }
    }
    Ok(applying.applied)
}

/// The AS to mark the routes of `update` with as ones that may from here
/// on only go down to customers, if they are to be marked.
///
/// A route from above or beside us is such a route. If nobody has marked
/// it so yet, we do (RFC 9234 §5, ingress, 3): the mark is part of the
/// route as it is kept, and goes wherever the route goes.
fn only_to_customer_mark(import: &Import<'_>, update: &UpdateMessage) -> Option<u32> {
    import
        .from
        .role
        .filter(|_| !import.from.ibgp)
        .and_then(|role| role.only_to_customer_on_import(import.from.peer_asn))
        .filter(|_| {
            !update
                .path_attributes
                .iter()
                .any(|attribute| matches!(attribute, PathAttribute::OnlyToCustomer(_)))
        })
}

/// One UPDATE on its way into the table.
struct Applying<'a> {
    table: &'a RibTable,
    import: &'a Import<'a>,
    held: &'a mut Held,
    update: &'a UpdateMessage,
    attributes_wire: Option<&'a [u8]>,
    applied: Applied,
    /// What the policy makes of the path, and the AS the routes originate
    /// from if it passes.
    path_verdict: Result<Option<u32>, Rejection>,
    /// Whether the ASPAs find the path invalid: asked only if a route
    /// gets as far as that mattering.
    path_invalid: OnceCell<bool>,
    /// The path the routes came with, and all their attributes, for the
    /// record of what is refused: made only if something is.
    as_path: OnceCell<Arc<[u32]>>,
    refused_attributes: OnceCell<Arc<[PathAttribute]>>,
    /// See [`only_to_customer_mark`].
    mark: Option<u32>,
    /// The UPDATE's attributes as the policy changes them, if it does:
    /// made when the first route is taken.
    changed: OnceCell<Vec<PathAttribute>>,
    /// For each of the policy's rules, whether it holds for the UPDATE's
    /// routes as far as their communities go: asked once, if there are
    /// rules.
    rules_holding: OnceCell<Vec<bool>>,
}

impl Applying<'_> {
    /// Takes out what the UPDATE withdraws.
    fn withdraw(&mut self) {
        let peer = self.import.peer;
        let unreachable =
            self.update
                .path_attributes
                .iter()
                .filter_map(|attribute| match attribute {
                    PathAttribute::MpUnreachNlri(mp) => Some(&mp.prefixes),
                    _ => None,
                });
        for routes in std::iter::once(&self.update.withdrawn_routes).chain(unreachable) {
            if !routes.is_empty() {
                self.applied.outcome += self.table.withdraw(peer, routes);
                self.held.forget(routes);
            }
        }
    }

    /// Puts in the routes of `nlri` that the policy takes. `next_hop` is
    /// theirs if they came in MP_REACH_NLRI.
    fn announce(
        &mut self,
        nlri: &[Nlri],
        next_hop: Option<&NlriNextHop>,
    ) -> Result<(), EncodeError> {
        let Some(first) = nlri.first() else {
            return Ok(());
        };
        // Whatever becomes of them, these replace what was held of the
        // same routes.
        self.held.forget(nlri);
        let family = first.prefix.address_family();
        let (origin, used_next_hop) = match self.judge_together(family, next_hop) {
            Ok(judged) => judged,
            Err(rejection) => {
                self.refuse(nlri, rejection, next_hop);
                return Ok(());
            }
        };
        // Asked once for all of them: they have the one next hop.
        let near = self
            .import
            .checks
            .next_hops
            .is_none_or(|next_hops| next_hops.next_hop_is_near(family, &used_next_hop));
        // Nearly always every route is taken, and then nothing is copied.
        let mut accepted = Cow::Borrowed(nlri);
        let mut invalid = Vec::new();
        for (at, route) in nlri.iter().enumerate() {
            if self.import.policy.originated.allows(&route.prefix) {
                if let Cow::Borrowed(_) = accepted {
                    accepted = Cow::Owned(nlri[..at].to_vec());
                }
                self.ignore(route);
                continue;
            }
            match self.judge(route, origin, near) {
                Ok(()) => {
                    if let Cow::Owned(accepted) = &mut accepted {
                        accepted.push(*route);
                    }
                }
                Err(rejection) => {
                    if let Cow::Borrowed(_) = accepted {
                        accepted = Cow::Owned(nlri[..at].to_vec());
                    }
                    if rejection.may_pass() {
                        invalid.push(*route);
                    }
                    self.refuse(std::slice::from_ref(route), rejection, next_hop);
                }
            }
        }
        if accepted.is_empty() && invalid.is_empty() {
            return Ok(());
        }
        if self.import.policy.rules.is_empty() {
            let attrs = self.attributes(next_hop, None)?;
            for route in invalid {
                self.held.hold(route, &attrs);
            }
            if !accepted.is_empty() {
                self.applied.outcome += self.table.announce(self.import.peer, &attrs, &accepted);
            }
            return Ok(());
        }
        // A rule may hold for some of the routes and not for others: they
        // go in with an attribute set for each rule that holds, and one
        // for those for which none does.
        let mut by_rule: Vec<(Option<usize>, Vec<Nlri>, Vec<Nlri>)> = Vec::new();
        let routes = accepted.iter().map(|route| (route, true));
        for (route, taken) in routes.chain(invalid.iter().map(|route| (route, false))) {
            let rule = self.rule_for(&route.prefix);
            let at = match by_rule.iter().position(|(of, ..)| *of == rule) {
                Some(at) => at,
                None => {
                    by_rule.push((rule, Vec::new(), Vec::new()));
                    by_rule.len() - 1
                }
            };
            if taken {
                by_rule[at].1.push(*route);
            } else {
                by_rule[at].2.push(*route);
            }
        }
        for (rule, taken, held) in by_rule {
            let attrs = self.attributes(next_hop, rule)?;
            for route in held {
                self.held.hold(route, &attrs);
            }
            if !taken.is_empty() {
                self.applied.outcome += self.table.announce(self.import.peer, &attrs, &taken);
            }
        }
        Ok(())
    }

    /// The first of the policy's rules that holds for the route for
    /// `prefix`, if one does.
    fn rule_for(&self, prefix: &IpPrefix) -> Option<usize> {
        let rules = &self.import.policy.rules;
        let holding = self.rules_holding.get_or_init(|| {
            // Rules look at the route as the policy keeps it: with the
            // communities the session puts on.
            let kept = self.changed();
            rules.iter().map(|rule| rule.holds_for(kept)).collect()
        });
        rules
            .iter()
            .zip(holding)
            .position(|(rule, holds)| *holds && rule.holds_for_prefix(prefix))
    }

    /// The UPDATE's attributes as its routes are kept, but for what a
    /// rule does: as the policy changes them, without the prefixes, and
    /// marked "only to customer" if they are to be.
    fn changed(&self) -> &Vec<PathAttribute> {
        self.changed.get_or_init(|| {
            let mut attributes = self.update.path_attributes.clone();
            // Not the route's own: they hold the prefixes.
            attributes.retain(|attribute| {
                !matches!(
                    attribute,
                    PathAttribute::MpReachNlri(_) | PathAttribute::MpUnreachNlri(_)
                )
            });
            self.import.policy.change(&mut attributes);
            if let Some(marked_by) = self.mark {
                attributes.push(PathAttribute::OnlyToCustomer(marked_by));
            }
            attributes
        })
    }

    /// Judges what routes announced together have in common: the path,
    /// and the next hop. `routes` is their IP version. Returns the AS they
    /// originate from, if the path says, and their next hop.
    fn judge_together(
        &self,
        routes: IpVersion,
        next_hop: Option<&NlriNextHop>,
    ) -> Result<(Option<u32>, NlriNextHop), Rejection> {
        let origin = self.path_verdict?;
        let in_attributes = || {
            self.update
                .path_attributes
                .iter()
                .find_map(|attribute| match attribute {
                    PathAttribute::NextHop(next_hop) => Some(NlriNextHop::Ipv4(next_hop.0)),
                    _ => None,
                })
        };
        // The decoder lets no route through without one.
        let next_hop = next_hop
            .copied()
            .or_else(in_attributes)
            .ok_or(Rejection::NextHopUnusable)?;
        self.import
            .policy
            .check_next_hop(&next_hop, routes, &self.import.from)?;
        Ok((origin, next_hop))
    }

    /// Judges one route, whose origin is `origin`; `near` says whether
    /// its next hop can be reached. What is wrong with the route for good
    /// comes before what may pass.
    fn judge(&self, route: &Nlri, origin: Option<u32>, near: bool) -> Result<(), Rejection> {
        let import = self.import;
        import.policy.check_prefix(&route.prefix)?;
        let origin_invalid = import
            .checks
            .vrps
            .is_some_and(|vrps| vrps.validate(&route.prefix, origin) == Validity::Invalid);
        if origin_invalid {
            return Err(Rejection::RpkiInvalid);
        }
        if *self
            .path_invalid
            .get_or_init(|| path_is_invalid(import, self.update))
        {
            return Err(Rejection::AspaInvalid);
        }
        if import
            .checks
            .next_hops
            .is_some_and(|next_hops| next_hops.is_own_network(&route.prefix))
        {
            return Err(Rejection::OwnNetwork);
        }
        if !near {
            return Err(Rejection::NextHopUnreachable);
        }
        Ok(())
    }

    /// Takes no notice of `route`, a prefix this router originates itself,
    /// from a peer in our own AS (see `ImportPolicy::originated`): it is
    /// not taken, and not counted as refused. What the peer had announced
    /// for it before is taken out.
    fn ignore(&mut self, route: &Nlri) {
        self.applied.outcome += self
            .table
            .withdraw(self.import.peer, std::slice::from_ref(route));
    }

    /// Notes that `refused`, whose next hop is `next_hop` if they came in
    /// MP_REACH_NLRI, were not taken, and takes out what the peer had
    /// announced for them before.
    fn refuse(&mut self, refused: &[Nlri], rejection: Rejection, next_hop: Option<&NlriNextHop>) {
        let came = Came {
            as_path: self.as_path.get_or_init(|| path_of(self.update)),
            attributes: self
                .refused_attributes
                .get_or_init(|| attributes_of(self.update)),
            next_hop,
        };
        self.import.rejections.note(refused, rejection, &came);
        self.applied.rejected += refused.len();
        self.applied.outcome += self.table.withdraw(self.import.peer, refused);
    }

    /// The attribute set the UPDATE's routes go into the table with.
    fn attributes(
        &self,
        next_hop: Option<&NlriNextHop>,
        rule: Option<usize>,
    ) -> Result<Arc<AttrSet>, EncodeError> {
        let store = self.table.attrs();
        if let Some(rule) = rule {
            let mut attributes = self.changed().clone();
            self.import.policy.rules[rule].apply(&mut attributes);
            return store.intern(&attributes, next_hop);
        }
        // What the policy puts on the peer's routes (communities that say
        // where they were learned, a preference) is the same for all of
        // an UPDATE's routes, and worked out once for them.
        if self.import.policy.changes(&self.update.path_attributes) {
            return store.intern(self.changed(), next_hop);
        }
        let Some(marked_by) = self.mark else {
            if let Some(set) = self
                .attributes_wire
                .and_then(|wire| store.intern_received(wire, next_hop))
            {
                return Ok(set);
            }
            return store.intern(&self.update.path_attributes, next_hop);
        };
        let added = only_to_customer_wire(marked_by);
        if let Some(set) = self
            .attributes_wire
            .and_then(|wire| store.intern_received_with(wire, &added, next_hop))
        {
            return Ok(set);
        }
        let mut marked = self.update.path_attributes.clone();
        marked.push(PathAttribute::OnlyToCustomer(marked_by));
        store.intern(&marked, next_hop)
    }
}

/// Whether the ASPAs find the path of `update`'s routes invalid.
fn path_is_invalid(import: &Import<'_>, update: &UpdateMessage) -> bool {
    let Some((aspas, check)) = import.checks.aspas else {
        return false;
    };
    // Without ASPAs nothing is known about any path, and there are none
    // until a validator that has them is found: not worth reading the
    // path for.
    if aspas.is_empty() {
        return false;
    }
    let mut path = Vec::new();
    let has_set = policy::effective_path(&update.path_attributes, &mut path);
    check.verify(aspas, &path, has_set) == PathValidity::Invalid
}

/// The UPDATE's path attributes without its prefixes: MP_REACH_NLRI and
/// MP_UNREACH_NLRI, which can hold thousands, are left out.
fn attributes_of(update: &UpdateMessage) -> Arc<[PathAttribute]> {
    update
        .path_attributes
        .iter()
        .filter(|attribute| {
            !matches!(
                attribute,
                PathAttribute::MpReachNlri(_) | PathAttribute::MpUnreachNlri(_)
            )
        })
        .cloned()
        .collect()
}

/// The AS numbers of the UPDATE's AS_PATH.
fn path_of(update: &UpdateMessage) -> Arc<[u32]> {
    let path = update
        .path_attributes
        .iter()
        .find_map(|attribute| match attribute {
            PathAttribute::AsPath(path) => Some(path),
            _ => None,
        });
    policy::asns(path).collect()
}

/// Looks again at a peer's routes after the validation data has changed.
///
/// `suspects` are routes of the peer's in the table that may now be
/// invalid. Those that are (they may have changed since they were listed)
/// are taken out and held. Every route that was held and is no longer
/// invalid goes into the table.
pub(crate) fn revalidate(
    table: &RibTable,
    import: &Import<'_>,
    held: &mut Held,
    suspects: &[Nlri],
) -> Applied {
    let mut applied = Applied::default();
    let checks = &import.checks;
    if checks.is_none() {
        return applied;
    }
    let mut path = Vec::new();

    for suspect in suspects {
        // As it is in the table now.
        let Some(route) = table
            .routes(&suspect.prefix)
            .into_iter()
            .find(|route| route.peer == import.peer && route.path_id == suspect.path_id)
        else {
            continue;
        };
        let Some(rejection) = checks.rejection(&suspect.prefix, &route.attrs, &mut path) else {
            continue;
        };
        applied.outcome += table.withdraw(import.peer, std::slice::from_ref(suspect));
        applied.rejected += 1;
        route.attrs.as_path(&mut path);
        let (attributes, next_hop) = route.attrs.attributes().unwrap_or_default();
        import.rejections.note(
            std::slice::from_ref(suspect),
            rejection,
            &Came {
                as_path: &Arc::from(&path[..]),
                attributes: &attributes.into(),
                next_hop: next_hop.as_ref(),
            },
        );
        held.hold(*suspect, &route.attrs);
    }

    let released: Vec<Nlri> = held
        .routes
        .iter()
        .filter(|(route, attrs)| checks.rejection(&route.prefix, attrs, &mut path).is_none())
        .map(|(route, _)| *route)
        .collect();
    for route in released {
        if let Some(attrs) = held.routes.remove(&route) {
            applied.outcome += table.announce(import.peer, &attrs, std::slice::from_ref(&route));
        }
    }
    applied
}