use std::collections::HashMap;
use std::time::Instant;

pub struct RateLimiter {
    counts: HashMap<String, (usize, Instant)>,
    per_second: usize,
}

impl RateLimiter {
    pub fn new(per_second: usize) -> Self {
        Self {
            counts: HashMap::new(),
            per_second,
        }
    }

    pub fn allow(&mut self, agent_id: &str) -> bool {
        let entry = self
            .counts
            .entry(agent_id.to_string())
            .or_insert((0, Instant::now()));
        if entry.1.elapsed().as_secs() >= 1 {
            *entry = (0, Instant::now());
        }
        if entry.0 < self.per_second {
            entry.0 += 1;
            true
        } else {
            false
        }
    }
}
