use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

/// A cheap, cloneable, thread-safe counter.
///
/// Clones share the same underlying atomic value, so increments made through
/// any clone are visible through all others.
#[derive(Clone, Debug, Default)]
pub struct ThreadSafeCounter {
    value: Arc<AtomicU64>,
}

impl ThreadSafeCounter {
    /// Creates a counter starting at zero.
    pub fn new() -> Self {
        Self::with_value(0)
    }

    /// Creates a counter with an explicit initial value.
    pub fn with_value(initial: u64) -> Self {
        Self {
            value: Arc::new(AtomicU64::new(initial)),
        }
    }

    /// Returns the current counter value.
    pub fn get(&self) -> u64 {
        self.value.load(Ordering::SeqCst)
    }

    /// Increments the counter by one and returns the new value.
    pub fn increment(&self) -> u64 {
        self.add(1)
    }

    /// Adds `amount` to the counter and returns the new value.
    pub fn add(&self, amount: u64) -> u64 {
        self.value.fetch_add(amount, Ordering::SeqCst) + amount
    }

    /// Sets the counter to zero and returns the previous value.
    pub fn reset(&self) -> u64 {
        self.value.swap(0, Ordering::SeqCst)
    }
}

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

    #[test]
    fn basic_operations_work() {
        let counter = ThreadSafeCounter::with_value(5);

        assert_eq!(counter.get(), 5);
        assert_eq!(counter.increment(), 6);
        assert_eq!(counter.add(4), 10);
        assert_eq!(counter.reset(), 10);
        assert_eq!(counter.get(), 0);
    }

    #[test]
    fn increments_are_thread_safe() {
        let counter = ThreadSafeCounter::new();
        let threads = 8;
        let increments_per_thread = 10_000;

        let handles: Vec<_> = (0..threads)
            .map(|_| {
                let counter = counter.clone();
                thread::spawn(move || {
                    for _ in 0..increments_per_thread {
                        counter.increment();
                    }
                })
            })
            .collect();

        for handle in handles {
            handle.join().expect("worker thread panicked");
        }

        assert_eq!(counter.get(), threads * increments_per_thread);
    }
}