KPN/include/kpn/node.hpp

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

#include <array>
#include <atomic>
#include <chrono>
#include <cstddef>
#include <functional>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <thread>
#include <tuple>
#include <type_traits>

namespace kpn {

// Called when a node's function throws. Return true to skip the failed
// invocation and keep running, false to stop the node.
using NodeErrorHandler = std::function<bool(std::string_view node_name, std::exception_ptr)>;

// ── INode — type-erased interface for Network / watchdog ─────────────────────

struct INode {
    virtual ~INode()       = default;
    virtual void start()   = 0;
    virtual void stop()    = 0;
    virtual bool running() const = 0;
    virtual const NodeStats& stats() const = 0;
    virtual NodeSnapshot node_snapshot(const std::string& name, double elapsed_s) const = 0;
    virtual void set_name(std::string name) = 0;
};

// ── Node ─────────────────────────────────────────────────────────────────────
//
// Template parameters:
//   Func        — the wrapped function (auto NTTP, deduced as a function pointer)
//   InputNames  — optional kpn::in<"a","b"> tag type (at most one)
//   OutputNames — optional kpn::out<"x","y"> tag type (at most one)

template<auto Func,
         typename InputTag  = in<>,
         typename OutputTag = out<>,
         fixed_string Label    = "",
         std::size_t UniqueTag = 0>
class Node;

// Specialisation that unpacks the in<>/out<> tag packs
template<auto Func, fixed_string... InNames, fixed_string... OutNames,
         fixed_string Label, std::size_t UniqueTag>
class Node<Func, in<InNames...>, out<OutNames...>, Label, UniqueTag> : public INode {
public:
    using F            = decltype(Func);
    using args_tuple   = args_t<F>;
    using return_raw   = return_t<F>;
    using return_tuple = normalised_return_t<return_raw>;

    // Identity accessors — used by StaticNetwork for diagnostics and type-level uniqueness
    static constexpr std::string_view label()     { return Label.view(); }
    static constexpr std::size_t      unique_tag  = UniqueTag;

    static constexpr std::size_t input_count  = arity_v<F>;
    static constexpr std::size_t output_count = std::tuple_size_v<return_tuple>;

    static_assert(
        sizeof...(InNames) == 0 || sizeof...(InNames) == input_count,
        "make_node: number of input names must match function arity, or provide none"
    );
    static_assert(
        sizeof...(OutNames) == 0 || sizeof...(OutNames) == output_count,
        "make_node: number of output names must match return tuple size, or provide none"
    );

    explicit Node(std::size_t fifo_capacity = 5)
        : fifo_capacity_(fifo_capacity)
    {
        init_input_channels(std::make_index_sequence<input_count>{});
    }

    ~Node() override { stop(); }

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

    void start() override {
        enable_inputs(std::make_index_sequence<input_count>{});
        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);
        // Disable all input channels: drops queued items and unblocks waiting pop()
        disable_inputs(std::make_index_sequence<input_count>{});
        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); }

    void set_error_handler(NodeErrorHandler h) { error_handler_ = std::move(h); }

    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 — by index ────────────────────────────────────────────────

    template<std::size_t I>
    InputPort<Node, I> input() {
        static_assert(I < input_count, "input index out of range");
        return {*this};
    }

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

    // ── Port access — by name ─────────────────────────────────────────────────

    template<fixed_string Name>
    auto input() {
        constexpr std::size_t idx = index_of<Name, InNames...>();
        static_assert(idx != npos, "unknown input port name");
        return input<idx>();
    }

    template<fixed_string Name>
    auto output() {
        constexpr std::size_t idx = index_of<Name, OutNames...>();
        static_assert(idx != npos, "unknown output port name");
        return output<idx>();
    }

    // ── Internal channel accessors (used by Network at connect time) ──────────

    template<std::size_t I>
    Channel<std::tuple_element_t<I, args_tuple>>& input_channel() {
        return *std::get<I>(input_channels_);
    }

    // Replace the owned input channel with an externally provided one.
    // Used by VariantNodeWrapper to share a Channel<T> with a VariantChannel adapter.
    template<std::size_t I>
    void set_input_channel(
        std::shared_ptr<Channel<std::tuple_element_t<I, args_tuple>>> ch) {
        std::get<I>(input_channels_) = std::move(ch);
    }

    template<std::size_t I>
    void set_output_channel(
        Channel<std::tuple_element_t<I, return_tuple>>* ch) {
        std::get<I>(output_channels_) = ch;
    }

