#include "prism_engineering/worker_pool.hpp"
#include <algorithm>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <thread>

namespace {
std::size_t checks{};
void require(bool value, const char *description) {
  ++checks;
  if (!value)
    throw std::runtime_error{description};
}
void contracts() {
  using namespace prism_engineering;
  auto report =
      run_batch({{4, {1, 2, 3}, false}, {4, {10, 20}, false}}, {2, {}});
  require(report.has_value(), "normal batch completes");
  require(report->results.size() == 2 && report->results[0].checksum == 6 &&
              report->results[1].checksum == 30,
          "result slots preserve submission order and owned data");
  require(report->peak_active >= 1 && report->peak_active <= 2,
          "active jobs never exceed worker bound");
  require(accepts(report->results[0], 4),
          "current successful generation can publish");
  require(!accepts(report->results[0], 5),
          "stale successful generation cannot publish");
  require(!accepts({4, Status::failed, 0}, 4),
          "failed current generation cannot publish");
  require(!run_batch({}, {0, {}}), "zero workers rejected before launch");
  require(!run_batch({}, {max_workers + 1, {}}),
          "worker limit enforced before launch");
  auto empty = run_batch({}, {1, {}});
  require(empty && empty->results.empty() && empty->peak_active == 0,
          "empty batch has no active jobs");
  std::vector<Job> maximum(max_jobs, Job{9, {2}, false});
  auto maximum_report = run_batch(maximum, {max_workers, {}});
  require(maximum_report && maximum_report->results.size() == max_jobs &&
              std::ranges::all_of(maximum_report->results,
                                  [](const Result &r) {
                                    return r.status == Status::succeeded &&
                                           r.checksum == 2;
                                  }),
          "exact job and worker bounds are supported");
  maximum.push_back({9, {2}, false});
  require(!run_batch(std::move(maximum)), "one job beyond limit is rejected");
  const auto largest = std::numeric_limits<std::uint32_t>::max();
  auto maximum_payload = run_batch(
      {{2, std::vector<std::uint32_t>(max_samples_per_job, largest), false}});
  require(maximum_payload &&
              maximum_payload->results[0].checksum ==
                  static_cast<std::uint64_t>(largest) * max_samples_per_job,
          "bounded maximum integer sum fits the proven uint64 domain");
  require(!run_batch({{1, std::vector<std::uint32_t>(max_samples_per_job + 1),
                       false}}),
          "one sample beyond limit rejected");
  auto failure = run_batch({{3, {5}, true}, {3, {8}, false}}, {1, {}});
  require(failure && failure->results[0].status == Status::failed &&
              failure->results[0].checksum == 0,
          "worker exception becomes a failed result instead of escaping thread "
          "entry");
  require(failure->results[1].status == Status::succeeded &&
              failure->results[1].checksum == 8,
          "worker can finish independent work after a failed job");
  std::stop_source already_cancelled;
  already_cancelled.request_stop();
  auto pre =
      run_batch({{3, {5}, false}}, {2, {}}, already_cancelled.get_token());
  require(pre && pre->results[0].status == Status::cancelled &&
              pre->peak_active == 0,
          "pre-cancelled batch performs no computation");

  auto no_worker = run_batch({{3, {5}, false}}, {2, {}, 0});
  require(!no_worker && no_worker.error() == Error::thread_start,
          "failure before first worker reports the launch error");
  auto failed_gate = std::make_shared<WorkGate>();
  auto partial_start = run_batch({{3, {5}, false}}, {2, failed_gate, 1});
  require(!partial_start && partial_start.error() == Error::thread_start,
          "partial launch failure stops and joins the first gated worker");
  require(failed_gate->wait_for_arrivals(1, already_cancelled.get_token()),
          "the launched worker reached its closed gate before joined return");
  auto after_failure = run_batch({{3, {8}, false}}, {1, {}});
  require(after_failure && after_failure->results[0].status == Status::succeeded &&
              after_failure->results[0].checksum == 8,
          "launch failure leaves a subsequent independent batch usable");

  // No sleep or scheduler assumption: all workers report arrival at a closed,
  // stop-aware gate before the controller requests cancellation.
  auto gate = std::make_shared<WorkGate>();
  std::stop_source cancellation;
  std::jthread controller{[gate, cancellation](std::stop_token stop) mutable {
    if (gate->wait_for_arrivals(3, stop))
      cancellation.request_stop();
  }};
  auto cancelled = run_batch(std::vector<Job>(8, Job{10, {1, 2, 3}, false}),
                             {3, gate}, cancellation.get_token());
  controller.request_stop();
  controller.join();
  require(cancelled.has_value(), "gated cancellation returns a report");
  require(cancelled->peak_active == 0,
          "gate cancellation occurs before queued payload access");
  require(std::ranges::all_of(cancelled->results,
                              [](const Result &r) {
                                return r.status == Status::cancelled &&
                                       r.checksum == 0;
                              }),
          "all unstarted slots become cancelled after join");
  require(!accepts(cancelled->results[0], 10),
          "cancellation never publishes a checksum");

  auto open_gate = std::make_shared<WorkGate>();
  open_gate->open();
  auto admitted = run_batch({{6, {7}, false}}, {1, open_gate});
  require(admitted && admitted->results[0].status == Status::succeeded &&
              admitted->results[0].checksum == 7,
          "open admission gate releases work");
  auto stopped_open = run_batch({{6, {7}, false}}, {1, open_gate},
                                already_cancelled.get_token());
  require(stopped_open && stopped_open->peak_active == 0 &&
              stopped_open->results[0].status == Status::cancelled,
          "an open wait predicate does not override an already requested stop");
}
} // namespace
int main() {
  try {
    contracts();
    std::cout << checks << " engineering checks passed.\n";
    return 0;
  } catch (const std::exception &e) {
    std::cerr << e.what() << '\n';
    return 1;
  }
}