#include <catch2/catch_test_macros.hpp>
#include <kpn/scheduler.hpp>
#include <kpn/pool_node.hpp>
#include <kpn/interrupt_node.hpp>
#include <atomic>
#include <chrono>
#include <mutex>
#include <thread>

using namespace kpn;

static int double_it(int x) { return x * 2; }
static std::tuple<int, float> split_it(int x) { return {x, float(x) * 0.5f}; }
static void consume_it(int x) { (void)x; }

// ── ThreadPool ────────────────────────────────────────────────────────────────

TEST_CASE("thread pool starts and stops cleanly", "[scheduler]") {
    ThreadPool pool(2);
    pool.start();
    pool.stop();
}

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

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

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

TEST_CASE("thread pool drain waits for all tasks", "[scheduler]") {
    ThreadPool pool(1);
    pool.start();

    std::atomic<bool> done{false};
    pool.submit([&] {
        std::this_thread::sleep_for(std::chrono::milliseconds(20));
        done.store(true);
    });
    pool.drain();
    REQUIRE(done.load());
    pool.stop();
}

TEST_CASE("thread pool priority: higher priority tasks run first", "[scheduler]") {
    ThreadPool pool(1);  // single thread so order is deterministic
    pool.start();

    // Submit a task that blocks the worker, then queue two tasks with
    // different priorities. When the blocker finishes, the high-priority
    // task should run before the low-priority one.
    std::vector<int> order;
    std::mutex order_mutex;

    std::atomic<bool> blocker_done{false};
    pool.submit([&] {
        std::this_thread::sleep_for(std::chrono::milliseconds(30));
        blocker_done.store(true);
    }, 0.5f);

    // Wait until blocker is running, then enqueue the two ordered tasks.
    while (!blocker_done.load()) std::this_thread::sleep_for(std::chrono::milliseconds(1));

    pool.submit([&] { std::lock_guard g(order_mutex); order.push_back(1); }, 0.1f);
    pool.submit([&] { std::lock_guard g(order_mutex); order.push_back(2); }, 0.9f);

    pool.drain();
    REQUIRE(order == std::vector<int>{2, 1});
    pool.stop();
}

// ── PoolNode ──────────────────────────────────────────────────────────────────

TEST_CASE("pool node input/output counts", "[pool_node]") {
    STATIC_REQUIRE(PoolNode<double_it>::input_count  == 1);
    STATIC_REQUIRE(PoolNode<double_it>::output_count == 1);
    STATIC_REQUIRE(PoolNode<split_it>::output_count  == 2);
    STATIC_REQUIRE(PoolNode<consume_it>::output_count == 0);
}

TEST_CASE("pool node processes items end-to-end", "[pool_node]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<double_it>(pool);
    Channel<int> out_ch(10);
    node.set_output_channel<0>(&out_ch);
    node.start();

    node.input_channel<0>().push(21);
    int result = out_ch.pop();

    node.stop();
    pool->stop();

    REQUIRE(result == 42);
}

TEST_CASE("pool node processes multiple items in order", "[pool_node]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<double_it>(pool, 20);  // capacity 20
    Channel<int> out_ch(20);
    node.set_output_channel<0>(&out_ch);
    node.start();

    constexpr int N = 10;
    for (int i = 0; i < N; ++i)
        node.input_channel<0>().push(i);

    std::vector<int> results;
    for (int i = 0; i < N; ++i)
        results.push_back(out_ch.pop());

    node.stop();
    pool->stop();

    REQUIRE(results.size() == N);
    for (int i = 0; i < N; ++i)
        REQUIRE(results[i] == i * 2);
}

TEST_CASE("pool node stop is clean with no deadlock", "[pool_node]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<double_it>(pool);
    node.start();
    // Node is idle (no input pushed) — stop must return without deadlock.
    node.stop();
    REQUIRE_FALSE(node.running());

    pool->stop();
}

