//! Round-trip times to peers. Mostly the paths layer's (`transport::paths`): the
//! smoothed round-trip time of the path a peer's messages go over, measured by PINGs
//! (keepalives, and the paths layer's own), taken whenever a PONG from the peer arrives
//! (`measured`). Before a peer's path is measured (while joining: no keepalives yet),
//! HELLO → HELLO_REPLY gives one, smoothed (exponential moving average, weight 1/4);
//! a HELLO_REPLY only counts when exactly one HELLO was outstanding: with two, we can't
//! tell which one the answer belongs to (an answer to the older one would look too
//! fast).
//!
//! Without a recent value (nodes move; an old value says little) it is unknown. A lost
//! PING or HELLO is no sample: a peer that stops answering goes dormant, and is given up
//! as a catch-up source, by itself. Silence alone doesn't count: we don't ask
//! non-neighbours anything, so their quiet is ours.

use std::collections::HashMap;
use std::time::Duration;

use tokio::time::Instant;

#[derive(Debug, Default)]
pub(crate) struct RttTable {
    hellos: HashMap<String, Outstanding>,
    /// From the paths layer, as last measured.
    from_paths: HashMap<String, Sample>,
    /// From HELLOs, smoothed.
    from_hellos: HashMap<String, Sample>,
}

#[derive(Debug, Clone, Copy)]
struct Sample {
    rtt: Duration,
    at: Instant,
}

#[derive(Debug, Clone, Copy)]
struct Outstanding {
    sent: Instant,
    count: u32,
}

impl RttTable {
    /// Keeps only the peers `keep` says (those still known: else every peer ever
    /// greeted or measured would stay).
    pub(crate) fn retain(&mut self, mut keep: impl FnMut(&str) -> bool) {
        self.hellos.retain(|peer, _| keep(peer));
        self.from_paths.retain(|peer, _| keep(peer));
        self.from_hellos.retain(|peer, _| keep(peer));
    }

    /// The paths layer's round-trip time to `peer` (`Transport::rtt`), as a PONG from
    /// it arrived.
    pub(crate) fn measured(&mut self, peer: &str, rtt: Duration, now: Instant) {
        self.from_paths
            .insert(peer.to_owned(), Sample { rtt, at: now });
    }

    pub(crate) fn hello_sent(&mut self, peer: &str, now: Instant) {
        self.hellos
            .entry(peer.to_owned())
            .and_modify(|outstanding| outstanding.count += 1)
            .or_insert(Outstanding {
                sent: now,
                count: 1,
            });
    }

    pub(crate) fn hello_reply_received(&mut self, peer: &str, now: Instant) {
        let Some(outstanding) = self.hellos.remove(peer) else {
            return;
        };
        if outstanding.count != 1 {
            return;
        }
        let sample = now.saturating_duration_since(outstanding.sent);
        let rtt = match self.from_hellos.get(peer) {
            Some(old) => (old.rtt * 3 + sample) / 4,
            None => sample,
        };
        self.from_hellos
            .insert(peer.to_owned(), Sample { rtt, at: now });
    }

    /// The round-trip time: the paths layer's if measured within `fresh_for`, else
    /// HELLOs' if sampled within it; None without either.
    pub(crate) fn effective(
        &self,
        peer: &str,
        now: Instant,
        fresh_for: Duration,
    ) -> Option<Duration> {
        let fresh = |sample: &&Sample| now.saturating_duration_since(sample.at) < fresh_for;
        self.from_paths
            .get(peer)
            .filter(fresh)
            .or_else(|| self.from_hellos.get(peer).filter(fresh))
            .map(|sample| sample.rtt)
    }

    /// The round-trip time of every peer that has one.
    pub(crate) fn all_effective(
        &self,
        now: Instant,
        fresh_for: Duration,
    ) -> impl Iterator<Item = (&str, Duration)> {
        let mut peers: Vec<&String> = self
            .from_paths
            .keys()
            .chain(self.from_hellos.keys())
            .collect();
        peers.sort_unstable();
        peers.dedup();
        peers.into_iter().filter_map(move |peer| {
            self.effective(peer, now, fresh_for)
                .map(|rtt| (peer.as_str(), rtt))
        })
    }
}

/// Upper bounds of the round-trip buckets in which members tell how far those they
/// reach are (`swarm::reach`), in ms; the last bucket is above them all. Doubling: a
/// writer's tree only needs to tell which of a member's possible parents is clearly the
/// faster (`swarm::tree`), and a coarse value changes seldom, and with it what is told.
const BUCKET_BOUNDS_MS: [u64; 6] = [25, 50, 100, 200, 400, 800];

/// The bucket of a link without a round trip measured: in the middle, neither fast
/// nor slow.
pub(crate) const UNKNOWN_BUCKET: u8 = 3;

