#include "fesa/assembly/parallel_for.hpp"

#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& parallelFor : parallelForBackends(serial, tbb)) {
        std::atomic<std::size_t> zeroVisits{0U};
        parallelFor.execute(0U, [&zeroVisits](std::size_t) {
            zeroVisits.fetch_add(1U, std::memory_order_relaxed);
        });
        EXPECT_EQ(zeroVisits.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);
            }
            parallelFor.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>> serialOutput(count);
    std::vector<std::atomic<std::size_t>> tbbOutput(count);
    std::vector<std::atomic<std::size_t>> serialVisits(count);
    std::vector<std::atomic<std::size_t>> tbbVisits(count);
    for (std::size_t index = 0; index < count; ++index) {
        serialOutput[index].store(0U, std::memory_order_relaxed);
        tbbOutput[index].store(0U, std::memory_order_relaxed);
        serialVisits[index].store(0U, std::memory_order_relaxed);
        tbbVisits[index].store(0U, std::memory_order_relaxed);
    }
    const auto valueForIndex = [](std::size_t index) {
        return (index + 17U) * (index + 3U);
    };

    const SerialParallelFor serial;
    serial.execute(count, [&serialOutput, &serialVisits, &valueForIndex](std::size_t index) {
        serialOutput[index].store(valueForIndex(index), std::memory_order_relaxed);
        serialVisits[index].fetch_add(1U, std::memory_order_relaxed);
    });

    const TbbParallelFor tbb;
    tbb.execute(count, [&tbbOutput, &tbbVisits, &valueForIndex](std::size_t index) {
        tbbOutput[index].store(valueForIndex(index), std::memory_order_relaxed);
        tbbVisits[index].fetch_add(1U, std::memory_order_relaxed);
    });

    for (std::size_t index = 0; index < count; ++index) {
        EXPECT_EQ(serialVisits[index].load(std::memory_order_relaxed), 1U);
        EXPECT_EQ(tbbVisits[index].load(std::memory_order_relaxed), 1U);
        EXPECT_EQ(
            tbbOutput[index].load(std::memory_order_relaxed),
            serialOutput[index].load(std::memory_order_relaxed));
        EXPECT_EQ(
            tbbOutput[index].load(std::memory_order_relaxed),
            valueForIndex(index));
    }
}

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

    for (const ParallelFor& parallelFor : parallelForBackends(serial, tbb)) {
        try {
            // Every iteration throws the same value so the assertion is independent
            // of which oneTBB task reports the cancellation-triggering exception.
            parallelFor.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