private:
    // ── Channel storage ───────────────────────────────────────────────────────

    // Input channels — shared ownership so VariantChannel adapters can share them
    template<std::size_t... Is>
    void init_input_channels(std::index_sequence<Is...>) {
        ((std::get<Is>(input_channels_) =
              std::make_shared<Channel<std::tuple_element_t<Is, args_tuple>>>(fifo_capacity_)),
         ...);
    }

    template<std::size_t... Is>
    void enable_inputs(std::index_sequence<Is...>) {
        (std::get<Is>(input_channels_)->enable(), ...);
    }

    template<std::size_t... Is>
    void disable_inputs(std::index_sequence<Is...>) {
        (std::get<Is>(input_channels_)->disable(), ...);
    }

    template<typename Tup, std::size_t... Is>
    static auto make_input_channel_tuple(std::index_sequence<Is...>)
        -> std::tuple<std::shared_ptr<Channel<std::tuple_element_t<Is, Tup>>>...>;

    using input_channels_t  = decltype(make_input_channel_tuple<args_tuple>(
        std::make_index_sequence<input_count>{}));

    // Output channels — non-owning pointers, set at connect time
    template<typename Tup, std::size_t... Is>
    static auto make_output_channel_tuple(std::index_sequence<Is...>)
        -> std::tuple<Channel<std::tuple_element_t<Is, Tup>>*...>;

    using output_channels_t = decltype(make_output_channel_tuple<return_tuple>(
        std::make_index_sequence<output_count>{}));

    // ── run_loop ──────────────────────────────────────────────────────────────

    void run_loop() {
        while (!stop_flag_.load(std::memory_order_relaxed)) {
            try {
                auto t0 = clock_t::now();
                auto args = pop_inputs(std::make_index_sequence<input_count>{});
                auto t1 = clock_t::now();
                auto cpu0 = NodeStats::cpu_now();

                if constexpr (std::is_void_v<return_raw>) {
                    std::apply(Func, args);
                } else {
                    auto result = std::apply(Func, args);
                    push_outputs(normalise(std::move(result)),
                                 std::make_index_sequence<output_count>{});
                }

                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;
            } catch (const ChannelOverflowError& e) {
                std::cerr << "[kpn] overflow: " << e.what() << "\n";
            } catch (...) {
                if (error_handler_ && error_handler_(name_, std::current_exception()))
                    continue;
                break;
            }
        }
        stop_flag_.store(true, std::memory_order_relaxed);
    }

    // Pop all inputs into a tuple of argument values
    template<std::size_t... Is>
    args_tuple pop_inputs(std::index_sequence<Is...>) {
        return {std::get<Is>(input_channels_)->pop()...};
    }

    // Normalise return value to tuple (handles void and single-value returns)
    template<typename R = return_raw>
    static return_tuple normalise(R&& r) {
        if constexpr (is_tuple_v<R>)
            return std::move(r);
        else
            return std::make_tuple(std::move(r));
    }

    static return_tuple normalise_void() { return {}; }

    // Push each output element to its connected channel (if connected)
    template<std::size_t... Is>
    void push_outputs(return_tuple&& result, std::index_sequence<Is...>) {
        (push_one<Is>(std::get<Is>(std::move(result))), ...);
    }

    template<std::size_t I>
    void push_one(std::tuple_element_t<I, return_tuple>&& val) {
        auto* ch = std::get<I>(output_channels_);
        if (!ch) return;
        try {
            ch->push(std::move(val));
        } catch (const ChannelOverflowError&) {
            throw ChannelOverflowError(ch->capacity(), "node '" + name_ + "' " + output_port_label<I>());
        }
    }

    template<std::size_t I>
    static std::string output_port_label() {
        if constexpr (sizeof...(OutNames) > 0) {
            constexpr std::array<std::string_view, sizeof...(OutNames)> names{OutNames.view()...};
            return std::string("output['") + std::string(names[I]) + "']";
        } else {
            return "output[" + std::to_string(I) + "]";
        }
    }

    // ── State ─────────────────────────────────────────────────────────────────

    std::string         name_;
    std::size_t         fifo_capacity_;
    input_channels_t    input_channels_;
    output_channels_t   output_channels_{};
    std::atomic<bool>   stop_flag_{false};
    std::jthread        thread_;
    NodeStats           stats_;
    NodeErrorHandler    error_handler_;
};

// ── ObjectNode — wraps a callable object (functor / class with operator()) ────
//
// Use this when the node needs state initialised in a constructor.
// The object must outlive the ObjectNode (stored by reference).
//
// Usage:
//   MyFunctor obj(...);
//   auto node = make_node(obj, in<"x">{}, out<"y">{}, capacity);

template<typename Obj,
         typename InputTag  = in<>,
         typename OutputTag = out<>,
         fixed_string Label    = "",
         std::size_t UniqueTag = 0>
class ObjectNode;

