#include"thread_pool.h"
#include<assert.h>

// Initialize thread pool and start workers.
ThreadPool::ThreadPool(int thread_count)
{
   workers_.reserve(thread_count);
   for(int i = 0; i < thread_count; i++)
      workers_.emplace_back([this]() { WorkerThread(); });
}

// Cleanup workers.
ThreadPool::~ThreadPool()
{
   mu_.lock();
   end_workers_ = true;
   mu_.unlock();
   cv_.notify_all();
   for(auto &t : workers_)
      t.join();
   assert(work_queue_.empty());
}

// Take work off of work_queue_ and run it.
void ThreadPool::WorkerThread()
{
   bool end_this_worker = false;
   while (!end_this_worker)
   {
      std::unique_lock<std::mutex> l(mu_);
      cv_.wait(l,
               [this]() {
                  return !work_queue_.empty() || end_workers_;
               });

      // Always execute next work to be done if there is any work left.
      if( !work_queue_.empty() )
      {
         std::function<void()> f = std::move(work_queue_.front());
         work_queue_.pop();
         mu_.unlock();
         cv_.notify_all();
         f();
         continue;
      }

      // Only exit if all work is done.
      if( end_workers_ )
         end_this_worker = true;
   }
}

// Add work to thread pool.
void ThreadPool::Add(std::function<void()> &&func)
{
   std::unique_lock<std::mutex> l(mu_);
   work_queue_.emplace(func);
   cv_.notify_one();
}