#include "testHarness.h"
#include "vf/JobSystem.h"

#include <atomic>
#include <condition_variable>
#include <cstring>
#include <mutex>
#include <stdexcept>
#include <thread>
#include <vector>

using namespace vf;

namespace {

// One-shot latch. Every wait in this file goes through it, so no test ever
// depends on a sleep or on how fast a worker happens to start.
class Latch {
public:
    void open() {
        {
            std::lock_guard<std::mutex> lock(mutex);
            opened = true;
        }
        cv.notify_all();
    }

    void await() {
        std::unique_lock<std::mutex> lock(mutex);
        cv.wait(lock, [this] { return opened; });
    }

private:
    std::mutex mutex;
    std::condition_variable cv;
    bool opened = false;
};

int payload(int index) {
    int value = index * 7919 + 13;
    for (int step = 0; step < 24; ++step) value = (value * 31 + step) % 1000003;
    return value;
}

std::vector<int> runPayloads(unsigned threadCount, int jobCount) {
    std::vector<int> output(static_cast<std::size_t>(jobCount), 0);
    JobSystem jobs(threadCount);
    JobBatch batch;
    for (int i = 0; i < jobCount; ++i) {
        std::size_t slot = static_cast<std::size_t>(i);
        batch.add(jobs.submit([&output, slot, i] { output[slot] = payload(i); }));
    }
    batch.wait();
    jobs.waitAll();
    return output;
}

} // namespace

VF_TEST(serialPoolRunsEveryJobOnTheCallingThread) {
    JobSystem jobs(0);
    CHECK(jobs.serialMode());
    CHECK_EQ(jobs.threadCount(), 0u);

    std::thread::id caller = std::this_thread::get_id();
    std::thread::id observed;
    int ran = 0;
    JobHandle handle = jobs.submit([&] {
        observed = std::this_thread::get_id();
        ++ran;
    });
    CHECK_EQ(ran, 1);                 // already executed at submit time
    CHECK(handle.done());
    CHECK(observed == caller);
    CHECK(!handle.failed());
}

VF_TEST(serialAndParallelPoolsProduceTheSameResults) {
    const int jobCount = 256;
    std::vector<int> serial = runPayloads(0, jobCount);
    std::vector<int> parallel = runPayloads(4, jobCount);
    CHECK_EQ(serial.size(), parallel.size());
    CHECK_EQ(std::memcmp(serial.data(), parallel.data(), serial.size() * sizeof(int)), 0);
    CHECK(serial[0] != serial[1]);    // the payload really varies per index
}

VF_TEST(handleReportsCompletionOnlyAfterTheJobBody) {
    JobSystem jobs(2);
    Latch release;
    std::atomic<bool> finished{false};
    JobHandle handle = jobs.submit([&] {
        release.await();
        finished.store(true);
    });
    CHECK(!handle.done());            // cannot finish before the latch opens
    release.open();
    handle.wait();
    CHECK(handle.done());
    CHECK(finished.load());
}

VF_TEST(batchWaitsForEveryHandle) {
    JobSystem jobs(4);
    Latch release;
    std::atomic<int> completed{0};
    JobBatch batch;
    for (int i = 0; i < 16; ++i) {
        batch.add(jobs.submit([&] {
            release.await();
            completed.fetch_add(1);
        }));
    }
    CHECK_EQ(batch.size(), static_cast<std::size_t>(16));
    release.open();
    batch.wait();
    CHECK(batch.done());
    CHECK_EQ(completed.load(), 16);
    CHECK_EQ(batch.failedCount(), static_cast<std::size_t>(0));
}

// The nodes are registered in the exact reverse of their dependency order, so a
// scheduler that ignored the edges and drained the queue front to back would be
// caught even on a single worker.
VF_TEST(graphNeverStartsNodeBeforeItsDependencies) {
    for (unsigned threads : {0u, 1u, 4u}) {
        JobSystem jobs(threads);
        std::atomic<int> ticket{0};
        const std::size_t nodeCount = 8;
        std::vector<int> entered(nodeCount, -1);
        std::vector<int> exited(nodeCount, -1);

        JobGraph graph;
        for (std::size_t i = 0; i < nodeCount; ++i) {
            graph.addNode("node" + std::to_string(i), [&, i] {
                entered[i] = ticket.fetch_add(1);
                exited[i] = ticket.fetch_add(1);
            });
        }
        const std::size_t edges[][2] = {{0, 1}, {1, 2}, {2, 3}, {3, 4},
                                        {4, 5}, {5, 6}, {6, 7}, {3, 7}};
        for (const auto& edge : edges) graph.addDependency(edge[0], edge[1]);

        CHECK(!graph.hasCycle());
        JobBatch batch = jobs.runGraph(graph);
        CHECK(batch.valid());
        batch.wait();

        for (std::size_t i = 0; i < nodeCount; ++i) CHECK(entered[i] >= 0);
        for (const auto& edge : edges)
            CHECK(exited[edge[1]] < entered[edge[0]]);
        CHECK_EQ(entered[7], 0);      // the only root goes first
        CHECK_EQ(entered[0], 14);     // and the deepest node goes last
        CHECK_EQ(batch.failedCount(), static_cast<std::size_t>(0));
    }
}