//! The routing table: every route every peer has announced, and for each
//! prefix the one chosen as best.
//!
//! This is the Adj-RIBs-In and the Loc-RIB of RFC 4271 §3.2 in one
//! structure, organised by prefix: a prefix's row holds the routes for it
//! from all peers, side by side, and which of them is best. Choosing again
//! after a change looks at one row and nothing else.
//!
//! It is built for a table of hundreds of full feeds, hundreds of millions
//! of routes, taken in on many cores at once:
//!
//! - **A route is small**: sixteen bytes. A pointer to its attribute set,
//!   which is shared with every other route that has the same attributes
//!   ([`attrs`]), and one word that holds the peer it is from and the first
//!   steps of route selection as a single number to compare.
//! - **The table is in shards**, by hash of the prefix, each behind its own
//!   lock. A session applies an UPDATE itself, on its own thread, taking
//!   the locks of the few shards its prefixes fall in. There is no queue in
//!   front of the table and no thread that owns it; sessions that touch
//!   different shards do not wait for each other, and one that has to wait
//!   is slowed down, which is the back-pressure that is wanted.
//! - **Readers do not stop the table.** Reading takes a shard's lock for
//!   reading, one shard at a time, and looks at the rows in place.
//! - **Followers go at their own pace.** What acts on the best routes reads
//!   the changes from a [`journal`], and a slow reader holds nothing up.

pub mod attrs;
mod journal;
mod select;
#[cfg(test)]
mod tests;

use std::hash::BuildHasher;
use std::net::{IpAddr, Ipv4Addr};
use std::sync::atomic::{AtomicU64, AtomicUsize};
use std::sync::{Arc, Mutex, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard};

use bgp_primitives::ip_prefix::{IpPrefix, Nlri};
use foldhash::fast::RandomState;
use hashbrown::HashTable;
use hashbrown::hash_table::Entry;
use tokio::sync::Notify;

pub use attrs::{AttrSet, AttrStore, SelectionKey};
pub use journal::{Event, Progress, Reader, Routes};

/// What route selection needs to know about a peer.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PeerInfo {
    pub asn: u32,
    /// The peer's BGP Identifier.
    pub router_id: Ipv4Addr,
    pub address: IpAddr,
    /// Whether the peer is in our own AS.
    pub ibgp: bool,
}

/// A peer's number in a table, given out by [`RibTable::add_peer`]. Small,
/// because every route carries one.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct PeerIndex(u16);

impl PeerIndex {
    pub fn as_usize(self) -> usize {
        self.0.into()
    }
}

/// The LOCAL_PREF of a route that does not carry one: every route from an
/// external peer, until import policy says otherwise.
pub const DEFAULT_LOCAL_PREF: u32 = 100;

/// What a change to the table did.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Outcome {
    /// Routes that were not in the table before.
    pub added: usize,
    /// Routes that took the place of one from the same peer.
    pub replaced: usize,
    /// Routes that were taken out.
    pub removed: usize,
    /// Prefixes whose best route is now another one, or none.
    pub best_changed: usize,
}

impl std::ops::AddAssign for Outcome {
    fn add_assign(&mut self, other: Self) {
        self.added += other.added;
        self.replaced += other.replaced;
        self.removed += other.removed;
        self.best_changed += other.best_changed;
    }
}

/// One route of a prefix: sixteen bytes.
struct Path {
    /// Two things in one word. The high 48 bits are the first steps of
    /// route selection as one number, the larger the better; see
    /// [`select::preference`]. The low 16 are the peer the route is from.
    word: u64,
    attrs: Arc<AttrSet>,
}

impl Path {
    fn new(preference: u64, peer: PeerIndex, attrs: &Arc<AttrSet>) -> Self {
        Self {
            word: (preference << 16) | u64::from(peer.0),
            attrs: Arc::clone(attrs),
        }
    }

    fn preference(&self) -> u64 {
        self.word >> 16
    }

    fn peer(&self) -> PeerIndex {
        PeerIndex(self.word as u16)
    }
}

