#pragma once
#include "channel.hpp"
#include "diagnostics.hpp"
#include "inode.hpp"
#include "port.hpp"
#include "traits.hpp"

#include <array>
#include <atomic>
#include <functional>
#include <memory>
#include <thread>

namespace kpn {

// ── RouterNode ────────────────────────────────────────────────────────────────
//
// Reads one item and pushes it to exactly one of N output channels, chosen by
// selector(item). If selector returns >= N the item is silently dropped.
//
// Usage:
//   auto router = make_router<Image, 3>(
//       [](const Image& img) -> std::size_t { return img.stream_id % 3; });
//   net.connect("src",    src.output<0>(),     "router", router.input<0>())
//      .connect("router", router.output<0>(),  "nodeA",  nodeA.input<0>())
//      .connect("router", router.output<1>(),  "nodeB",  nodeB.input<0>())
//      .connect("router", router.output<2>(),  "nodeC",  nodeC.input<0>());

template<typename T, std::size_t N, std::size_t Id = 0>
class RouterNode : public INode {
public:
    using Selector     = std::function<std::size_t(const T&)>;
    using args_tuple   = std::tuple<T>;
    using return_tuple = repeat_tuple_t<T, N>;
    using return_raw   = return_tuple;

    static constexpr std::size_t input_count    = 1;
    static constexpr std::size_t output_count   = N;
    static constexpr std::size_t unique_tag     = Id;
    static constexpr bool        is_router_node = true;

    explicit RouterNode(Selector sel, std::size_t fifo_capacity = 5)
        : selector_(std::move(sel))
        , fifo_capacity_(fifo_capacity)
    {
        input_ch_ = std::make_shared<Channel<T>>(fifo_capacity);
    }

    ~RouterNode() override { stop(); }

    // ── INode ─────────────────────────────────────────────────────────────────

    void start() override {
        input_ch_->enable();
        stop_flag_.store(false, std::memory_order_relaxed);
        thread_ = std::jthread([this](std::stop_token) { run_loop(); });
    }

    void stop() override {
        stop_flag_.store(true, std::memory_order_relaxed);
        input_ch_->disable();
        if (thread_.joinable()) thread_.request_stop(), thread_.join();
    }

    bool running() const override {
        return thread_.joinable() && !stop_flag_.load(std::memory_order_relaxed);
    }

    void set_name(std::string name) override { name_ = std::move(name); }

    const NodeStats& stats() const override { return stats_; }

    NodeSnapshot node_snapshot(const std::string& name, double elapsed_s) const override {
        uint64_t frames   = stats_.frames_processed.load(std::memory_order_relaxed);
        double exec_ms    = stats_.ema_exec_us.load(std::memory_order_relaxed) / 1000.0;
        double blocked_ms = stats_.total_blocked_us.load(std::memory_order_relaxed) / 1000.0;
        double total_ms   = exec_ms + blocked_ms;
        return {name, frames, exec_ms,
                stats_.max_exec_us.load(std::memory_order_relaxed) / 1000.0,
                blocked_ms,
                elapsed_s > 0 ? frames / elapsed_s : 0.0,
                stats_.total_cpu_us.load(std::memory_order_relaxed) / 1000.0,
                total_ms > 0 ? 100.0 * exec_ms / total_ms : 0.0};
    }

    // ── Port access ───────────────────────────────────────────────────────────

    template<std::size_t I = 0>
    InputPort<RouterNode, I> input() {
        static_assert(I == 0, "RouterNode has exactly one input");
        return {*this};
    }

    template<std::size_t I>
    OutputPort<RouterNode, I> output() {
        static_assert(I < N, "RouterNode output index out of range");
        return {*this};
    }

    // ── Internal channel accessors (called by Network::connect) ───────────────

    template<std::size_t I>
    Channel<T>& input_channel() {
        static_assert(I == 0);
        return *input_ch_;
    }

    template<std::size_t I>
    void set_input_channel(std::shared_ptr<Channel<T>> ch) {
        static_assert(I == 0);
        input_ch_ = std::move(ch);
    }

    template<std::size_t I>
    void set_output_channel(Channel<T>* ch) {
        static_assert(I < N);
        out_channels_[I] = ch;
    }

private:
    void run_loop() {
        while (!stop_flag_.load(std::memory_order_relaxed)) {
            try {
                auto        t0  = clock_t::now();
                T           val = input_ch_->pop();
                auto        t1  = clock_t::now();
                auto        cpu0 = NodeStats::cpu_now();

                std::size_t idx = selector_(val);
                if (idx < N && out_channels_[idx]) {
                    try { out_channels_[idx]->push(val); }
                    catch (const ChannelOverflowError&) {}
                }

                auto cpu1 = NodeStats::cpu_now();
                auto t2   = clock_t::now();
                stats_.record_exec(duration_t(t2 - t1), duration_t(t1 - t0), cpu0, cpu1);
            } catch (const ChannelClosedError&) {
                break;
            }
        }
    }

