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

use event_listener::Event;

struct Inner {
    count: AtomicUsize,
    event: Event,
}

/// A synchronization primitive that waits for a group of tasks to complete.
#[derive(Clone)]
pub struct WaitGroup {
    inner: Arc<Inner>,
}

impl WaitGroup {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Inner {
                count: AtomicUsize::new(0),
                event: Event::new(),
            }),
        }
    }

    /// Increment the counter by `n`.
    pub fn add(&self, n: usize) {
        self.inner.count.fetch_add(n, Ordering::AcqRel);
    }

    /// Decrement the counter by 1. Wakes all waiters when counter reaches 0.
    pub fn done(&self) {
        let prev = self.inner.count.fetch_sub(1, Ordering::AcqRel);
        if prev == 1 {
            self.inner.event.notify(usize::MAX);
        }
    }

    /// Wait until the counter reaches zero.
    pub async fn wait(&self) {
        loop {
            if self.inner.count.load(Ordering::Acquire) == 0 {
                return;
            }
            let listener = self.inner.event.listen();
            if self.inner.count.load(Ordering::Acquire) == 0 {
                return;
            }
            listener.await;
        }
    }

    /// Create a guard that calls `done()` when dropped.
    pub fn guard(&self) -> WaitGroupGuard {
        WaitGroupGuard { wg: self.clone() }
    }
}

impl Default for WaitGroup {
    fn default() -> Self {
        Self::new()
    }
}

/// A guard that calls `WaitGroup::done()` when dropped.
pub struct WaitGroupGuard {
    wg: WaitGroup,
}

impl Drop for WaitGroupGuard {
    fn drop(&mut self) {
        self.wg.done();
    }
}