TEST_CASE("pool node two-stage pipeline produces correct count", "[pool_node]") {
    auto pool = std::make_shared<ThreadPool>(4);
    pool->start();

    auto src       = make_pool_node<double_it>(pool);
    auto transform = make_pool_node<double_it>(pool);
    Channel<int> out_ch(20);

    // Wire src output → transform input channel, transform output → out_ch.
    src.set_output_channel<0>(&transform.input_channel<0>());
    transform.set_output_channel<0>(&out_ch);

    src.start();
    transform.start();

    constexpr int N = 5;
    for (int i = 1; i <= N; ++i)
        src.input_channel<0>().push(i);

    std::vector<int> results;
    for (int i = 0; i < N; ++i)
        results.push_back(out_ch.pop());

    // Stop nodes before they (and their channels) go out of scope.
    src.stop();
    transform.stop();
    pool->stop();

    REQUIRE(results.size() == static_cast<std::size_t>(N));
    for (int i = 0; i < N; ++i)
        REQUIRE(results[i] == (i + 1) * 4);  // double_it twice
}

// ── InterruptNode ─────────────────────────────────────────────────────────────

namespace {
static std::atomic<int> g_interrupt_counter{0};
static int interrupt_produce() { return g_interrupt_counter.fetch_add(1); }
} // namespace

TEST_CASE("interrupt node fires on each trigger", "[interrupt_node]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    g_interrupt_counter.store(0);
    auto node = make_interrupt_node<interrupt_produce>(pool, out<>{});
    Channel<int> out_ch(20);
    node.set_output_channel<0>(&out_ch);
    node.start();

    auto trigger = node.get_trigger();
    constexpr int N = 5;
    for (int i = 0; i < N; ++i) trigger();

    std::vector<int> results;
    for (int i = 0; i < N; ++i)
        results.push_back(out_ch.pop());

    node.stop();
    pool->stop();

    REQUIRE(results.size() == static_cast<std::size_t>(N));
}

TEST_CASE("interrupt node does not fire without trigger", "[interrupt_node]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    g_interrupt_counter.store(0);
    auto node = make_interrupt_node<interrupt_produce>(pool, out<>{});
    Channel<int> out_ch(5);
    node.set_output_channel<0>(&out_ch);
    node.start();

    std::this_thread::sleep_for(std::chrono::milliseconds(30));
    // No trigger fired — output channel should be empty.
    REQUIRE(out_ch.approx_size() == 0);

    node.stop();
    pool->stop();
}

TEST_CASE("interrupt node: trigger after stop is ignored", "[interrupt_node]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    g_interrupt_counter.store(0);
    auto node = make_interrupt_node<interrupt_produce>(pool, out<>{});
    Channel<int> out_ch(5);
    node.set_output_channel<0>(&out_ch);
    node.start();

    auto trigger = node.get_trigger();
    node.stop();
    trigger();  // should be a no-op

    std::this_thread::sleep_for(std::chrono::milliseconds(10));
    REQUIRE(out_ch.approx_size() == 0);
    pool->stop();
}

// ── Overflow callback ─────────────────────────────────────────────────────────

TEST_CASE("pool node overflow callback fires on full output channel", "[pool_node][overflow]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<double_it>(pool);
    // Pre-fill a tiny channel so every node push overflows.
    Channel<int> full_ch(1);
    full_ch.push(99);
    node.set_output_channel<0>(&full_ch);

    std::atomic<int> overflow_count{0};
    node.set_overflow_callback([&](auto) { overflow_count.fetch_add(1); });

    node.start();
    node.input_channel<0>().push(1);
    node.input_channel<0>().push(2);

    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    node.stop();
    pool->stop();

    REQUIRE(overflow_count.load() > 0);
}

TEST_CASE("pool node overflow callback is independent per instance", "[pool_node][overflow]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto nodeA = make_pool_node<double_it>(pool);
    auto nodeB = make_pool_node<double_it>(pool);

    std::atomic<int> a_overflows{0}, b_overflows{0};
    nodeA.set_overflow_callback([&](auto) { a_overflows.fetch_add(1); });

    Channel<int> full_ch(1);
    full_ch.push(0);
    nodeA.set_output_channel<0>(&full_ch);

    Channel<int> ok_ch(20);
    nodeB.set_output_channel<0>(&ok_ch);

    nodeA.start();
    nodeB.start();

    nodeA.input_channel<0>().push(1);
    nodeA.input_channel<0>().push(2);
    nodeB.input_channel<0>().push(10);

    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    nodeA.stop();
    nodeB.stop();
    pool->stop();

    REQUIRE(a_overflows.load() > 0);
    REQUIRE(b_overflows.load() == 0);
}

