#include "workers.hpp"

Workers::Workers(size_t worker_count) : worker_count(worker_count) {
  start();
}

Workers::~Workers() {
  stop();
}

void Workers::start() {
  std::unique_lock<std::mutex> lock(start_stop_mutex);
  for(size_t i = 0; i < worker_count; ++i) {
    threads.emplace_back([this] {
      while(true) {
        std::function<void()> task;
        {
          std::unique_lock<std::mutex> lock(mutex);
          while(tasks.empty() && !stop_threads)
            cv.wait(lock);
          if(tasks.empty())
            return;
          task = std::move(tasks.front());
          tasks.pop_front();
        }
        task();

        std::unique_lock<std::mutex> lock(mutex);
        if(stop_threads_on_completed_tasks && tasks.empty()) {
          stop_threads = true;
          lock.unlock();
          cv.notify_all();
          return;
        }
      }
    });
  }
}

void Workers::post(std::function<void()> &&task) {
  {
    std::unique_lock<std::mutex> lock(mutex);
    tasks.emplace_back(std::move(task));
  }
  cv.notify_one();
}

void Workers::stop() {
  std::unique_lock<std::mutex> lock(start_stop_mutex);

  if(threads.empty())
    return;

  {
    std::unique_lock<std::mutex> lock(mutex);
    if(tasks.empty()) {
      stop_threads = true;
      lock.unlock();
      cv.notify_all();
    }
    else
      stop_threads_on_completed_tasks = true;
  }

  for(auto &thread : threads)
    thread.join();
  threads.clear();

  {
    std::unique_lock<std::mutex> lock(mutex);
    stop_threads = false;
    stop_threads_on_completed_tasks = false;
  }
}

void Workers::restart() {
  stop();
  start();
}