/// Everything about one prefix.
struct Row {
    prefix: IpPrefix,
    /// Index of the best route in `paths`. A row always has at least one
    /// route, and so always a best one; and the best one always has the
    /// highest [`Path::preference`] in the row, which is what lets most
    /// changes be settled without looking at the other routes
    /// ([`select`]).
    best: u32,
    /// The sequence numbers of the latest change of this row's best route
    /// and of any of its routes, by which the journals tell their entries
    /// for it that still say something from those that do not.
    seqs: journal::Seqs,
    /// The routes, in one run of memory: finding a peer's route and
    /// choosing the best are scans of it. One from each peer that has
    /// announced the prefix, or several from a peer that sends more than one
    /// (ADD-PATH, RFC 7911).
    paths: Vec<Path>,
    /// The path identifier of each route in `paths`, for telling apart the
    /// routes of a peer that sends several. Empty while every route has
    /// the identifier zero, which is all routes of peers that send one
    /// each: they pay nothing for this.
    path_ids: Vec<u32>,
}

impl Row {
    fn new(prefix: IpPrefix, path: Path, path_id: u32, seqs: journal::Seqs) -> Self {
        Self {
            prefix,
            best: 0,
            seqs,
            paths: vec![path],
            path_ids: if path_id == 0 { vec![] } else { vec![path_id] },
        }
    }

    fn path_id(&self, index: usize) -> u32 {
        self.path_ids.get(index).copied().unwrap_or(0)
    }

    /// Where the route from `peer` with the identifier `path_id` is.
    fn position(&self, peer: PeerIndex, path_id: u32) -> Option<usize> {
        self.paths
            .iter()
            .enumerate()
            .position(|(index, path)| path.peer() == peer && self.path_id(index) == path_id)
    }

    /// Where a route from `peer` is, whatever its identifier.
    fn position_of_peer(&self, peer: PeerIndex) -> Option<usize> {
        self.paths.iter().position(|path| path.peer() == peer)
    }

    fn push(&mut self, path: Path, path_id: u32) {
        if path_id != 0 && self.path_ids.is_empty() {
            self.path_ids = vec![0; self.paths.len()];
        }
        push_tight(&mut self.paths, path);
        if !self.path_ids.is_empty() {
            push_tight(&mut self.path_ids, path_id);
        }
    }

    /// Takes out the route at `at`; the last route moves into its place.
    fn swap_remove(&mut self, at: usize) -> Path {
        if !self.path_ids.is_empty() {
            self.path_ids.swap_remove(at);
        }
        self.paths.swap_remove(at)
    }

    fn best(&self) -> usize {
        self.best as usize
    }

    fn set_best(&mut self, index: usize) {
        // An index into `paths`, which cannot have that many: each is
        // sixteen bytes.
        self.best = index as u32;
    }

    fn route(&self, index: usize) -> RouteRef<'_> {
        let path = &self.paths[index];
        RouteRef {
            prefix: self.prefix,
            peer: path.peer(),
            path_id: self.path_id(index),
            attrs: &path.attrs,
            best: index == self.best(),
        }
    }

    fn routes(&self) -> impl Iterator<Item = RouteRef<'_>> {
        (0..self.paths.len()).map(|index| self.route(index))
    }
}

/// A route in the table, looked at in place.
#[derive(Debug, Clone, Copy)]
pub struct RouteRef<'a> {
    pub prefix: IpPrefix,
    /// The peer the route is from.
    pub peer: PeerIndex,
    /// The peer's identifier for the route: zero unless the peer sends
    /// several routes for a prefix (ADD-PATH, RFC 7911).
    pub path_id: u32,
    pub attrs: &'a Arc<AttrSet>,
    /// Whether it is the route chosen for its prefix.
    pub best: bool,
}

impl RouteRef<'_> {
    /// The route, to keep.
    pub fn to_owned(&self) -> Route {
        Route {
            prefix: self.prefix,
            peer: self.peer,
            path_id: self.path_id,
            attrs: Arc::clone(self.attrs),
            best: self.best,
        }
    }
}