/// How far past a bound a round trip goes before its bucket changes: a round trip close
/// to a bound would otherwise move between two buckets, and every move is told to every
/// neighbour and may change who passes changes to whom.
const BUCKET_HYSTERESIS: f64 = 0.2;

/// The round-trip bucket of a link (`BUCKET_BOUNDS_MS`): `rtt`'s, unless it is still
/// within `BUCKET_HYSTERESIS` of the bucket `before` (as last told). Without a round
/// trip, as before; never measured, `UNKNOWN_BUCKET`.
pub(crate) fn bucket(rtt: Option<Duration>, before: Option<u8>) -> u8 {
    let Some(rtt) = rtt else {
        return before.unwrap_or(UNKNOWN_BUCKET);
    };
    let ms = rtt.as_secs_f64() * 1000.0;
    let plain = BUCKET_BOUNDS_MS
        .iter()
        .take_while(|bound| ms >= **bound as f64)
        .count() as u8;
    let Some(before) = before.filter(|before| *before != plain) else {
        return plain;
    };
    let index = usize::from(before);
    let lower = index
        .checked_sub(1)
        .and_then(|below| BUCKET_BOUNDS_MS.get(below))
        .map_or(0.0, |bound| *bound as f64 * (1.0 - BUCKET_HYSTERESIS));
    let upper = BUCKET_BOUNDS_MS.get(index).map_or(f64::INFINITY, |bound| {
        *bound as f64 * (1.0 + BUCKET_HYSTERESIS)
    });
    if (lower..upper).contains(&ms) {
        before
    } else {
        plain
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const FRESH: Duration = Duration::from_secs(300);

    #[test]
    fn hellos_sample_only_unambiguous_answers_and_smooth() {
        let mut table = RttTable::default();
        let start = Instant::now();
        let at = |ms| start + Duration::from_millis(ms);
        table.hello_sent("a", start);
        table.hello_reply_received("a", at(100));
        assert_eq!(
            table.effective("a", at(100), FRESH),
            Some(Duration::from_millis(100))
        );
        table.hello_sent("a", at(100));
        table.hello_reply_received("a", at(300));
        assert_eq!(
            table.effective("a", at(300), FRESH),
            Some(Duration::from_millis(125))
        );
        // Two HELLOs out: the HELLO_REPLY can't be attributed.
        table.hello_sent("a", at(300));
        table.hello_sent("a", at(400));
        table.hello_reply_received("a", at(410));
        assert_eq!(
            table.effective("a", at(410), FRESH),
            Some(Duration::from_millis(125))
        );
        // An unsolicited HELLO_REPLY changes nothing.
        table.hello_reply_received("b", start);
        assert_eq!(table.effective("b", start, FRESH), None);
    }

    #[test]
    fn the_path_wins_nothing_waits_and_old_values_expire() {
        let mut table = RttTable::default();
        let start = Instant::now();
        table.hello_sent("a", start);
        table.hello_reply_received("a", start + Duration::from_millis(80));
        table.measured("a", Duration::from_millis(20), start);
        assert_eq!(
            table.effective("a", start, FRESH),
            Some(Duration::from_millis(20))
        );
        // A HELLO unanswered for long raises nothing.
        table.hello_sent("a", start);
        assert_eq!(
            table.effective("a", start + Duration::from_secs(40), FRESH),
            Some(Duration::from_millis(20))
        );
        assert_eq!(
            table.all_effective(start, FRESH).collect::<Vec<_>>(),
            [("a", Duration::from_millis(20))]
        );
        // No value within the freshness window: unknown.
        let stale = start + FRESH + Duration::from_secs(1);
        assert_eq!(table.effective("a", stale, FRESH), None);
    }

    #[test]
    fn buckets_double_and_hold_near_their_bounds() {
        let ms = |ms| Some(Duration::from_millis(ms));
        assert_eq!(bucket(ms(10), None), 0);
        assert_eq!(bucket(ms(25), None), 1);
        assert_eq!(bucket(ms(150), None), 3);
        assert_eq!(bucket(ms(5000), None), 6);
        assert_eq!(bucket(None, None), UNKNOWN_BUCKET);
        assert_eq!(bucket(None, Some(1)), 1, "unmeasured now: as before");
        // Just past a bound: still the bucket told before.
        assert_eq!(bucket(ms(110), Some(2)), 2);
        assert_eq!(bucket(ms(45), Some(2)), 2);
        // Well past it: the new one.
        assert_eq!(bucket(ms(130), Some(2)), 3);
        assert_eq!(bucket(ms(35), Some(2)), 1);
        assert_eq!(bucket(ms(2000), Some(6)), 6);
    }
}