KPN/include/kpn/branch.hpp

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

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

namespace kpn {

// ── Lossless single-output delivery ───────────────────────────────────────────
//
// Shared by RouterNode and FilterNode, which each deliver a value to exactly one
// channel. Both previously did
//
//     try { ch->push(val); } catch (const ChannelOverflowError&) {}
//
// which discards the value whenever the consumer is behind. 6595e6e made node
// outputs lossless, 28e0667 stopped them parking a worker, and a8cfe73 did the
// same for FanoutNode — these two were in none of them, and were the last
// remaining users of the throwing push() on a data path.
//
// A dropped item does not degrade a downstream result, it silently changes one.
// Worse, a dropped *sentinel* wedges the pipeline outright: EOF is what tells
// every downstream node to shut down, and there is nothing after it to retry.
// A filter that passes EOF by predicate but drops it by backpressure is a
// pipeline that never terminates.
//
// So sentinels go out-of-band via push_sentinel (a dedicated slot that consumes
// no ring capacity and cannot overflow), and everything else is retried until
// taken. Like FanoutNode and unlike a pool node, these own a private thread, so
// waiting here costs no scheduler worker and needs no space-callback park.
// stop_flag_ is rechecked every pass so teardown cannot hang on a full output.
//
// `parked` receives the time spent waiting, which the caller charges to blocked
// rather than exec — a parked node is idle, and charging it to exec reports the
// node as busy exactly when it is the one being held up.
//
// Returns false if stopped with the value undelivered.
template<typename T>
bool deliver_one(Channel<T>* ch, T& val, const std::atomic<bool>& stop_flag,
                 duration_t& parked) {
    if (is_sentinel_value(val)) {
        ch->push_sentinel(std::move(val));
        return true;
    }
    const auto park_from = clock_t::now();
    for (;;) {
        switch (ch->try_push(val)) {
            case Channel<T>::PushResult::Taken:
                parked = duration_t(clock_t::now() - park_from);
                return true;
            case Channel<T>::PushResult::Closed:
                // Nobody is listening any more; the channel has recorded the
                // drop. Retrying would spin until teardown noticed.
                parked = duration_t(clock_t::now() - park_from);
                return false;
            case Channel<T>::PushResult::Full:
                break;                       // fall through to the retry logic
        }
        if (stop_flag.load(std::memory_order_relaxed)) {
            // Teardown with work in hand and the output still full. One last
            // throwing push, purely so the channel's own stats record the
            // overflow — the point of the lossless path is that a loss is never
            // invisible, and a silent return here would reintroduce exactly the
            // hole this function exists to close.
            try { ch->push(std::move(val)); }
            catch (const ChannelOverflowError&) {}
            parked = duration_t(clock_t::now() - park_from);
            return false;
        }
        std::this_thread::sleep_for(std::chrono::microseconds(50));
    }
}

// ── 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();

                // An out-of-range selector still drops by design (documented on
                // the class): the item was routed nowhere, not lost to a full
                // channel. Only the latter is what deliver_one exists to stop.
                std::size_t idx = selector_(val);
                duration_t  parked{0};
                bool        delivered = true;
                if (idx < N && out_channels_[idx])
                    delivered = deliver_one(out_channels_[idx], val, stop_flag_, parked);

                auto cpu1 = NodeStats::cpu_now();
                auto t2   = clock_t::now();
                stats_.record_exec(duration_t(t2 - t1) - parked,
                                   duration_t(t1 - t0) + parked, cpu0, cpu1);
                if (!delivered) break;
            } 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();

                // A value the predicate rejects is dropped by design and is not
                // counted as a processed frame. One it accepts is now delivered
                // losslessly — including a sentinel, which a filter typically
                // passes unconditionally so downstream can shut down, and which
                // the old throwing push discarded whenever the output was full.
                if (pred_(val) && out_ch_) {
                    duration_t parked{0};
                    const bool delivered = deliver_one(out_ch_, val, stop_flag_, parked);
                    auto cpu1 = NodeStats::cpu_now();
                    auto t2   = clock_t::now();
                    stats_.record_exec(duration_t(t2 - t1) - parked,
                                       duration_t(t1 - t0) + parked, cpu0, cpu1);
                    if (!delivered) break;
                }
            } 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