//! Round-trip times to peers, measured from traffic we send anyway: the keepalive PING
//! → PONG (`swarm::peer::keepalive`), and HELLO → HELLO_REPLY.
//!
//! A PONG answers the PING with its nonce, so it is attributed exactly; a PONG to an
//! older PING (one a newer one replaced) is no sample. 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). 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 a HELLO still unanswered (the round trip takes at least that long), so a peer that
//! stops answering its handshake drifts away by itself. An unanswered PING isn't: a lost
//! PING or PONG is no sample, not a slow one (the next keepalive measures again, and a
//! peer that is gone goes dormant). 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, Ping>,
    hellos: HashMap<String, Outstanding>,
    smoothed: HashMap<String, Sample>,
}

/// The PING out to a peer.
#[derive(Debug, Clone, Copy)]
struct Ping {
    nonce: u64,
    sent: Instant,
}

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

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

impl RttTable {
    /// A PING to `peer`: it replaces one still unanswered.
    pub(crate) fn ping_sent(&mut self, peer: &str, nonce: u64, now: Instant) {
        self.pings
            .insert(peer.to_owned(), Ping { nonce, sent: now });
    }

    pub(crate) fn pong_received(&mut self, peer: &str, nonce: u64, now: Instant) {
        if self.pings.get(peer).is_some_and(|ping| ping.nonce == nonce)
            && let Some(ping) = self.pings.remove(peer)
        {
            self.record(peer, now.saturating_duration_since(ping.sent), 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
            .hellos
            .get(peer)
            .map(|outstanding| now.saturating_duration_since(outstanding.sent))
            .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", 1, start);
        table.pong_received("a", 1, at(100));
        assert_eq!(
            table.effective("a", at(100), FRESH),
            Some(Duration::from_millis(100))
        );

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

        // A newer PING replaced it: a PONG to the older one is no sample.
        table.ping_sent("a", 3, at(300));
        table.ping_sent("a", 4, at(400));
        table.pong_received("a", 3, at(410));
        assert_eq!(
            table.effective("a", at(410), FRESH),
            Some(Duration::from_millis(125))
        );
        table.pong_received("a", 4, at(500));
        assert_eq!(
            table.effective("a", at(500), FRESH),
            Some(Duration::from_micros(118_750)) // (125 * 3 + 100) / 4
        );

        // Two HELLOs out: the HELLO_REPLY can't be attributed.
        table.hello_sent("c", start);
        table.hello_reply_received("c", at(100));
        table.hello_sent("c", at(100));
        table.hello_sent("c", at(200));
        table.hello_reply_received("c", at(210));
        assert_eq!(
            table.effective("c", at(210), FRESH),
            Some(Duration::from_millis(100))
        );

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

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

        // PINGs unanswered (lost): no sample, nothing raised.
        let sent = start + Duration::from_secs(60);
        table.ping_sent("a", 2, sent);
        table.ping_sent("a", 3, sent + Duration::from_secs(30));
        let later = sent + Duration::from_secs(40);
        assert_eq!(
            table.effective("a", later, FRESH),
            Some(Duration::from_millis(50))
        );

        // A HELLO unanswered for 40 s: at least 40 s.
        table.hello_sent("a", sent);
        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);
    }
}