//! Round-trip times to peers, measured from traffic we send anyway: PING → PONG and
//! HELLO → HELLO_REPLY. No probes of our own.
//!
//! A sample only counts when exactly one request was outstanding: with two, we can't tell
//! which one the answer belongs to (an answer to the older one would look too fast).
//! Samples are smoothed (exponential moving average, weight 1/4) to ride out jitter.
//!
//! What we use is the *effective* round-trip time: the smoothed value, raised to the age
//! of the oldest request still unanswered (the round trip takes at least that long), so
//! a peer that stops answering drifts away by itself. Without a recent sample (nodes
//! move; an old value says little) it is unknown. 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 {
    pings: HashMap<String, Outstanding>,
    hellos: HashMap<String, Outstanding>,
    smoothed: 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 {
    pub(crate) fn ping_sent(&mut self, peer: &str, now: Instant) {
        Self::sent(&mut self.pings, peer, now);
    }

    pub(crate) fn pong_received(&mut self, peer: &str, now: Instant) {
        if let Some(sample) = Self::answered(&mut self.pings, peer, now) {
            self.record(peer, sample, now);
        }
    }

    pub(crate) fn hello_sent(&mut self, peer: &str, now: Instant) {
        Self::sent(&mut self.hellos, peer, now);
    }

    pub(crate) fn hello_reply_received(&mut self, peer: &str, now: Instant) {
        if let Some(sample) = Self::answered(&mut self.hellos, peer, now) {
            self.record(peer, sample, now);
        }
    }

    /// The effective round-trip time: the smoothed value if sampled within `fresh_for`,
    /// raised to the age of the oldest unanswered request; None without a recent sample.
    pub(crate) fn effective(
        &self,
        peer: &str,
        now: Instant,
        fresh_for: Duration,
    ) -> Option<Duration> {
        let sample = self
            .smoothed
            .get(peer)
            .filter(|sample| now.saturating_duration_since(sample.at) < fresh_for)?;
        let waiting = [self.pings.get(peer), self.hellos.get(peer)]
            .into_iter()
            .flatten()
            .map(|outstanding| now.saturating_duration_since(outstanding.sent))
            .max()
            .unwrap_or_default();
        Some(sample.rtt.max(waiting))
    }

    /// The effective 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)> {
        self.smoothed.keys().filter_map(move |peer| {
            self.effective(peer, now, fresh_for)
                .map(|rtt| (peer.as_str(), rtt))
        })
    }

    fn sent(map: &mut HashMap<String, Outstanding>, peer: &str, now: Instant) {
        map.entry(peer.to_owned())
            .and_modify(|outstanding| outstanding.count += 1)
            .or_insert(Outstanding {
                sent: now,
                count: 1,
            });
    }

    fn answered(
        map: &mut HashMap<String, Outstanding>,
        peer: &str,
        now: Instant,
    ) -> Option<Duration> {
        let outstanding = map.remove(peer)?;
        (outstanding.count == 1).then(|| now.saturating_duration_since(outstanding.sent))
    }

    fn record(&mut self, peer: &str, sample: Duration, now: Instant) {
        let rtt = match self.smoothed.get(peer) {
            Some(old) => (old.rtt * 3 + sample) / 4,
            None => sample,
        };
        self.smoothed
            .insert(peer.to_owned(), Sample { rtt, at: now });
    }
}

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

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

    #[test]
    fn samples_only_unambiguous_answers_and_smooths() {
        let mut table = RttTable::default();
        let start = Instant::now();
        let at = |ms| start + Duration::from_millis(ms);
        table.ping_sent("a", start);
        table.pong_received("a", at(100));
        assert_eq!(
            table.effective("a", at(100), FRESH),
            Some(Duration::from_millis(100))
        );

        table.ping_sent("a", at(100));
        table.pong_received("a", at(300));
        assert_eq!(
            table.effective("a", at(300), FRESH),
            Some(Duration::from_millis(125))
        );

        // Two PINGs out: the PONG can't be attributed.
        table.ping_sent("a", at(300));
        table.ping_sent("a", at(400));
        table.pong_received("a", at(450));
        assert_eq!(
            table.effective("a", at(450), FRESH),
            Some(Duration::from_millis(125))
        );

        // An unsolicited PONG changes nothing.
        table.pong_received("b", start);
        assert_eq!(table.effective("b", start, FRESH), None);
    }

    #[test]
    fn unanswered_requests_raise_it_and_old_samples_expire() {
        let mut table = RttTable::default();
        let start = Instant::now();
        table.ping_sent("a", start);
        table.pong_received("a", start + Duration::from_millis(50));

        // A PING unanswered for 40 s: at least 40 s.
        let sent = start + Duration::from_secs(60);
        table.ping_sent("a", sent);
        let later = sent + Duration::from_secs(40);
        assert_eq!(
            table.effective("a", later, FRESH),
            Some(Duration::from_secs(40))
        );

        // No sample within the freshness window: unknown.
        let stale = start + FRESH + Duration::from_secs(1);
        assert_eq!(table.effective("a", stale, FRESH), None);
    }
}