#pragma once

#include <mutex>
#include <twist/stdlike/atomic.hpp>
#include <wheels/logging/logging.hpp>

namespace exe::support {
class WaitGroup {
 public:
  explicit WaitGroup(uint32_t count = 0) : counter_(count) {
  }

  void Add() {
    zero_notifier_.store(0);
    counter_.fetch_add(1);
  }

  void Done() {
    if (counter_.fetch_sub(1) == 1) {
      zero_notifier_.fetch_add(1);
      zero_notifier_.notify_all();
    }
  }

  void Wait() {
    uint32_t old_zero_notifier = 0;
    while (counter_.load() != 0) {
      zero_notifier_.wait(old_zero_notifier);
    }
  }

 private:
  twist::stdlike::atomic<uint32_t> counter_{0};
  twist::stdlike::atomic<uint32_t> zero_notifier_{0};
};
}  // namespace exe::support