/// A route taken out of the table to keep; see [`RouteRef`].
#[derive(Debug, Clone)]
pub struct Route {
    pub prefix: IpPrefix,
    pub peer: PeerIndex,
    pub path_id: u32,
    pub attrs: Arc<AttrSet>,
    pub best: bool,
}

struct Shard {
    rows: HashTable<Row>,
    routes: usize,
    /// A copy of the peer list, so that route selection reads it under the
    /// lock it already holds.
    peers: Vec<Option<PeerInfo>>,
    hasher: RandomState,
    /// Reused between selections.
    scratch: select::Scratch,
    /// Which prefixes have changed, for those who follow.
    journals: journal::Journals,
}

/// The routing table. Shared between threads by reference; every method
/// takes `&self`.
pub struct RibTable {
    shards: Box<[RwLock<Shard>]>,
    hasher: RandomState,
    attrs: AttrStore,
    /// Which peer numbers are in use. Changing the peer list also takes
    /// every shard's lock in turn; this serialises that.
    peers: Mutex<Vec<Option<PeerInfo>>>,
    /// Which reader slots are in use; see [`journal`].
    reader_slots: Mutex<Vec<bool>>,
    /// How many readers there are, to skip waking them when there are none.
    readers: AtomicUsize,
    /// Goes up whenever something is written to a journal that has
    /// readers.
    version: AtomicU64,
    wake: Notify,
}

/// No more peers can be added: every [`PeerIndex`] is in use.
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("the table has no room for another peer")]
pub struct TooManyPeers;

impl RibTable {
    /// A table in `shards` shards. More shards let more sessions write at
    /// once; a few times the number of cores is plenty.
    pub fn new(shards: usize) -> Self {
        Self::with_journal_limits(shards, journal::Limits::default())
    }

    /// A table whose journals are tidied, and found too long, at the given
    /// sizes. Tests use small ones to get there quickly.
    fn with_journal_limits(shards: usize, limits: journal::Limits) -> Self {
        let hasher = RandomState::default();
        Self {
            shards: (0..shards.max(1))
                .map(|_| {
                    RwLock::new(Shard {
                        rows: HashTable::new(),
                        routes: 0,
                        peers: Vec::new(),
                        hasher: hasher.clone(),
                        scratch: select::Scratch::default(),
                        journals: journal::Journals::new(limits),
                    })
                })
                .collect(),
            hasher,
            attrs: AttrStore::new(),
            peers: Mutex::default(),
            reader_slots: Mutex::default(),
            readers: AtomicUsize::new(0),
            version: AtomicU64::new(0),
            wake: Notify::new(),
        }
    }

    /// Where attribute sets for this table's routes are made.
    pub fn attrs(&self) -> &AttrStore {
        &self.attrs
    }

    /// Makes room for a peer's routes and returns the number to announce
    /// them under.
    ///
    /// The number is the peer's until [`RibTable::remove_peer`] is called
    /// with it, and may then be given to another. Whoever was given it must
    /// not use it after that: the table cannot tell a late announcement
    /// from one of the number's next owner.
    pub fn add_peer(&self, info: PeerInfo) -> Result<PeerIndex, TooManyPeers> {
        let mut peers = self.peers.lock().unwrap_or_else(PoisonError::into_inner);
        let index = match peers.iter().position(Option::is_none) {
            Some(free) => free,
            None => {
                // A peer number shares a word with other things; see `Path`.
                if peers.len() > usize::from(u16::MAX) {
                    return Err(TooManyPeers);
                }
                peers.push(None);
                peers.len() - 1
            }
        };
        peers[index] = Some(info);
        for shard in &self.shards {
            let mut shard = write(shard);
            if shard.peers.len() <= index {
                shard.peers.resize(index + 1, None);
            }
            shard.peers[index] = Some(info);
        }
        Ok(PeerIndex(index as u16))
    }

