//! Per-(user, room) send budget.
//!
//! Chat needs its own limiter rather than borrowing either host's. Multithreaded
//! caps posting at 15 per 60 seconds, which is a sensible forum-post rate and an
//! unusable chat rate. The per-IP `tower_governor` layers do not fit either: they
//! key on the connection, and chat's connection is long-lived and shared by every
//! message on it.
//!
//! A token bucket rather than a fixed window, because the shape of chat traffic is
//! bursty by nature. Someone answering a question types four short lines in five
//! seconds and then goes quiet for a minute; a fixed window either refuses the
//! fourth line or has to be set so loose it stops being a limit.
//!
//! # Eviction is not optional
//!
//! One bucket per (user, room) pair, held in a map, is unbounded state keyed by
//! something an attacker influences. Under a hard 512M cgroup cap that is a memory
//! exhaustion vector, not untidiness, so [`RateLimiter::evict_idle`] exists and
//! the host is expected to call it on a timer. Multithreaded's per-IP governor
//! already sweeps its own buckets every 300 seconds; this wants the same
//! treatment.

use std::time::{Duration, Instant};

use dashmap::DashMap;

use crate::ids::{RoomId, UserId};

/// Send budget.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RateLimits {
    /// Messages allowed back to back from a standing start.
    pub burst: u32,
    /// Long-run rate, in messages per minute.
    pub sustain_per_min: u32,
}

impl Default for RateLimits {
    fn default() -> Self {
        // Ten in a row, then one every two seconds. Fast enough that a person
        // talking normally never meets it, slow enough that a script is throttled
        // to something a room can absorb.
        Self {
            burst: 10,
            sustain_per_min: 30,
        }
    }
}

impl RateLimits {
    fn refill_per_sec(self) -> f64 {
        f64::from(self.sustain_per_min) / 60.0
    }
}

#[derive(Debug, Clone, Copy)]
struct Bucket {
    tokens: f64,
    last: Instant,
}

/// Outcome of a send check.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RateDecision {
    Allow,
    /// Refused. `retry_after` is when a token next becomes available, so the
    /// caller can tell the client rather than leaving it to guess.
    Deny {
        retry_after: Duration,
    },
}

impl RateDecision {
    pub fn is_allowed(&self) -> bool {
        matches!(self, Self::Allow)
    }
}

/// Token buckets keyed by (user, room).
///
/// Keyed by the pair rather than by user alone so that being talkative in one
/// room does not silence someone in another, and so a room's budget cannot be
/// consumed from outside it.
pub struct RateLimiter {
    buckets: DashMap<(UserId, RoomId), Bucket>,
    limits: RateLimits,
}

impl RateLimiter {
    pub fn new(limits: RateLimits) -> Self {
        Self {
            buckets: DashMap::new(),
            limits,
        }
    }

    /// Take a token for one message.
    ///
    /// `now` is passed in rather than read here so the behavior is testable
    /// without sleeping, and so a caller batching work can use one consistent
    /// instant.
    pub fn check(&self, user: UserId, room: RoomId, now: Instant) -> RateDecision {
        let capacity = f64::from(self.limits.burst);
        let refill = self.limits.refill_per_sec();

        let mut bucket = self.buckets.entry((user, room)).or_insert(Bucket {
            tokens: capacity,
            last: now,
        });

        // `saturating_duration_since` rather than subtraction: a caller passing a
        // `now` older than the last one (batching, or a clock read from a
        // different task) would otherwise panic.
        let elapsed = now.saturating_duration_since(bucket.last).as_secs_f64();
        bucket.tokens = (bucket.tokens + elapsed * refill).min(capacity);
        bucket.last = now;

        if bucket.tokens >= 1.0 {
            bucket.tokens -= 1.0;
            RateDecision::Allow
        } else {
            let deficit = 1.0 - bucket.tokens;
            RateDecision::Deny {
                retry_after: Duration::from_secs_f64(deficit / refill),
            }
        }
    }

    /// Drop buckets untouched for longer than `idle_for`.
    ///
    /// A bucket at full tokens carries no information: recreating it produces the
    /// same answer as keeping it. So eviction cannot make the limiter more
    /// permissive as long as `idle_for` is at least the time it takes a bucket to
    /// refill, which is what [`RateLimiter::recommended_idle_timeout`] returns.
    ///
    /// Returns how many were dropped.
    pub fn evict_idle(&self, now: Instant, idle_for: Duration) -> usize {
        let before = self.buckets.len();
        self.buckets
            .retain(|_, bucket| now.saturating_duration_since(bucket.last) < idle_for);
        before - self.buckets.len()
    }

    /// The shortest idle timeout that cannot make eviction a bypass: long enough
    /// for an empty bucket to have refilled completely.
    pub fn recommended_idle_timeout(&self) -> Duration {
        Duration::from_secs_f64(f64::from(self.limits.burst) / self.limits.refill_per_sec())
    }