template<typename Obj, fixed_string... InNames, fixed_string... OutNames,
         fixed_string Label, std::size_t UniqueTag>
class ObjectNode<Obj, in<InNames...>, out<OutNames...>, Label, UniqueTag> : public INode {
public:
    using F            = decltype(&Obj::operator());
    using args_tuple   = args_t<F>;
    using return_raw   = return_t<F>;
    using return_tuple = normalised_return_t<return_raw>;

    static constexpr std::string_view label()    { return Label.view(); }
    static constexpr std::size_t      unique_tag = UniqueTag;

    static constexpr std::size_t input_count  = arity_v<F>;
    static constexpr std::size_t output_count = std::tuple_size_v<return_tuple>;

    static_assert(
        sizeof...(InNames) == 0 || sizeof...(InNames) == input_count,
        "make_node: number of input names must match operator() arity, or provide none"
    );
    static_assert(
        sizeof...(OutNames) == 0 || sizeof...(OutNames) == output_count,
        "make_node: number of output names must match return tuple size, or provide none"
    );

    explicit ObjectNode(Obj& obj, std::size_t fifo_capacity = 5)
        : obj_(obj), fifo_capacity_(fifo_capacity)
    {
        init_input_channels(std::make_index_sequence<input_count>{});
    }

    ~ObjectNode() override { stop(); }

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

    void start() override {
        enable_inputs(std::make_index_sequence<input_count>{});
        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);
        disable_inputs(std::make_index_sequence<input_count>{});
        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); }

    void set_error_handler(NodeErrorHandler h) { error_handler_ = std::move(h); }

    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>
    InputPort<ObjectNode, I> input() {
        static_assert(I < input_count, "input index out of range");
        return {*this};
    }

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

    template<fixed_string Name>
    auto input() {
        constexpr std::size_t idx = index_of<Name, InNames...>();
        static_assert(idx != npos, "unknown input port name");
        return input<idx>();
    }

    template<fixed_string Name>
    auto output() {
        constexpr std::size_t idx = index_of<Name, OutNames...>();
        static_assert(idx != npos, "unknown output port name");
        return output<idx>();
    }

    template<std::size_t I>
    Channel<std::tuple_element_t<I, args_tuple>>& input_channel() {
        return *std::get<I>(input_channels_);
    }

    template<std::size_t I>
    void set_input_channel(
        std::shared_ptr<Channel<std::tuple_element_t<I, args_tuple>>> ch) {
        std::get<I>(input_channels_) = std::move(ch);
    }

    template<std::size_t I>
    void set_output_channel(Channel<std::tuple_element_t<I, return_tuple>>* ch) {
        std::get<I>(output_channels_) = ch;
    }

private:
    template<std::size_t... Is>
    void init_input_channels(std::index_sequence<Is...>) {
        ((std::get<Is>(input_channels_) =
              std::make_shared<Channel<std::tuple_element_t<Is, args_tuple>>>(fifo_capacity_)),
         ...);
    }

    template<std::size_t... Is>
    void enable_inputs(std::index_sequence<Is...>) {
        (std::get<Is>(input_channels_)->enable(), ...);
    }

    template<std::size_t... Is>
    void disable_inputs(std::index_sequence<Is...>) {
        (std::get<Is>(input_channels_)->disable(), ...);
    }

    template<typename Tup, std::size_t... Is>
    static auto make_input_channel_tuple(std::index_sequence<Is...>)
        -> std::tuple<std::shared_ptr<Channel<std::tuple_element_t<Is, Tup>>>...>;

    using input_channels_t = decltype(make_input_channel_tuple<args_tuple>(
        std::make_index_sequence<input_count>{}));

    template<typename Tup, std::size_t... Is>
    static auto make_output_channel_tuple(std::index_sequence<Is...>)
        -> std::tuple<Channel<std::tuple_element_t<Is, Tup>>*...>;

    using output_channels_t = decltype(make_output_channel_tuple<return_tuple>(
        std::make_index_sequence<output_count>{}));

    void run_loop() {
        while (!stop_flag_.load(std::memory_order_relaxed)) {
            try {
                auto t0   = clock_t::now();
                auto args = pop_inputs(std::make_index_sequence<input_count>{});
                auto t1   = clock_t::now();
                auto cpu0 = NodeStats::cpu_now();

                if constexpr (std::is_void_v<return_raw>) {
                    std::apply([this](auto&&... a) { obj_(std::forward<decltype(a)>(a)...); }, args);
                } else {
                    auto result = std::apply([this](auto&&... a) { return obj_(std::forward<decltype(a)>(a)...); }, args);
                    push_outputs(normalise(std::move(result)),
                                 std::make_index_sequence<output_count>{});
                }

                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;
            } catch (const ChannelOverflowError& e) {
                std::cerr << "[kpn] overflow: " << e.what() << "\n";
            } catch (...) {
                if (error_handler_ && error_handler_(name_, std::current_exception()))
                    continue;
                break;
            }
        }
        stop_flag_.store(true, std::memory_order_relaxed);
    }

    template<std::size_t... Is>
    args_tuple pop_inputs(std::index_sequence<Is...>) {
        return {std::get<Is>(input_channels_)->pop()...};
    }

    template<typename R = return_raw>
    static return_tuple normalise(R&& r) {
        if constexpr (is_tuple_v<R>) return std::move(r);
        else                         return std::make_tuple(std::move(r));
    }

    template<std::size_t... Is>
    void push_outputs(return_tuple&& result, std::index_sequence<Is...>) {
        (push_one<Is>(std::get<Is>(std::move(result))), ...);
    }

    template<std::size_t I>
    void push_one(std::tuple_element_t<I, return_tuple>&& val) {
        auto* ch = std::get<I>(output_channels_);
        if (!ch) return;
        try {
            ch->push(std::move(val));
        } catch (const ChannelOverflowError&) {
            throw ChannelOverflowError(ch->capacity(), "node '" + name_ + "' " + output_port_label<I>());
        }
    }

    template<std::size_t I>
    static std::string output_port_label() {
        if constexpr (sizeof...(OutNames) > 0) {
            constexpr std::array<std::string_view, sizeof...(OutNames)> names{OutNames.view()...};
            return std::string("output['") + std::string(names[I]) + "']";
        } else {
            return "output[" + std::to_string(I) + "]";
        }
    }

    Obj&                obj_;
    std::string         name_;
    std::size_t         fifo_capacity_;
    input_channels_t    input_channels_;
    output_channels_t   output_channels_{};
    std::atomic<bool>   stop_flag_{false};
    std::jthread        thread_;
    NodeStats           stats_;
    NodeErrorHandler    error_handler_;
};