    /// Takes out every route from `peer` and gives its number back for
    /// reuse (RFC 4271 §8.2.2: the routes of a session that ended are
    /// deleted).
    ///
    /// This looks at every prefix in the table. The list of peers is not
    /// held while it does, so other peers can come and go meanwhile.
    pub fn remove_peer(&self, peer: PeerIndex) -> Outcome {
        let mut outcome = Outcome::default();
        let mut news = false;
        for shard in &self.shards {
            let mut shard = write(shard);
            let before = shard.journals.mark();
            outcome += shard.remove_peer(peer);
            // From here on the shard takes no more routes from the peer.
            if let Some(slot) = shard.peers.get_mut(peer.as_usize()) {
                *slot = None;
            }
            shard.tidy_journals();
            news |= shard.journals.has_news_since(before);
        }
        // Only now can the number be given out again: no shard knows it.
        let mut peers = self.peers.lock().unwrap_or_else(PoisonError::into_inner);
        if let Some(slot) = peers.get_mut(peer.as_usize()) {
            *slot = None;
        }
        drop(peers);
        if news {
            self.wake_readers();
        }
        outcome
    }

    /// Announces `routes` from `peer`, all with the attributes `attrs`:
    /// what one UPDATE does. A route is known by its prefix and path
    /// identifier together ([`Nlri`]; a plain [`IpPrefix`] is the route
    /// with identifier zero, which is every route of a peer that sends one
    /// for each prefix). Announcing a route the peer had announced before
    /// replaces it.
    pub fn announce<N>(&self, peer: PeerIndex, attrs: &Arc<AttrSet>, routes: &[N]) -> Outcome
    where
        N: Copy + Into<Nlri>,
    {
        let mut outcome = Outcome::default();
        self.by_shard(routes, |shard, placed| {
            let Some(info) = shard.peers.get(peer.as_usize()).copied().flatten() else {
                // The peer was removed while this was on its way.
                return;
            };
            let preference = select::preference(attrs.key(), &info);
            for (hash, nlri) in placed {
                outcome += shard.announce(*hash, *nlri, peer, preference, attrs);
            }
        });
        outcome
    }

    /// Withdraws `routes` of `peer`. Withdrawing a route that is not there
    /// does nothing.
    pub fn withdraw<N>(&self, peer: PeerIndex, routes: &[N]) -> Outcome
    where
        N: Copy + Into<Nlri>,
    {
        let mut outcome = Outcome::default();
        self.by_shard(routes, |shard, placed| {
            for (hash, nlri) in placed {
                outcome += shard.withdraw(*hash, nlri, peer);
            }
        });
        outcome
    }

    fn shard_of(&self, hash: u64) -> usize {
        spread(hash, self.shards.len())
    }

    /// Calls `apply` once for each shard that some of `routes` fall in,
    /// with that shard locked and those routes. Wakes the readers if a
    /// journal one of them follows was written to.
    fn by_shard<N>(&self, routes: &[N], mut apply: impl FnMut(&mut Shard, &[(u64, Nlri)]))
    where
        N: Copy + Into<Nlri>,
    {
        let Some(first) = routes.first() else {
            return;
        };
        let hashed = |route: &N| {
            let nlri: Nlri = (*route).into();
            (self.hasher.hash_one(nlri.prefix), nlri)
        };
        // An UPDATE has a handful of prefixes; they are sorted on the stack.
        const INLINE: usize = 32;
        let mut inline = [hashed(first); INLINE];
        let mut spilled = Vec::new();
        let placed: &mut [(u64, Nlri)] = if routes.len() <= INLINE {
            for (slot, route) in inline.iter_mut().zip(routes) {
                *slot = hashed(route);
            }
            &mut inline[..routes.len()]
        } else {
            spilled.extend(routes.iter().map(hashed));
            &mut spilled
        };
        placed.sort_unstable_by_key(|(hash, _)| self.shard_of(*hash));

        let mut news = false;
        for run in placed.chunk_by(|a, b| self.shard_of(a.0) == self.shard_of(b.0)) {
            let mut shard = write(&self.shards[self.shard_of(run[0].0)]);
            let before = shard.journals.mark();
            apply(&mut shard, run);
            shard.tidy_journals();
            news |= shard.journals.has_news_since(before);
        }
        if news {
            self.wake_readers();
        }
    }

