//! In-memory sliding-window failure counters for brute-force throttling.
//!
//! Each key keeps the timestamps of its recent failures. A key is limited while it has `limit`
//! failures inside the window. Failures are only recorded for rejected credentials, so a locked
//! key stops growing. Memory stays bounded: expired timestamps and empty keys are pruned
//! periodically, and when the key cap is reached the key with the oldest failure is evicted.

use std::collections::{HashMap, VecDeque};
use std::net::IpAddr;
use std::sync::{Mutex, MutexGuard};
use std::time::{Duration, Instant};

/// Failure window shared by login and git authentication.
pub const WINDOW: Duration = Duration::from_secs(15 * 60);
/// Failed logins allowed per folded username inside [`WINDOW`].
pub const LOGIN_USER_LIMIT: usize = 10;
/// Failed logins allowed per client IP inside [`WINDOW`].
pub const LOGIN_IP_LIMIT: usize = 30;
/// Failed git Basic/token authentications allowed per client IP inside [`WINDOW`].
pub const GIT_IP_LIMIT: usize = 30;
/// Failed second-factor checks (login and settings) allowed per user inside [`WINDOW`].
pub const TWO_FACTOR_USER_LIMIT: usize = 10;
/// Maximum number of distinct keys tracked at once.
const MAX_KEYS: usize = 100_000;
/// Run a full prune every this many recorded failures.
const PRUNE_EVERY: u32 = 256;

pub struct Throttle {
    window: Duration,
    max_keys: usize,
    state: Mutex<State>,
}

#[derive(Default)]
struct State {
    failures: HashMap<String, VecDeque<Instant>>,
    recorded: u32,
}

impl Default for Throttle {
    fn default() -> Self {
        Self::with_limits(WINDOW, MAX_KEYS)
    }
}

impl Throttle {
    pub fn with_limits(window: Duration, max_keys: usize) -> Self {
        Self {
            window,
            max_keys: max_keys.max(1),
            state: Mutex::new(State::default()),
        }
    }

    /// Returns how long until `key` drops below `limit` failures, or `None` when it is not limited.
    pub fn check(&self, key: &str, limit: usize, now: Instant) -> Option<Duration> {
        let mut state = self.lock();
        let timestamps = state.failures.get_mut(key)?;
        drop_expired(timestamps, self.window, now);
        if timestamps.len() < limit {
            return None;
        }
        // The entry at this index expires first and brings the count below `limit`.
        let index = timestamps.len().checked_sub(limit)?;
        let oldest = *timestamps.get(index)?;
        Some((oldest + self.window).saturating_duration_since(now))
    }

    /// Records one failure for `key` at `now`.
    pub fn record_failure(&self, key: &str, now: Instant) {
        let mut state = self.lock();
        state.recorded = state.recorded.wrapping_add(1);
        if state.recorded.is_multiple_of(PRUNE_EVERY) {
            self.prune_locked(&mut state, now);
        }
        if !state.failures.contains_key(key) && state.failures.len() >= self.max_keys {
            self.evict_oldest_locked(&mut state);
        }
        let timestamps = state.failures.entry(key.to_owned()).or_default();
        drop_expired(timestamps, self.window, now);
        timestamps.push_back(now);
    }

    /// Forgets all failures for `key`, for example after a successful login.
    pub fn clear(&self, key: &str) {
        self.lock().failures.remove(key);
    }

    /// Drops expired failures and empty keys.
    #[cfg(test)]
    pub fn prune(&self, now: Instant) {
        let mut state = self.lock();
        self.prune_locked(&mut state, now);
    }

    /// Number of keys currently tracked.
    #[cfg(test)]
    pub fn tracked_keys(&self) -> usize {
        self.lock().failures.len()
    }

    fn prune_locked(&self, state: &mut State, now: Instant) {
        let window = self.window;
        state.failures.retain(|_, timestamps| {
            drop_expired(timestamps, window, now);
            !timestamps.is_empty()
        });
    }

