#include "prism_engineering/worker_pool.hpp"
#include <algorithm>
#include <atomic>
#include <limits>
#include <new>
#include <stdexcept>
#include <system_error>
#include <thread>

namespace prism_engineering {
bool WorkGate::arrive_and_wait(std::stop_token stop) {
  std::unique_lock lock{mutex_};
  ++arrivals_;
  changed_.notify_all();
  return changed_.wait(lock, stop, [this] { return open_; });
}
bool WorkGate::wait_for_arrivals(std::size_t count, std::stop_token stop) {
  std::unique_lock lock{mutex_};
  return changed_.wait(lock, stop,
                       [this, count] { return arrivals_ >= count; });
}
void WorkGate::open() {
  {
    std::lock_guard lock{mutex_};
    open_ = true;
  }
  changed_.notify_all();
}

bool accepts(const Result &result, std::uint64_t current_generation) noexcept {
  return result.status == Status::succeeded &&
         result.generation == current_generation;
}

std::expected<Report, Error> run_batch(std::vector<Job> jobs, Options options,
                                       std::stop_token cancellation) {
  if (options.workers == 0 || options.workers > max_workers)
    return std::unexpected(Error::invalid_worker_count);
  if (jobs.size() > max_jobs)
    return std::unexpected(Error::too_many_jobs);
  for (const auto &job : jobs)
    if (job.samples.size() > max_samples_per_job)
      return std::unexpected(Error::payload_too_large);
  // All jobs fit before any worker launches. Payload allocation by a caller
  // occurs before transfer; the accepted batch is at most 32 × 4096 samples.
  // Vector capacity can exceed size, so this is a work bound, not a byte cap.
  static_assert(max_samples_per_job <
                std::numeric_limits<std::uint64_t>::max() /
                    std::numeric_limits<std::uint32_t>::max());
  try {
    Report report{std::vector<Result>(jobs.size()), 0};
    for (std::size_t i = 0; i < jobs.size(); ++i)
      report.results[i].generation = jobs[i].generation;
    std::mutex state_mutex;
    std::size_t next{}, active{};
    std::stop_source stop;
    std::atomic<bool> runtime_failed{false};
    std::stop_callback external_cancel{cancellation,
                                       [&stop] { stop.request_stop(); }};
    // Declare workers last: every borrowed local above outlives all threads.
    std::vector<std::jthread> workers;
    workers.reserve(options.workers);
    try {
      for (std::size_t worker = 0; worker < options.workers; ++worker) {
        if (options.fail_launch_after == worker)
          throw std::system_error{
              std::make_error_code(std::errc::resource_unavailable_try_again),
              "injected worker launch failure"};
        workers.emplace_back([&](std::stop_token thread_stop) {
          try {
            // Individual jthread stop requests relay to the shared batch.
            std::stop_callback relay{thread_stop,
                                     [&stop] { stop.request_stop(); }};
            const auto token = stop.get_token();
            if (options.gate && !options.gate->arrive_and_wait(token))
              return;
            while (!token.stop_requested()) {
              std::size_t index;
              {
                std::lock_guard lock{state_mutex};
                if (next == jobs.size() || token.stop_requested())
                  return;
                index = next++;
                report.peak_active = std::max(report.peak_active, ++active);
              }
              Result result{jobs[index].generation, Status::succeeded, 0};
              try {
                if (jobs[index].inject_failure)
                  throw std::runtime_error{"injected job failure"};
                for (const auto sample : jobs[index].samples) {
                  if (token.stop_requested()) {
                    result.status = Status::cancelled;
                    result.checksum = 0;
                    break;
                  }
                  result.checksum += sample;
                }
              } catch (...) {
                // No exception may escape a thread entry point.
                result.status = Status::failed;
                result.checksum = 0;
              }
              {
                std::lock_guard lock{state_mutex};
                report.results[index] = result;
                --active;
              }
            }
          } catch (...) {
            // Gate/locking failures also stay inside the thread.
            runtime_failed.store(true);
            stop.request_stop();
          }
        });
      }
    } catch (const std::system_error &) {
      stop.request_stop();
      for (auto &worker : workers)
        worker.join();
      return std::unexpected(Error::thread_start);
    } catch (...) {
      stop.request_stop();
      for (auto &worker : workers)
        worker.join();
      throw;
    }
    // Explicit join, before inspecting results or destroying borrowed jobs.
    // Calling clear here would request jthread stop and cancel valid work.
    for (auto &worker : workers)
      worker.join();
    if (runtime_failed.load())
      return std::unexpected(Error::worker_runtime);
    for (auto &result : report.results)
      if (result.status == Status::pending)
        result.status = Status::cancelled;
    return report;
  } catch (const std::bad_alloc &) {
    return std::unexpected(Error::allocation);
  }
}
} // namespace prism_engineering