    /// How many routes the table holds, from all peers.
    pub fn len(&self) -> usize {
        self.shards.iter().map(|shard| read(shard).routes).sum()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// How many prefixes have at least one route.
    pub fn prefix_count(&self) -> usize {
        self.shards.iter().map(|shard| read(shard).rows.len()).sum()
    }

    /// The best route for `prefix`: who it is from, and its attributes.
    pub fn best(&self, prefix: &IpPrefix) -> Option<(PeerIndex, Arc<AttrSet>)> {
        self.with_row(prefix, |row| {
            let best = &row.paths[row.best()];
            (best.peer(), Arc::clone(&best.attrs))
        })
    }

    /// Every route for `prefix`.
    pub fn routes(&self, prefix: &IpPrefix) -> Vec<Route> {
        self.with_row(prefix, |row| row.routes().map(|r| r.to_owned()).collect())
            .unwrap_or_default()
    }

    fn with_row<T>(&self, prefix: &IpPrefix, read_row: impl FnOnce(&Row) -> T) -> Option<T> {
        let hash = self.hasher.hash_one(prefix);
        let shard = read(&self.shards[self.shard_of(hash)]);
        shard
            .rows
            .find(hash, |row| row.prefix == *prefix)
            .map(read_row)
    }

    /// Calls `visit` with the best route of every prefix, in no particular
    /// order. The table is not stopped: one shard at a time is held for
    /// reading, so the routes seen are not all of one moment.
    ///
    /// `visit` is called with a shard held, and must not call into the
    /// table itself: taking the shard's lock again can wait for ever on a
    /// writer that is waiting for this hold to end.
    pub fn for_each_best(&self, mut visit: impl FnMut(&IpPrefix, PeerIndex, &Arc<AttrSet>)) {
        for shard in &self.shards {
            for row in &read(shard).rows {
                let best = &row.paths[row.best()];
                visit(&row.prefix, best.peer(), &best.attrs);
            }
        }
    }

    /// Up to `limit` routes: every route of as many prefixes as fit, in no
    /// particular order. For showing some of a table that is too large to
    /// show all of.
    pub fn some_routes(&self, limit: usize) -> Vec<Route> {
        let mut routes = Vec::new();
        for shard in &self.shards {
            for row in &read(shard).rows {
                for route in row.routes() {
                    if routes.len() == limit {
                        return routes;
                    }
                    routes.push(route.to_owned());
                }
            }
        }
        routes
    }

    /// Calls `visit` with every route in the table. As
    /// [`RibTable::for_each_best`]: the table is not stopped, and `visit`
    /// must not call into it.
    pub fn for_each_route(&self, mut visit: impl FnMut(RouteRef<'_>)) {
        for shard in &self.shards {
            for row in &read(shard).rows {
                row.routes().for_each(&mut visit);
            }
        }
    }
}

fn read(shard: &RwLock<Shard>) -> RwLockReadGuard<'_, Shard> {
    // A panic under the lock cannot leave a row half changed in a way that
    // is unsafe to read; the poison is ignored, as for the server's tables.
    shard.read().unwrap_or_else(PoisonError::into_inner)
}

fn write(shard: &RwLock<Shard>) -> RwLockWriteGuard<'_, Shard> {
    shard.write().unwrap_or_else(PoisonError::into_inner)
}

/// Picks one of `count` buckets by a hash, for spreading keys over shards
/// or lock stripes that each hold a hash table keyed by the same hash.
///
/// That table takes a key's slot from the low bits of the hash and, from
/// the top seven, the tag it compares before it looks at the key itself
/// (hashbrown). The bucket is taken from bits in between: taken from
/// either end, the keys of one bucket would all share those bits, and the
/// table inside would lose them.
fn spread(hash: u64, count: usize) -> usize {
    let middle = (hash << 7) >> 32;
    ((middle * count as u64) >> 32) as usize
}

