#include "scheduler/thread_pool.h"

namespace cf {

ThreadPool::ThreadPool(size_t threads) {
  if (threads == 0) {
    threads = 1;
  }
  for (size_t i = 0; i < threads; ++i) {
    workers_.emplace_back([this]() { worker_loop(); });
  }
}

ThreadPool::~ThreadPool() {
  {
    std::lock_guard<std::mutex> lock(mutex_);
    stop_ = true;
  }
  cv_.notify_all();
  for (auto& worker : workers_) {
    if (worker.joinable()) {
      worker.join();
    }
  }
}

void ThreadPool::enqueue(std::function<void()> fn) {
  {
    std::lock_guard<std::mutex> lock(mutex_);
    queue_.push(std::move(fn));
  }
  cv_.notify_one();
}

void ThreadPool::wait_for_idle() {
  std::unique_lock<std::mutex> lock(mutex_);
  idle_cv_.wait(lock, [this]() { return queue_.empty() && active_ == 0; });
}

void ThreadPool::worker_loop() {
  while (true) {
    std::function<void()> task;
    {
      std::unique_lock<std::mutex> lock(mutex_);
      cv_.wait(lock, [this]() { return stop_ || !queue_.empty(); });
      if (stop_ && queue_.empty()) {
        return;
      }
      task = std::move(queue_.front());
      queue_.pop();
      active_++;
    }
    task();
    {
      std::lock_guard<std::mutex> lock(mutex_);
      active_--;
      if (queue_.empty() && active_ == 0) {
        idle_cv_.notify_all();
      }
    }
  }
}

}  // namespace cf
