use super::{QuotaPolicy, QuotaPool};
use crate::refill_rate::RefillRate;
use std::sync::{Arc, Barrier};

#[test]
fn pool_consumes_initial_tokens_for_one_key() {
    let policy = QuotaPolicy::new()
        .set_capacity(10.0)
        .set_refill_rate(RefillRate::per_sec(3));

    let pool = QuotaPool::new(policy, 10);
    let mut results = vec![];

    for _ in 0..100 {
        results.push(pool.consume("testing", 1));
    }

    assert_eq!(results.iter().filter(|r| r.is_ok()).count(), 10);
    assert_eq!(results.iter().filter(|r| r.is_err()).count(), 90);
}

#[test]
fn pool_limits_one_key_across_threads() {
    let policy = QuotaPolicy::new().set_capacity(10.0);
    let pool = Arc::new(QuotaPool::new(policy, 10));
    let barrier = Arc::new(Barrier::new(100));
    let mut handles = Vec::new();

    for _ in 0..100 {
        let pool = Arc::clone(&pool);
        let barrier = Arc::clone(&barrier);
        handles.push(std::thread::spawn(move || {
            barrier.wait();
            pool.consume("testing", 1).is_ok()
        }));
    }

    let accepted = handles
        .into_iter()
        .map(|handle| handle.join().expect("worker thread panicked"))
        .filter(|accepted| *accepted)
        .count();

    assert_eq!(accepted, 10);
    assert_eq!(pool.len(), 1);
}

#[test]
fn pool_consumes_owned_key_and_keeps_entry() {
    let policy = QuotaPolicy::new().set_capacity(1.0);
    let pool = QuotaPool::new(policy, 1);

    assert!(pool.consume_owned(String::from("testing"), 1).is_ok());
    assert!(pool.consume("testing", 1).is_err());
    assert!(pool.contains_key("testing"));
}

#[test]
fn pool_tracks_spread_keys_independently() {
    let policy = QuotaPolicy::new().set_capacity(1.0);
    let keys: Vec<_> = (0..64).map(|id| format!("key-{id}")).collect();
    let pool = Arc::new(QuotaPool::with_capacity(policy, 1, keys.len()));
    pool.insert_keys(&keys);

    assert_eq!(pool.len(), keys.len());

    let mut handles = Vec::new();

    for key in keys {
        let pool = Arc::clone(&pool);
        handles.push(std::thread::spawn(move || {
            pool.consume(key.as_str(), 1).is_ok()
        }));
    }

    let accepted = handles
        .into_iter()
        .map(|handle| handle.join().expect("worker thread panicked"))
        .filter(|accepted| *accepted)
        .count();

    assert_eq!(accepted, 64);
    assert_eq!(pool.len(), 64);
}