// ── make_node overloads for callable objects ──────────────────────────────────

template<typename Obj>
auto make_node(Obj& obj, std::size_t fifo_capacity = 5) {
    return ObjectNode<Obj, in<>, out<>>(obj, fifo_capacity);
}

template<typename Obj, fixed_string... InNames>
auto make_node(Obj& obj, in<InNames...>, std::size_t fifo_capacity = 5) {
    return ObjectNode<Obj, in<InNames...>, out<>>(obj, fifo_capacity);
}

template<typename Obj, fixed_string... OutNames>
auto make_node(Obj& obj, out<OutNames...>, std::size_t fifo_capacity = 5) {
    return ObjectNode<Obj, in<>, out<OutNames...>>(obj, fifo_capacity);
}

template<typename Obj, fixed_string... InNames, fixed_string... OutNames>
auto make_node(Obj& obj, in<InNames...>, out<OutNames...>, std::size_t fifo_capacity = 5) {
    return ObjectNode<Obj, in<InNames...>, out<OutNames...>>(obj, fifo_capacity);
}

// ── make_node factory (NTTP) ──────────────────────────────────────────────────
//
// Usage:
//   make_node<func>(capacity)
//   make_node<func, "label">(capacity)
//   make_node<func, "label", 1>(capacity)          // UniqueTag=1
//   make_node<func>(in<"a","b">{}, capacity)
//   make_node<func, "label">(in<"a","b">{}, out<"x">{}, capacity)

// No port names
template<auto Func, fixed_string Label = "", std::size_t UniqueTag = 0>
auto make_node(std::size_t fifo_capacity = 5) {
    return Node<Func, in<>, out<>, Label, UniqueTag>(fifo_capacity);
}

// in<> only
template<auto Func, fixed_string Label = "", std::size_t UniqueTag = 0,
         fixed_string... InNames>
auto make_node(in<InNames...>, std::size_t fifo_capacity = 5) {
    return Node<Func, in<InNames...>, out<>, Label, UniqueTag>(fifo_capacity);
}

// out<> only
template<auto Func, fixed_string Label = "", std::size_t UniqueTag = 0,
         fixed_string... OutNames>
auto make_node(out<OutNames...>, std::size_t fifo_capacity = 5) {
    return Node<Func, in<>, out<OutNames...>, Label, UniqueTag>(fifo_capacity);
}

// in<> and out<>
template<auto Func, fixed_string Label = "", std::size_t UniqueTag = 0,
         fixed_string... InNames, fixed_string... OutNames>
auto make_node(in<InNames...>, out<OutNames...>, std::size_t fifo_capacity = 5) {
    return Node<Func, in<InNames...>, out<OutNames...>, Label, UniqueTag>(fifo_capacity);
}

} // namespace kpn