    /// Removes the key whose most recent failure is the oldest. Runs only when the cap is hit.
    fn evict_oldest_locked(&self, state: &mut State) {
        let victim = state
            .failures
            .iter()
            .filter_map(|(key, timestamps)| timestamps.back().map(|last| (key, *last)))
            .min_by_key(|(_, last)| *last)
            .map(|(key, _)| key.clone());
        if let Some(key) = victim {
            state.failures.remove(&key);
        }
    }

    fn lock(&self) -> MutexGuard<'_, State> {
        // A poisoned lock still holds consistent counters, so keep using it.
        self.state.lock().unwrap_or_else(|e| e.into_inner())
    }
}

fn drop_expired(timestamps: &mut VecDeque<Instant>, window: Duration, now: Instant) {
    while let Some(front) = timestamps.front() {
        if now.saturating_duration_since(*front) < window {
            break;
        }
        timestamps.pop_front();
    }
}

/// Throttle key for a folded username.
pub fn user_key(folded: &str) -> String {
    format!("user:{folded}")
}

/// Throttle key for second-factor checks of one user, by user id.
pub fn two_factor_key(user_id: &str) -> String {
    format!("2fa:{user_id}")
}

/// Throttle key for an IP under a named scope, such as `login-ip` or `git-ip`.
pub fn ip_key(scope: &str, ip: IpAddr) -> String {
    format!("{scope}:{ip}")
}

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

    fn base() -> Instant {
        Instant::now()
    }

    #[test]
    fn limits_after_threshold_and_reports_retry_after() {
        let throttle = Throttle::with_limits(Duration::from_secs(60), 16);
        let t0 = base();
        for i in 0..3 {
            assert!(
                throttle
                    .check("k", 3, t0 + Duration::from_secs(i))
                    .is_none()
            );
            throttle.record_failure("k", t0 + Duration::from_secs(i));
        }
        let retry = throttle
            .check("k", 3, t0 + Duration::from_secs(10))
            .expect("limited");
        assert_eq!(retry, Duration::from_secs(50));
        // Once the first failure leaves the window, the key is free again.
        assert!(
            throttle
                .check("k", 3, t0 + Duration::from_secs(60))
                .is_none()
        );
    }

    #[test]
    fn clear_resets_a_key() {
        let throttle = Throttle::with_limits(Duration::from_secs(60), 16);
        let t0 = base();
        throttle.record_failure("k", t0);
        throttle.clear("k");
        assert!(throttle.check("k", 1, t0).is_none());
        assert_eq!(throttle.tracked_keys(), 0);
    }

    #[test]
    fn pruning_keeps_the_map_bounded() {
        let throttle = Throttle::with_limits(Duration::from_secs(10), 8);
        let t0 = base();
        for i in 0..1000u64 {
            let now = t0 + Duration::from_millis(i);
            throttle.record_failure(&format!("key-{i}"), now);
            assert!(throttle.tracked_keys() <= 8);
        }
        assert_eq!(throttle.tracked_keys(), 8);
    }

    #[test]
    fn prune_drops_expired_keys() {
        let throttle = Throttle::with_limits(Duration::from_secs(10), 8);
        let t0 = base();
        throttle.record_failure("a", t0);
        throttle.record_failure("b", t0 + Duration::from_secs(5));
        throttle.prune(t0 + Duration::from_secs(12));
        assert_eq!(throttle.tracked_keys(), 1);
        throttle.prune(t0 + Duration::from_secs(20));
        assert_eq!(throttle.tracked_keys(), 0);
    }

    #[test]
    fn keys_are_scoped_by_ip() {
        let ip: IpAddr = "192.0.2.1".parse().unwrap();
        assert_eq!(ip_key("git-ip", ip), "git-ip:192.0.2.1");
        assert_eq!(user_key("alice"), "user:alice");
    }
}