#include "raid/core/thread_pool.hpp"
#include <cstdlib>
#include <algorithm>
#include <memory>

namespace raid {

ThreadPool::ThreadPool(unsigned n) {
  if (n == 0) {
    if (const char* e = std::getenv("RAID_THREADS")) n = static_cast<unsigned>(std::max(1, std::atoi(e)));
    else n = std::max(1u, std::thread::hardware_concurrency());
  }
  for (unsigned i = 0; i < n; ++i) workers_.emplace_back([this] { worker(); });
}

ThreadPool::~ThreadPool() {
  { std::lock_guard<std::mutex> lk(m_); stop_ = true; }
  cv_.notify_all();
  for (auto& t : workers_) t.join();
}

void ThreadPool::worker() {
  for (;;) {
    std::function<void()> job;
    { std::unique_lock<std::mutex> lk(m_); cv_.wait(lk, [this] { return stop_ || !queue_.empty(); });
      if (stop_ && queue_.empty()) return;
      job = std::move(queue_.front()); queue_.pop_front(); }
    job();
  }
}

void ThreadPool::parallel_for(usize n, const std::function<void(usize, usize)>& fn, usize grain) {
  if (n == 0) return;
  const unsigned nt = size();
  if (nt <= 1 || n == 1) { fn(0, n); return; }
  if (grain == 0) grain = std::max<usize>(1, (n + nt * 4 - 1) / (nt * 4));
  const usize chunks = (n + grain - 1) / grain;
  // Shared completion state: jobs may outlive this frame's locals by a few instructions, so the state lives on the heap
  // and is released by the last reference (job or caller).
  struct Done { std::mutex m; std::condition_variable cv; usize remaining; std::exception_ptr first_error; };
  auto done = std::make_shared<Done>(); done->remaining = chunks;
  for (usize c = 0; c < chunks; ++c) {
    const usize b = c * grain, e = std::min(n, b + grain);
    std::function<void()> job = [&fn, done, b, e] {
      std::exception_ptr err;
      try { fn(b, e); } catch (...) { err = std::current_exception(); }
      std::lock_guard<std::mutex> lk(done->m);
      if (err && !done->first_error) done->first_error = err;
      if (--done->remaining == 0) done->cv.notify_all();
    };
    { std::lock_guard<std::mutex> lk(m_); queue_.push_back(std::move(job)); }
    cv_.notify_one();
  }
  // Help from the calling thread as well, then wait.
  for (;;) {
    std::function<void()> job;
    { std::lock_guard<std::mutex> lk(m_); if (!queue_.empty()) { job = std::move(queue_.front()); queue_.pop_front(); } }
    if (!job) break;
    job();
  }
  std::exception_ptr first_error;
  { std::unique_lock<std::mutex> lk(done->m); done->cv.wait(lk, [&] { return done->remaining == 0; }); first_error = done->first_error; }
  if (first_error) std::rethrow_exception(first_error);
}

ThreadPool& ThreadPool::global() { static ThreadPool pool; return pool; }

}  // namespace raid