//! Sync throughput with peers, measured from sync exchanges we have anyway: from sending
//! a sync message until the peer's answer arrives, the sync bytes that went either way
//! with that peer. No transfers of our own.
//!
//! Only exchanges of at least `MIN_BYTES` count; smaller ones measure the round trip,
//! not the throughput. The round trip (`swarm::rtt`) is taken off the elapsed time, but
//! never more than half of it, so a stale round trip can't inflate a sample. What
//! remains also includes the time the peer took to apply and answer: this is the rate a
//! peer delivers at in practice, not the raw capacity of the path.
//!
//! Samples are smoothed (exponential moving average, weight 1/4). Large exchanges are
//! rare (catching up, big writes), so estimates are kept however old; they come with
//! their age.

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

use tokio::time::Instant;

use crate::transport::traffic::Traffic;

/// Smallest exchange that counts: some dozens of Mycelium chunks.
const MIN_BYTES: u64 = 32 * 1024;
const WEIGHT: f64 = 0.25;

/// Sync bytes exchanged with a peer so far (from the traffic counters).
#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct SyncBytes {
    sent: u64,
    received: u64,
}

impl SyncBytes {
    pub(crate) fn of(traffic: &Traffic, peer: &str) -> Self {
        traffic
            .peer(peer)
            .map(|t| Self {
                sent: t.sent.sync.bytes,
                received: t.received.sync.bytes,
            })
            .unwrap_or_default()
    }

    /// Bytes either way since `earlier` (0 if the counters were dropped meanwhile).
    fn since(&self, earlier: &SyncBytes) -> u64 {
        self.sent.saturating_sub(earlier.sent) + self.received.saturating_sub(earlier.received)
    }
}

#[derive(Debug, Clone, Copy)]
pub(crate) struct Estimate {
    pub(crate) bytes_per_sec: f64,
    pub(crate) samples: u32,
    /// When the last sample was taken.
    pub(crate) at: Instant,
}

#[derive(Debug, Default)]
pub(crate) struct ThroughputTable {
    peers: HashMap<String, Estimate>,
}

impl ThroughputTable {
    /// A sync exchange with `peer` finished: it started at `started` with the counters
    /// at `before`, and its answer arrived at `arrived` with the counters at `after`.
    pub(crate) fn exchange(
        &mut self,
        peer: &str,
        (started, before): (Instant, SyncBytes),
        (arrived, after): (Instant, SyncBytes),
        rtt: Option<Duration>,
    ) {
        let bytes = after.since(&before);
        if bytes < MIN_BYTES {
            return;
        }
        let elapsed = arrived.saturating_duration_since(started);
        let transfer = elapsed
            .saturating_sub(rtt.unwrap_or_default())
            .max(elapsed / 2);
        if transfer.is_zero() {
            return;
        }
        let sample = bytes as f64 / transfer.as_secs_f64();
        let estimate = self.peers.entry(peer.to_owned()).or_insert(Estimate {
            bytes_per_sec: sample,
            samples: 0,
            at: arrived,
        });
        if estimate.samples > 0 {
            estimate.bytes_per_sec += WEIGHT * (sample - estimate.bytes_per_sec);
        }
        estimate.samples += 1;
        estimate.at = arrived;
    }

    pub(crate) fn get(&self, peer: &str) -> Option<Estimate> {
        self.peers.get(peer).copied()
    }

    pub(crate) fn retain(&mut self, mut keep: impl FnMut(&str) -> bool) {
        self.peers.retain(|peer, _| keep(peer));
    }
}

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

    fn bytes(sent: u64, received: u64) -> SyncBytes {
        SyncBytes { sent, received }
    }

    #[test]
    fn small_exchanges_dont_count() {
        let mut table = ThroughputTable::default();
        let t0 = Instant::now();
        table.exchange(
            "a",
            (t0, bytes(0, 0)),
            (t0 + Duration::from_millis(200), bytes(500, 10_000)),
            None,
        );
        assert!(table.get("a").is_none());
    }

    #[test]
    fn round_trip_is_taken_off_but_at_most_half() {
        let mut table = ThroughputTable::default();
        let t0 = Instant::now();
        // 1 MB in 2 s, of which 1 s round trip: 1 MB/s.
        table.exchange(
            "a",
            (t0, bytes(100, 0)),
            (t0 + Duration::from_secs(2), bytes(200, 999_900)),
            Some(Duration::from_secs(1)),
        );
        let a = table.get("a").unwrap();
        assert!((a.bytes_per_sec - 1_000_000.0).abs() < 1.0);

        // A round trip longer than the exchange takes off only half.
        table.exchange(
            "b",
            (t0, bytes(0, 0)),
            (t0 + Duration::from_secs(2), bytes(0, 1_000_000)),
            Some(Duration::from_secs(5)),
        );
        assert!((table.get("b").unwrap().bytes_per_sec - 1_000_000.0).abs() < 1.0);
    }

    #[test]
    fn smoothed() {
        let mut table = ThroughputTable::default();
        let t0 = Instant::now();
        for received in [500_000u64, 1_000_000] {
            table.exchange(
                "a",
                (t0, bytes(0, 0)),
                (t0 + Duration::from_secs(1), bytes(0, received)),
                None,
            );
        }
        // Samples of 500 kB/s, then 1 MB/s: 500 kB/s + (1 MB/s - 500 kB/s) / 4.
        let a = table.get("a").unwrap();
        assert_eq!(a.samples, 2);
        assert!((a.bytes_per_sec - 625_000.0).abs() < 1.0);
    }
}