#include "fesa/assembly/parallel_for.h"

#include <gtest/gtest.h>

#include <array>
#include <atomic>
#include <cstddef>
#include <functional>
#include <stdexcept>
#include <string>
#include <vector>

namespace fesa {
namespace {

class ParallelForBodyError final : public std::runtime_error {
 public:
  using std::runtime_error::runtime_error;
};

std::array<std::reference_wrapper<const ParallelFor>, 2> ParallelForBackends(
    const SerialParallelFor& serial, const TbbParallelFor& tbb) {
  return {std::cref(serial), std::cref(tbb)};
}

TEST(ParallelFor, ZeroOneManyExecuteExactlyOnce) {
  const SerialParallelFor serial;
  const TbbParallelFor tbb;

  for (const ParallelFor& parallel_for : ParallelForBackends(serial, tbb)) {
    std::atomic<std::size_t> zero_visits{0U};
    parallel_for.Execute(0U, [&zero_visits](std::size_t) {
      zero_visits.fetch_add(1U, std::memory_order_relaxed);
    });
    EXPECT_EQ(zero_visits.load(std::memory_order_relaxed), 0U);

    for (const std::size_t count : {1U, 257U}) {
      std::vector<std::atomic<std::size_t>> visits(count);
      for (auto& visit : visits) {
        visit.store(0U, std::memory_order_relaxed);
      }
      parallel_for.Execute(count, [&visits](std::size_t index) {
        visits[index].fetch_add(1U, std::memory_order_relaxed);
      });
      for (std::size_t index = 0; index < count; ++index) {
        EXPECT_EQ(visits[index].load(std::memory_order_relaxed), 1U);
      }
    }
  }
}

TEST(ParallelFor, SerialAndTbbProduceStableIndexedOutput) {
  constexpr std::size_t count = 1024U;
  std::vector<std::atomic<std::size_t>> serial_output(count);
  std::vector<std::atomic<std::size_t>> tbb_output(count);
  std::vector<std::atomic<std::size_t>> serial_visits(count);
  std::vector<std::atomic<std::size_t>> tbb_visits(count);
  for (std::size_t index = 0; index < count; ++index) {
    serial_output[index].store(0U, std::memory_order_relaxed);
    tbb_output[index].store(0U, std::memory_order_relaxed);
    serial_visits[index].store(0U, std::memory_order_relaxed);
    tbb_visits[index].store(0U, std::memory_order_relaxed);
  }
  const auto value_for_index = [](std::size_t index) {
    return (index + 17U) * (index + 3U);
  };

  const SerialParallelFor serial;
  serial.Execute(count, [&serial_output, &serial_visits,
                         &value_for_index](std::size_t index) {
    serial_output[index].store(value_for_index(index),
                               std::memory_order_relaxed);
    serial_visits[index].fetch_add(1U, std::memory_order_relaxed);
  });

  const TbbParallelFor tbb;
  tbb.Execute(count, [&tbb_output, &tbb_visits,
                      &value_for_index](std::size_t index) {
    tbb_output[index].store(value_for_index(index), std::memory_order_relaxed);
    tbb_visits[index].fetch_add(1U, std::memory_order_relaxed);
  });

  for (std::size_t index = 0; index < count; ++index) {
    EXPECT_EQ(serial_visits[index].load(std::memory_order_relaxed), 1U);
    EXPECT_EQ(tbb_visits[index].load(std::memory_order_relaxed), 1U);
    EXPECT_EQ(tbb_output[index].load(std::memory_order_relaxed),
              serial_output[index].load(std::memory_order_relaxed));
    EXPECT_EQ(tbb_output[index].load(std::memory_order_relaxed),
              value_for_index(index));
  }
}

TEST(ParallelFor, PropagatesBodyExceptionByContract) {
  const SerialParallelFor serial;
  const TbbParallelFor tbb;

  for (const ParallelFor& parallel_for : ParallelForBackends(serial, tbb)) {
    try {
      // Every iteration throws the same value so the assertion is independent
      // of which oneTBB task reports the cancellation-triggering exception.
      parallel_for.Execute(64U, [](std::size_t) {
        throw ParallelForBodyError{"parallel-for-body-failure"};
      });
      ADD_FAILURE() << "ParallelFor swallowed the body exception.";
    } catch (const ParallelForBodyError& error) {
      EXPECT_EQ(std::string{error.what()}, "parallel-for-body-failure");
    } catch (...) {
      ADD_FAILURE() << "ParallelFor changed the body exception type.";
    }
  }
}

}  // namespace
}  // namespace fesa