//! Per-agent token-bucket rate limiter for session signal writes.
//!
//! Only `session_signal()` is rate limited. Non-session `signal()` writes
//! are NOT rate limited; their write pressure is handled by WAL backpressure.

use std::time::Instant;

use dashmap::DashMap;

// ── TokenBucket ───────────────────────────────────────────────────────────────

/// A single token bucket with lazy refill.
///
/// Tokens are added lazily at `try_acquire()` time based on the elapsed
/// wall-clock time since the last refill. No background thread or timer
/// is needed. This makes the bucket cheap to create and hold in a
/// `DashMap` shard with no per-bucket overhead when idle.
///
/// Invariant: `tokens <= capacity`. Enforced at construction and after
/// every refill.
pub(crate) struct TokenBucket {
    /// Current available tokens.
    tokens: f64,
    /// Maximum tokens (also the initial value).
    capacity: f64,
    /// Tokens added per second.
    refill_rate: f64,
    /// Wall-clock instant of the last refill computation.
    last_refill: Instant,
}

impl TokenBucket {
    /// Create a new bucket, full to capacity.
    pub(crate) fn new(capacity: f64, refill_rate: f64) -> Self {
        Self {
            tokens: capacity,
            capacity,
            refill_rate,
            last_refill: Instant::now(),
        }
    }

    /// Try to acquire one token. Returns `true` if the token was consumed,
    /// `false` if the bucket is empty.
    ///
    /// On each call, refills the bucket based on elapsed time since the
    /// last refill. This is O(1) with no syscall (`Instant::now` is vDSO
    /// on Linux, `mach_absolute_time` on macOS).
    pub(crate) fn try_acquire(&mut self) -> bool {
        self.refill();
        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            true
        } else {
            false
        }
    }

    /// Estimated milliseconds until one token is available.
    ///
    /// Returns 0 if the bucket has tokens. Used for the `retry_after_ms`
    /// field in `TidalError::RateLimited`.
    pub(crate) fn retry_after_ms(&self) -> u64 {
        if self.tokens >= 1.0 {
            return 0;
        }
        let deficit = 1.0 - self.tokens;
        let secs = deficit / self.refill_rate;
        // Ceil to the next millisecond so the caller never retries too early.
        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
        let ms = (secs * 1000.0).ceil() as u64;
        ms.max(1) // never suggest 0ms retry
    }

    /// Refill tokens based on elapsed time.
    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill);
        let added = elapsed.as_secs_f64() * self.refill_rate;
        self.tokens = (self.tokens + added).min(self.capacity);
        self.last_refill = now;
    }
}

// ── RateLimiterConfig ─────────────────────────────────────────────────────────

/// Configuration for the per-agent rate limiter.
#[derive(Debug, Clone, Copy)]
pub struct RateLimiterConfig {
    /// Maximum sustained signals per second per `(agent_id, session_id)` pair.
    ///
    /// Set to `f64::INFINITY` for no sustained-rate cap (unlimited mode).
    pub signals_per_second: f64,
    /// Burst capacity (max tokens in the bucket).
    ///
    /// Set to `f64::INFINITY` for no burst cap (unlimited mode).
    /// The default is `f64::INFINITY` (unlimited); configure explicitly to enforce limits.
    pub burst_capacity: f64,
}

impl RateLimiterConfig {
    /// Create a configuration that enforces limits.
    ///
    /// `signals_per_second` is the sustained rate; `burst_capacity` is the
    /// maximum burst (typically 2x the sustained rate).
    #[must_use]
    pub const fn limited(signals_per_second: f64, burst_capacity: f64) -> Self {
        Self {
            signals_per_second,
            burst_capacity,
        }
    }
}

impl Default for RateLimiterConfig {
    /// Unlimited by default — callers must opt-in to rate limiting via
    /// `RateLimiterConfig::limited(rate, burst)` passed to `TidalDb::builder()`.
    fn default() -> Self {
        Self {
            signals_per_second: f64::INFINITY,
            burst_capacity: f64::INFINITY,
        }
    }
}

// ── RateLimiter ───────────────────────────────────────────────────────────────

/// Per-agent, per-session token-bucket rate limiter.
///
/// Uses `DashMap` for concurrent access across threads. Each
/// `(agent_id, session_id)` pair gets its own `TokenBucket`, created
/// lazily on first access. The bucket refills at `signals_per_second`
/// tokens per second up to `burst_capacity`.
///
/// Design choice: keyed by `(agent_id, session_id)` rather than just
/// `agent_id` so that a single agent running multiple concurrent sessions
/// gets separate rate limits per session. If the agent opens 10 sessions,
/// it gets 10x the rate budget -- this is intentional because each session
/// represents a separate user interaction context.
pub struct RateLimiter {
    /// Per-(agent, session) token buckets.
    buckets: DashMap<(String, u64), TokenBucket>,
    config: RateLimiterConfig,
}

impl RateLimiter {
    /// Create a new rate limiter with the given configuration.
    #[must_use]
    pub fn new(config: RateLimiterConfig) -> Self {
        Self {
            buckets: DashMap::new(),
            config,
        }
    }

