#include <catch2/catch_test_macros.hpp>
#include <kpn/scheduler.hpp>

#include <atomic>
#include <chrono>
#include <thread>
#include <vector>
#include <mutex>

using namespace kpn;
using namespace std::chrono_literals;

// ── basic execution ───────────────────────────────────────────────────────────

TEST_CASE("scheduler runs submitted tasks", "[scheduler]") {
    ThreadPool pool(2);
    pool.start();

    std::atomic<int> counter{0};
    for (int i = 0; i < 100; ++i)
        pool.submit([&counter]{ counter.fetch_add(1, std::memory_order_relaxed); });

    pool.drain();
    REQUIRE(counter.load() == 100);
    pool.stop();
}

TEST_CASE("scheduler single thread executes all tasks", "[scheduler]") {
    ThreadPool pool(1);
    pool.start();

    std::atomic<int> counter{0};
    for (int i = 0; i < 50; ++i)
        pool.submit([&counter]{ counter.fetch_add(1, std::memory_order_relaxed); });

    pool.drain();
    REQUIRE(counter.load() == 50);
    pool.stop();
}

// ── drain ─────────────────────────────────────────────────────────────────────

TEST_CASE("drain returns immediately when pool is idle", "[scheduler]") {
    ThreadPool pool(2);
    pool.start();
    pool.drain();   // nothing submitted — should return immediately
    pool.stop();
}

TEST_CASE("drain waits for all tasks to complete", "[scheduler]") {
    ThreadPool pool(4);
    pool.start();

    std::atomic<int> counter{0};
    constexpr int N = 200;

    for (int i = 0; i < N; ++i) {
        pool.submit([&counter]{
            std::this_thread::sleep_for(1ms);
            counter.fetch_add(1, std::memory_order_relaxed);
        });
    }

    pool.drain();
    REQUIRE(counter.load() == N);
    pool.stop();
}

TEST_CASE("drain is safe to call multiple times", "[scheduler]") {
    ThreadPool pool(2);
    pool.start();

    std::atomic<int> counter{0};
    pool.submit([&counter]{ counter.fetch_add(1, std::memory_order_relaxed); });
    pool.drain();
    REQUIRE(counter.load() == 1);

    pool.submit([&counter]{ counter.fetch_add(1, std::memory_order_relaxed); });
    pool.drain();
    REQUIRE(counter.load() == 2);

    pool.stop();
}

// ── priority ordering ─────────────────────────────────────────────────────────

TEST_CASE("higher priority tasks run before lower priority on single thread", "[scheduler]") {
    // Single thread guarantees serial execution — we can observe order.
    ThreadPool pool(1);
    pool.start();

    // Pause the worker so we can fill the queue before it drains.
    std::mutex gate;
    gate.lock();
    pool.submit([&gate]{ std::lock_guard lg(gate); }); // blocks worker

    std::vector<float> order;
    std::mutex order_mx;

    for (float p : {0.1f, 0.9f, 0.5f, 0.8f, 0.2f}) {
        pool.submit([p, &order, &order_mx]{
            std::lock_guard lg(order_mx);
            order.push_back(p);
        }, p);
    }

    gate.unlock();   // release the blocking task
    pool.drain();
    pool.stop();

    // order should be descending by priority
    REQUIRE(order.size() == 5);
    for (std::size_t i = 1; i < order.size(); ++i)
        REQUIRE(order[i - 1] >= order[i]);
}

TEST_CASE("equal priority tasks execute in FIFO order on single thread", "[scheduler]") {
    ThreadPool pool(1);
    pool.start();

    std::mutex gate;
    gate.lock();
    pool.submit([&gate]{ std::lock_guard lg(gate); });

    std::vector<int> order;
    std::mutex order_mx;

    for (int i = 0; i < 5; ++i) {
        pool.submit([i, &order, &order_mx]{
            std::lock_guard lg(order_mx);
            order.push_back(i);
        }, 0.5f);   // all same priority
    }

    gate.unlock();
    pool.drain();
    pool.stop();

    REQUIRE(order == std::vector<int>{0, 1, 2, 3, 4});
}

// ── total_ / active_ accounting ───────────────────────────────────────────────

TEST_CASE("snapshot queue depth and active counts are consistent", "[scheduler]") {
    ThreadPool pool(2);
    pool.start();

    // While tasks are running, active should be > 0 and total >= active.
    std::atomic<bool> running{false};
    std::mutex gate;
    gate.lock();

    for (int i = 0; i < 4; ++i) {
        pool.submit([&gate, &running]{
            running.store(true, std::memory_order_relaxed);
            std::lock_guard lg(gate);
        });
    }

    // Spin until at least one task has started
    while (!running.load(std::memory_order_relaxed))
        std::this_thread::yield();

    auto snap = pool.snapshot("test");
    REQUIRE(snap.active_count > 0);
    REQUIRE(snap.queue_depth + snap.active_count > 0);

    gate.unlock();
    pool.drain();

    auto snap2 = pool.snapshot("test");
    REQUIRE(snap2.active_count == 0);
    REQUIRE(snap2.queue_depth  == 0);

    pool.stop();
}

TEST_CASE("submitted and completed counters are accurate", "[scheduler]") {
    ThreadPool pool(3);
    pool.start();

    constexpr int N = 60;
    for (int i = 0; i < N; ++i)
        pool.submit([]{ std::this_thread::yield(); });

    pool.drain();
    auto snap = pool.snapshot("test");
    REQUIRE(snap.tasks_submitted == static_cast<uint64_t>(N));
    REQUIRE(snap.tasks_completed == static_cast<uint64_t>(N));

    pool.stop();
}

// ── work stealing ─────────────────────────────────────────────────────────────

TEST_CASE("work stealing: all tasks complete with uneven initial distribution", "[scheduler]") {
    // 4-thread pool. Submit a burst to ensure some threads start empty and must steal.
    ThreadPool pool(4);
    pool.start();

    std::atomic<int> counter{0};
    constexpr int N = 400;

    for (int i = 0; i < N; ++i)
        pool.submit([&counter]{
            std::this_thread::sleep_for(100us);
            counter.fetch_add(1, std::memory_order_relaxed);
        });

    pool.drain();
    REQUIRE(counter.load() == N);
    pool.stop();
}

TEST_CASE("work stealing: tasks complete with more threads than initial queue targets", "[scheduler]") {
    // With round-robin, some threads may get no tasks initially and must steal.
    constexpr std::size_t THREADS = 8;
    ThreadPool pool(THREADS);
    pool.start();

    std::atomic<int> counter{0};
    // Submit fewer tasks than threads so most threads must steal
    for (int i = 0; i < 4; ++i)
        pool.submit([&counter]{ counter.fetch_add(1, std::memory_order_relaxed); });

    pool.drain();
    REQUIRE(counter.load() == 4);
    pool.stop();
}