KPN/include/kpn/fanout.hpp

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

#include <array>
#include <atomic>
#include <iostream>
#include <memory>
#include <thread>
#include <tuple>
#include <utility>

namespace kpn {

namespace detail {

// Produces std::tuple<T, T, ..., T> with N repetitions — used so that
// Network::connect can do its normal return_tuple type-check against FanoutNode.
template<typename T, std::size_t N, typename Seq = std::make_index_sequence<N>>
struct repeat_tuple;

template<typename T, std::size_t N, std::size_t... Is>
struct repeat_tuple<T, N, std::index_sequence<Is...>> {
    template<std::size_t> using always_T = T;
    using type = std::tuple<always_T<Is>...>;
};

template<typename T, std::size_t N>
using repeat_tuple_t = typename repeat_tuple<T, N>::type;

} // namespace detail

// ── FanoutNode ────────────────────────────────────────────────────────────────
//
// Reads one item from its single input channel and pushes a copy to each of
// N output channels. All N downstream nodes receive every item.
//
// Usage:
//   auto fan = make_fanout<Image, 2>(/*capacity=*/8);
//   net.connect("src",  src.output<0>(),   "fan",   fan.input<0>())
//      .connect("fan",  fan.output<0>(),   "nodeA", nodeA.input<0>())
//      .connect("fan",  fan.output<1>(),   "nodeB", nodeB.input<0>())

template<typename T, std::size_t N, std::size_t Id = 0>
class FanoutNode : public INode {
public:
    using args_tuple   = std::tuple<T>;
    using return_tuple = detail::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;

    explicit FanoutNode(std::size_t fifo_capacity = 5)
        : fifo_capacity_(fifo_capacity)
    {
        input_ch_ = std::make_shared<Channel<T>>(fifo_capacity);
    }

    ~FanoutNode() 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<FanoutNode, I> input() {
        static_assert(I == 0, "FanoutNode has exactly one input");
        return {*this};
    }

    template<std::size_t I>
    OutputPort<FanoutNode, I> output() {
        static_assert(I < N, "FanoutNode 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();

                for (std::size_t i = 0; i < N; ++i) {
                    if (out_channels_[i]) {
                        try { out_channels_[i]->push(val); }
                        catch (const ChannelOverflowError&) {}  // drop for this output independently
                    }
                }

                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_;
    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_;
};

// ── Factory ───────────────────────────────────────────────────────────────────

template<typename T, std::size_t N>
FanoutNode<T, N> make_fanout(std::size_t fifo_capacity = 5) {
    return FanoutNode<T, N, 0>(fifo_capacity);
}

} // namespace kpn