    /// Live bucket count. Diagnostics, and the number to watch if eviction is
    /// ever suspected of not running.
    pub fn bucket_count(&self) -> usize {
        self.buckets.len()
    }
}

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

    fn user() -> UserId {
        UserId(Uuid::new_v4())
    }

    fn room() -> RoomId {
        RoomId(Uuid::new_v4())
    }

    fn limiter(burst: u32, sustain_per_min: u32) -> RateLimiter {
        RateLimiter::new(RateLimits {
            burst,
            sustain_per_min,
        })
    }

    #[test]
    fn burst_is_spendable_immediately_then_refused() {
        let rl = limiter(3, 60);
        let (u, r, t0) = (user(), room(), Instant::now());

        for i in 0..3 {
            assert!(rl.check(u, r, t0).is_allowed(), "burst message {i}");
        }
        assert!(!rl.check(u, r, t0).is_allowed(), "burst exhausted");
    }

    #[test]
    fn refuses_with_a_usable_retry_after() {
        // 60/min is one per second, so an empty bucket is one second from a token.
        let rl = limiter(1, 60);
        let (u, r, t0) = (user(), room(), Instant::now());

        assert!(rl.check(u, r, t0).is_allowed());
        match rl.check(u, r, t0) {
            RateDecision::Deny { retry_after } => {
                let ms = retry_after.as_millis();
                assert!((900..=1100).contains(&ms), "retry_after was {ms}ms");
            }
            RateDecision::Allow => panic!("should have been refused"),
        }
    }

    #[test]
    fn tokens_come_back_over_time() {
        let rl = limiter(2, 60);
        let (u, r, t0) = (user(), room(), Instant::now());

        assert!(rl.check(u, r, t0).is_allowed());
        assert!(rl.check(u, r, t0).is_allowed());
        assert!(!rl.check(u, r, t0).is_allowed());

        assert!(rl.check(u, r, t0 + Duration::from_secs(1)).is_allowed());
    }

    #[test]
    fn refill_does_not_exceed_the_burst_ceiling() {
        // Idling for an hour must not bank an hour of messages.
        let rl = limiter(3, 60);
        let (u, r, t0) = (user(), room(), Instant::now());

        assert!(rl.check(u, r, t0).is_allowed());
        let long_later = t0 + Duration::from_hours(1);

        for i in 0..3 {
            assert!(rl.check(u, r, long_later).is_allowed(), "message {i}");
        }
        assert!(
            !rl.check(u, r, long_later).is_allowed(),
            "banked more than the burst"
        );
    }

    #[test]
    fn the_sustained_rate_is_actually_the_sustained_rate() {
        // Spend the burst, then send at exactly the refill interval forever.
        let rl = limiter(5, 60);
        let (u, r) = (user(), room());
        let mut t = Instant::now();

        for _ in 0..5 {
            assert!(rl.check(u, r, t).is_allowed());
        }
        for i in 0..100 {
            t += Duration::from_secs(1);
            assert!(rl.check(u, r, t).is_allowed(), "sustained message {i}");
        }
        // And faster than that is still refused.
        assert!(!rl.check(u, r, t).is_allowed());
    }

    #[test]
    fn budgets_are_per_user_and_per_room() {
        let rl = limiter(1, 60);
        let (u1, u2) = (user(), user());
        let (r1, r2) = (room(), room());
        let t0 = Instant::now();

        assert!(rl.check(u1, r1, t0).is_allowed());
        assert!(!rl.check(u1, r1, t0).is_allowed(), "same pair is spent");
        assert!(rl.check(u2, r1, t0).is_allowed(), "other user unaffected");
        assert!(rl.check(u1, r2, t0).is_allowed(), "other room unaffected");
    }

    #[test]
    fn a_backwards_now_does_not_panic() {
        let rl = limiter(2, 60);
        let (u, r) = (user(), room());
        let t = Instant::now() + Duration::from_secs(10);

        assert!(rl.check(u, r, t).is_allowed());
        // Earlier than the bucket's last touch: elapsed saturates to zero rather
        // than underflowing.
        assert!(
            rl.check(u, r, t.checked_sub(Duration::from_secs(5)).unwrap())
                .is_allowed()
        );
    }

    #[test]
    fn eviction_drops_only_idle_buckets() {
        let rl = limiter(5, 60);
        let (busy, idle) = (user(), user());
        let r = room();
        let t0 = Instant::now();

        rl.check(busy, r, t0);
        rl.check(idle, r, t0);
        assert_eq!(rl.bucket_count(), 2);

        let later = t0 + Duration::from_mins(1);
        rl.check(busy, r, later);

        assert_eq!(rl.evict_idle(later, Duration::from_secs(30)), 1);
        assert_eq!(rl.bucket_count(), 1);
    }

    #[test]
    fn eviction_at_the_recommended_timeout_is_not_a_bypass() {
        // The concern: spend the burst, wait out eviction, and get a fresh bucket
        // sooner than refilling would have allowed. At the recommended timeout the
        // bucket has fully refilled anyway, so the two paths agree.
        let rl = limiter(5, 60);
        let (u, r, t0) = (user(), room(), Instant::now());

        for _ in 0..5 {
            assert!(rl.check(u, r, t0).is_allowed());
        }
        assert!(!rl.check(u, r, t0).is_allowed());

        let timeout = rl.recommended_idle_timeout();
        assert_eq!(timeout, Duration::from_secs(5), "5 tokens at 1/sec");

        let after = t0 + timeout;
        assert_eq!(rl.evict_idle(after, timeout), 1);

        // A fresh bucket gives exactly the burst, which is what waiting the same
        // period without eviction would also have given.
        for i in 0..5 {
            assert!(rl.check(u, r, after).is_allowed(), "post-eviction {i}");
        }
        assert!(!rl.check(u, r, after).is_allowed(), "no extra budget");
    }

    #[test]
    fn default_limits_are_not_the_forum_post_limit() {
        // The whole reason this module exists: 15/min would make chat unusable.
        let d = RateLimits::default();
        assert!(
            d.sustain_per_min > 15,
            "chat sustain must exceed the forum post rate"
        );
        assert!(d.burst >= 5, "chat must absorb a short burst of typing");
    }
}