    std::string                 name_;
    std::size_t                 fifo_capacity_;
    Selector                    selector_;
    std::shared_ptr<Channel<T>> input_ch_;
    std::array<Channel<T>*, N>  out_channels_{};
    std::atomic<bool>           stop_flag_{false};
    std::jthread                thread_;
    NodeStats                   stats_;
};

// ── FilterNode ────────────────────────────────────────────────────────────────
//
// Reads one item and pushes it downstream only when pred(item) is true.
// Dropped items are not counted as processed frames.
//
// Usage:
//   auto filt = make_filter<Frame>([](const Frame& f) { return f.valid; });
//   net.connect("src",  src.output<0>(),  "filt", filt.input<0>())
//      .connect("filt", filt.output<0>(), "dst",  dst.input<0>());

template<typename T, std::size_t Id = 0>
class FilterNode : public INode {
public:
    using Predicate    = std::function<bool(const T&)>;
    using args_tuple   = std::tuple<T>;
    using return_tuple = std::tuple<T>;
    using return_raw   = return_tuple;

    static constexpr std::size_t input_count    = 1;
    static constexpr std::size_t output_count   = 1;
    static constexpr std::size_t unique_tag     = Id;
    static constexpr bool        is_filter_node = true;

    explicit FilterNode(Predicate pred, std::size_t fifo_capacity = 5)
        : pred_(std::move(pred))
        , fifo_capacity_(fifo_capacity)
    {
        input_ch_ = std::make_shared<Channel<T>>(fifo_capacity);
    }

    ~FilterNode() override { stop(); }

    // ── INode ─────────────────────────────────────────────────────────────────

    void start() override {
        input_ch_->enable();
        stop_flag_.store(false, std::memory_order_relaxed);
        thread_ = std::jthread([this](std::stop_token) { run_loop(); });
    }

    void stop() override {
        stop_flag_.store(true, std::memory_order_relaxed);
        input_ch_->disable();
        if (thread_.joinable()) thread_.request_stop(), thread_.join();
    }

    bool running() const override {
        return thread_.joinable() && !stop_flag_.load(std::memory_order_relaxed);
    }

    void set_name(std::string name) override { name_ = std::move(name); }

    const NodeStats& stats() const override { return stats_; }

    NodeSnapshot node_snapshot(const std::string& name, double elapsed_s) const override {
        uint64_t frames   = stats_.frames_processed.load(std::memory_order_relaxed);
        double exec_ms    = stats_.ema_exec_us.load(std::memory_order_relaxed) / 1000.0;
        double blocked_ms = stats_.total_blocked_us.load(std::memory_order_relaxed) / 1000.0;
        double total_ms   = exec_ms + blocked_ms;
        return {name, frames, exec_ms,
                stats_.max_exec_us.load(std::memory_order_relaxed) / 1000.0,
                blocked_ms,
                elapsed_s > 0 ? frames / elapsed_s : 0.0,
                stats_.total_cpu_us.load(std::memory_order_relaxed) / 1000.0,
                total_ms > 0 ? 100.0 * exec_ms / total_ms : 0.0};
    }

    // ── Port access ───────────────────────────────────────────────────────────

    template<std::size_t I = 0>
    InputPort<FilterNode, I> input() {
        static_assert(I == 0, "FilterNode has exactly one input");
        return {*this};
    }

    template<std::size_t I = 0>
    OutputPort<FilterNode, I> output() {
        static_assert(I == 0, "FilterNode has exactly one output");
        return {*this};
    }

    // ── Internal channel accessors (called by Network::connect) ───────────────

    template<std::size_t I>
    Channel<T>& input_channel() {
        static_assert(I == 0);
        return *input_ch_;
    }

    template<std::size_t I>
    void set_input_channel(std::shared_ptr<Channel<T>> ch) {
        static_assert(I == 0);
        input_ch_ = std::move(ch);
    }

    template<std::size_t I>
    void set_output_channel(Channel<T>* ch) {
        static_assert(I == 0);
        out_ch_ = ch;
    }

private:
    void run_loop() {
        while (!stop_flag_.load(std::memory_order_relaxed)) {
            try {
                auto t0  = clock_t::now();
                T    val = input_ch_->pop();
                auto t1  = clock_t::now();
                auto cpu0 = NodeStats::cpu_now();

                if (pred_(val) && out_ch_) {
                    try { out_ch_->push(val); }
                    catch (const ChannelOverflowError&) {}
                    auto cpu1 = NodeStats::cpu_now();
                    auto t2   = clock_t::now();
                    stats_.record_exec(duration_t(t2 - t1), duration_t(t1 - t0), cpu0, cpu1);
                }
            } catch (const ChannelClosedError&) {
                break;
            }
        }
    }

    std::string                 name_;
    std::size_t                 fifo_capacity_;
    Predicate                   pred_;
    std::shared_ptr<Channel<T>> input_ch_;
    Channel<T>*                 out_ch_{nullptr};
    std::atomic<bool>           stop_flag_{false};
    std::jthread                thread_;
    NodeStats                   stats_;
};

// ── Factories ─────────────────────────────────────────────────────────────────

template<typename T, std::size_t N>
RouterNode<T, N> make_router(std::function<std::size_t(const T&)> sel,
                              std::size_t capacity = 5) {
    return RouterNode<T, N, 0>(std::move(sel), capacity);
}

template<typename T>
FilterNode<T> make_filter(std::function<bool(const T&)> pred,
                           std::size_t capacity = 5) {
    return FilterNode<T, 0>(std::move(pred), capacity);
}

} // namespace kpn