TEST_CASE("interrupt node overflow callback fires on full output", "[interrupt_node][overflow]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    g_interrupt_counter.store(0);
    auto node = make_interrupt_node<interrupt_produce>(pool, out<>{});

    Channel<int> full_ch(1);
    full_ch.push(99);
    node.set_output_channel<0>(&full_ch);

    std::atomic<int> overflow_count{0};
    node.set_overflow_callback([&](auto) { overflow_count.fetch_add(1); });

    node.start();
    auto trigger = node.get_trigger();
    trigger(); trigger(); trigger();

    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    node.stop();
    pool->stop();

    REQUIRE(overflow_count.load() > 0);
}

// ── self_stop: disable inputs + outputs on crash ──────────────────────────────

static int always_throw(int) { throw std::runtime_error("node crashed"); return 0; }

TEST_CASE("pool node self_stop disables output on crash so downstream sees closed", "[pool_node][self_stop]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<always_throw>(pool, 5);
    Channel<int> out_ch(10);
    node.set_output_channel<0>(&out_ch);

    node.start();
    node.input_channel<0>().push(1);

    std::this_thread::sleep_for(std::chrono::milliseconds(100));

    REQUIRE_FALSE(out_ch.is_accepting());

    node.stop();
    pool->stop();
}

TEST_CASE("pool node self_stop disables input on crash", "[pool_node][self_stop]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<always_throw>(pool, 5);
    Channel<int> out_ch(5);
    node.set_output_channel<0>(&out_ch);

    node.start();
    node.input_channel<0>().push(1);

    std::this_thread::sleep_for(std::chrono::milliseconds(100));

    REQUIRE_FALSE(node.input_channel<0>().is_accepting());

    node.stop();
    pool->stop();
}

TEST_CASE("pool node closed callback fires on self_stop from crash", "[pool_node][self_stop]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<always_throw>(pool, 5);
    Channel<int> out_ch(5);
    node.set_output_channel<0>(&out_ch);

    std::atomic<bool> closed_fired{false};
    node.set_closed_callback([&](auto) { closed_fired.store(true); });

    node.start();
    node.input_channel<0>().push(1);

    std::this_thread::sleep_for(std::chrono::milliseconds(100));

    REQUIRE(closed_fired.load());

    node.stop();
    pool->stop();
}

// ── Network-level event callbacks ─────────────────────────────────────────────

TEST_CASE("network_overflow_callback fires on overflow", "[pool_node][network]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<double_it>(pool);

    Channel<int> full_ch(1);
    full_ch.push(0);
    node.set_output_channel<0>(&full_ch);

    std::atomic<int> net_overflows{0};
    node.set_network_overflow_callback([&](auto) { net_overflows.fetch_add(1); });

    node.start();
    node.input_channel<0>().push(1);
    node.input_channel<0>().push(2);

    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    node.stop();
    pool->stop();

    REQUIRE(net_overflows.load() > 0);
}

TEST_CASE("network_closed_callback fires on crash", "[pool_node][network]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<always_throw>(pool);
    Channel<int> out_ch(5);
    node.set_output_channel<0>(&out_ch);

    std::atomic<bool> net_closed{false};
    node.set_network_closed_callback([&](auto) { net_closed.store(true); });

    node.start();
    node.input_channel<0>().push(1);

    std::this_thread::sleep_for(std::chrono::milliseconds(100));

    REQUIRE(net_closed.load());

    node.stop();
    pool->stop();
}

TEST_CASE("per-node and network overflow callbacks both fire independently", "[pool_node][network]") {
    auto pool = std::make_shared<ThreadPool>(2);
    pool->start();

    auto node = make_pool_node<double_it>(pool);

    Channel<int> full_ch(1);
    full_ch.push(0);
    node.set_output_channel<0>(&full_ch);

    std::atomic<int> per_node{0}, network{0};
    node.set_overflow_callback([&](auto) { per_node.fetch_add(1); });
    node.set_network_overflow_callback([&](auto) { network.fetch_add(1); });

    node.start();
    node.input_channel<0>().push(1);
    node.input_channel<0>().push(2);

    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    node.stop();
    pool->stop();

    REQUIRE(per_node.load() > 0);
    REQUIRE(network.load() > 0);
}