//! 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.

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, Duration>,
}

#[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);
        }
    }

    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);
        }
    }

    /// The smoothed round-trip time, once measured.
    pub(crate) fn get(&self, peer: &str) -> Option<Duration> {
        self.smoothed.get(peer).copied()
    }

    /// All measured peers.
    pub(crate) fn all(&self) -> impl Iterator<Item = (&str, Duration)> {
        self.smoothed
            .iter()
            .map(|(peer, 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) {
        let smoothed = match self.smoothed.get(peer) {
            Some(old) => (*old * 3 + sample) / 4,
            None => sample,
        };
        self.smoothed.insert(peer.to_owned(), smoothed);
    }
}

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

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

        table.ping_sent("a", start);
        table.pong_received("a", start + Duration::from_millis(200));
        assert_eq!(table.get("a"), Some(Duration::from_millis(125)));

        // Two PINGs out: the PONG can't be attributed.
        table.ping_sent("a", start);
        table.ping_sent("a", start + Duration::from_secs(1));
        table.pong_received("a", start + Duration::from_secs(1));
        assert_eq!(table.get("a"), Some(Duration::from_millis(125)));

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