    /// Check and consume one token for the given agent-session pair.
    ///
    /// Returns `Ok(())` if the signal is allowed.
    /// Returns `Err((retry_after_ms, limit))` if rate-limited.
    ///
    /// Thread-safe: concurrent calls for different keys go to different
    /// `DashMap` shards. Concurrent calls for the same key contend on the
    /// shard lock, which serializes the token check. This is correct
    /// because the token bucket is not thread-safe on its own (`&mut self`).
    ///
    /// # Errors
    ///
    /// Returns `Err((retry_after_ms, limit))` when the bucket is empty.
    /// `retry_after_ms` is the estimated wait before one token refills;
    /// `limit` is the configured signals-per-second rate.
    pub fn check(&self, agent_id: &str, session_id: u64) -> Result<(), (u64, f64)> {
        let key = (agent_id.to_owned(), session_id);
        // Compute the result inside a block so the DashMap entry guard drops
        // before we return, releasing the shard lock as early as possible.
        let result = {
            let mut entry = self.buckets.entry(key).or_insert_with(|| {
                TokenBucket::new(self.config.burst_capacity, self.config.signals_per_second)
            });
            if entry.try_acquire() {
                Ok(())
            } else {
                Err(entry.retry_after_ms())
            }
        };
        result.map_err(|retry_ms| (retry_ms, self.config.signals_per_second))
    }

    /// Remove the bucket for a closed session.
    ///
    /// Called by `close_session()` and `close_session_internal()` to prevent
    /// unbounded growth of the bucket map. After a session is closed, its
    /// rate limit state is no longer needed.
    pub fn remove(&self, agent_id: &str, session_id: u64) {
        self.buckets.remove(&(agent_id.to_owned(), session_id));
    }

    /// Number of active buckets (for metrics/diagnostics).
    #[must_use]
    pub fn active_buckets(&self) -> usize {
        self.buckets.len()
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

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

    #[test]
    fn token_bucket_allows_burst() {
        let mut bucket = TokenBucket::new(5.0, 10.0);
        // Should allow 5 immediate acquires (burst capacity).
        for _ in 0..5 {
            assert!(bucket.try_acquire());
        }
        // 6th should fail (bucket empty, no time has passed).
        assert!(!bucket.try_acquire());
    }

    #[test]
    fn token_bucket_refills_over_time() {
        let mut bucket = TokenBucket::new(5.0, 100.0);
        // Drain all tokens.
        for _ in 0..5 {
            assert!(bucket.try_acquire());
        }
        assert!(!bucket.try_acquire());

        // Wait 50ms -> should refill ~5 tokens at 100/sec.
        std::thread::sleep(std::time::Duration::from_millis(50));
        assert!(bucket.try_acquire(), "bucket should have refilled");
    }

    #[test]
    fn token_bucket_retry_after_when_empty() {
        let mut bucket = TokenBucket::new(1.0, 10.0);
        assert!(bucket.try_acquire());
        assert!(!bucket.try_acquire());
        let retry = bucket.retry_after_ms();
        // At 10 tokens/sec, 1 token deficit = 100ms.
        assert!(retry > 0 && retry <= 100, "retry_after_ms={retry}");
    }

    #[test]
    fn rate_limiter_creates_bucket_lazily() {
        let rl = RateLimiter::new(RateLimiterConfig::default());
        assert_eq!(rl.active_buckets(), 0);
        let _ = rl.check("agent-a", 1);
        assert_eq!(rl.active_buckets(), 1);
    }

    #[test]
    fn rate_limiter_separate_buckets_per_session() {
        let config = RateLimiterConfig {
            signals_per_second: 1.0,
            burst_capacity: 1.0,
        };
        let rl = RateLimiter::new(config);

        // Session 1: consume the one token.
        assert!(rl.check("agent-a", 1).is_ok());
        assert!(rl.check("agent-a", 1).is_err());

        // Session 2: separate bucket, still has a token.
        assert!(rl.check("agent-a", 2).is_ok());
    }

    #[test]
    fn rate_limiter_remove_cleans_up() {
        let rl = RateLimiter::new(RateLimiterConfig::default());
        let _ = rl.check("agent-a", 1);
        assert_eq!(rl.active_buckets(), 1);
        rl.remove("agent-a", 1);
        assert_eq!(rl.active_buckets(), 0);
    }

    #[test]
    fn rate_limited_error_display() {
        let err = crate::schema::TidalError::RateLimited {
            agent_id: "planner".to_string(),
            limit: 100.0,
            retry_after_ms: 50,
        };
        let msg = err.to_string();
        assert!(msg.contains("planner"));
        assert!(msg.contains("100"));
        assert!(msg.contains("50"));
    }

    mod proptests {
        use proptest::prelude::*;

        use super::*;

        proptest! {
            #[test]
            fn burst_within_capacity_always_succeeds(
                capacity in 1u32..100,
                rate in 1u32..1000,
            ) {
                let config = RateLimiterConfig {
                    signals_per_second: f64::from(rate),
                    burst_capacity: f64::from(capacity),
                };
                let rl = RateLimiter::new(config);
                for _ in 0..capacity {
                    prop_assert!(rl.check("agent", 1).is_ok());
                }
            }
        }
    }
}