/// Adds to a vector without the doubling that `push` does: a row of a few
/// hundred routes should not hold room for as many again. Grows by a
/// quarter, which keeps adding cheap on average.
fn push_tight<T>(vec: &mut Vec<T>, value: T) {
    if vec.len() == vec.capacity() {
        vec.reserve_exact((vec.len() / 4).max(1));
    }
    vec.push(value);
}

impl Shard {
    fn tidy_journals(&mut self) {
        self.journals.tidy(&self.rows, &self.hasher);
    }

    fn announce(
        &mut self,
        hash: u64,
        nlri: Nlri,
        peer: PeerIndex,
        preference: u64,
        attrs: &Arc<AttrSet>,
    ) -> Outcome {
        let Nlri { prefix, path_id } = nlri;
        let hasher = &self.hasher;
        let entry = self.rows.entry(
            hash,
            |row| row.prefix == prefix,
            |row| hasher.hash_one(row.prefix),
        );
        let path = Path::new(preference, peer, attrs);
        let row = match entry {
            Entry::Vacant(vacant) => {
                let seqs = self.journals.appeared(prefix);
                vacant.insert(Row::new(prefix, path, path_id, seqs));
                self.routes += 1;
                return Outcome {
                    added: 1,
                    best_changed: 1,
                    ..Outcome::default()
                };
            }
            Entry::Occupied(occupied) => occupied.into_mut(),
        };

        let (outcome, best_changed) = match row.position(peer, path_id) {
            None => {
                row.push(path, path_id);
                self.routes += 1;
                let outcome = Outcome {
                    added: 1,
                    ..Outcome::default()
                };
                (
                    outcome,
                    select::after_add(row, &self.peers, &mut self.scratch),
                )
            }
            Some(at) => {
                let before = std::mem::replace(&mut row.paths[at], path);
                let outcome = Outcome {
                    replaced: 1,
                    ..Outcome::default()
                };
                if Arc::ptr_eq(&before.attrs, attrs) && before.preference() == preference {
                    // The same route again: nothing changed at all.
                    return outcome;
                }
                let best_changed = select::after_replace(
                    row,
                    at,
                    before.preference(),
                    &self.peers,
                    &mut self.scratch,
                );
                (outcome, best_changed)
            }
        };
        self.journals.changed(row, best_changed);
        Outcome {
            best_changed: best_changed.into(),
            ..outcome
        }
    }

    fn withdraw(&mut self, hash: u64, nlri: &Nlri, peer: PeerIndex) -> Outcome {
        let prefix = nlri.prefix;
        let Ok(mut entry) = self.rows.find_entry(hash, |row| row.prefix == prefix) else {
            return Outcome::default();
        };
        let row = entry.get_mut();
        let Some(at) = row.position(peer, nlri.path_id) else {
            return Outcome::default();
        };
        self.routes -= 1;
        if row.paths.len() == 1 {
            entry.remove();
            self.journals.disappeared(prefix);
            return Outcome {
                removed: 1,
                best_changed: 1,
                ..Outcome::default()
            };
        }
        let best_changed = select::after_remove(row, at, &self.peers, &mut self.scratch);
        self.journals.changed(row, best_changed);
        Outcome {
            removed: 1,
            best_changed: best_changed.into(),
            ..Outcome::default()
        }
    }

    fn remove_peer(&mut self, peer: PeerIndex) -> Outcome {
        let mut outcome = Outcome::default();
        let peers = &self.peers;
        let scratch = &mut self.scratch;
        let journals = &mut self.journals;
        self.rows.retain(|row| {
            let mut removed = 0;
            let mut best_changed = false;
            // A peer that sends several routes for a prefix has several
            // here.
            while let Some(at) = row.position_of_peer(peer) {
                removed += 1;
                if row.paths.len() == 1 {
                    outcome.removed += removed;
                    outcome.best_changed += 1;
                    journals.disappeared(row.prefix);
                    return false;
                }
                best_changed |= select::after_remove(row, at, peers, scratch);
            }
            if removed > 0 {
                outcome.removed += removed;
                outcome.best_changed += usize::from(best_changed);
                journals.changed(row, best_changed);
            }
            true
        });
        self.routes -= outcome.removed;
        outcome
    }
}