KPN/include/kpn/variant_node.hpp

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

#include <array>
#include <memory>
#include <stdexcept>
#include <string>
#include <typeindex>
#include <variant>

namespace kpn {

// ── unique_types TMP helper ───────────────────────────────────────────────────

namespace detail {

template<typename Result, typename... Ts>
struct unique_types_impl;

template<typename... Unique>
struct unique_types_impl<std::tuple<Unique...>> {
    using type = std::tuple<Unique...>;
};

template<typename... Unique, typename Head, typename... Tail>
struct unique_types_impl<std::tuple<Unique...>, Head, Tail...> {
    using type = std::conditional_t<
        (std::is_same_v<Head, Unique> || ...),
        typename unique_types_impl<std::tuple<Unique...>, Tail...>::type,
        typename unique_types_impl<std::tuple<Unique..., Head>, Tail...>::type
    >;
};

template<typename... Ts>
using unique_types_t = typename unique_types_impl<std::tuple<>, Ts...>::type;

template<typename Tup>
struct tuple_to_variant;

template<typename... Ts>
struct tuple_to_variant<std::tuple<Ts...>> {
    using type = std::variant<Ts...>;
};

} // namespace detail

// ── IVariantChannel ───────────────────────────────────────────────────────────
// Type-erased channel surface used by PyNetwork.
// The underlying FIFO stores raw T — variant conversion happens only at push/pop.

template<typename Variant>
class IVariantChannel {
public:
    virtual ~IVariantChannel() = default;
    virtual void            push(Variant v) = 0;
    // Lossless push with backpressure (waits instead of dropping when full).
    virtual void            push_blocking(Variant v) = 0;
    virtual Variant         pop()           = 0;
    virtual std::type_index type_index()    const = 0;
    virtual std::string     type_name()     const = 0;
    virtual void            enable()              = 0;
    virtual void            disable()             = 0;
};

// ── VariantChannel<T, Variant> ────────────────────────────────────────────────
// Shares a Channel<T> with the node that owns the input slot.
// push(): std::get<T> from Variant → raw T into the queue.
// pop():  raw T from queue → wrapped into Variant.

template<typename T, typename Variant>
class VariantChannel : public IVariantChannel<Variant> {
public:
    explicit VariantChannel(std::shared_ptr<Channel<T>> ch)
        : channel_(std::move(ch)) {}

    void push(Variant v) override {
        channel_->push(std::get<T>(std::move(v)));
    }
    void push_blocking(Variant v) override {
        channel_->push_blocking(std::get<T>(std::move(v)));
    }
    Variant pop() override {
        return Variant{ channel_->pop() };
    }
    std::type_index type_index() const override { return std::type_index(typeid(T)); }
    std::string     type_name()  const override { return typeid(T).name(); }
    void enable()  override { channel_->enable(); }
    void disable() override { channel_->disable(); }

    Channel<T>* raw_ptr() { return channel_.get(); }

private:
    std::shared_ptr<Channel<T>> channel_;
};

// ── IVariantNode ──────────────────────────────────────────────────────────────

template<typename Variant>
class IVariantNode : public INode,
                     public std::enable_shared_from_this<IVariantNode<Variant>> {
public:
    void set_name(std::string name) override { name_ = std::move(name); }
    const std::string& name() const { return name_; }

    virtual std::size_t     input_count()              const = 0;
    virtual std::size_t     output_count()             const = 0;
    virtual std::type_index input_type(std::size_t i)  const = 0;
    virtual std::type_index output_type(std::size_t i) const = 0;

    // Returns the IVariantChannel wrapping this node's input slot i.
    // PyNetwork calls this on the *destination* node to get the channel to wire upstream into.
    virtual std::shared_ptr<IVariantChannel<Variant>> input_channel(std::size_t i) = 0;

    // Wires this node's output slot i into ch (ch belongs to the downstream node's input).
    virtual void set_output_channel(std::size_t i,
                                    std::shared_ptr<IVariantChannel<Variant>> ch) = 0;
private:
    std::string name_;
};

// ── VariantNodeWrapper<Func, Variant, InTag, OutTag> ──────────────────────────
// Wraps a Node<Func,...> for use inside a PyNetwork.
//
// At construction: for each input port I, builds a shared_ptr<Channel<ArgI>> and
//   installs it in both the wrapped node (via set_input_channel) and a
//   VariantChannel<ArgI> adapter. PyNetwork passes the adapter to the upstream
//   node's set_output_channel so the upstream raw pointer points at the same Channel<T>.
//
// At connect (output side): downcasts the provided IVariantChannel to
//   VariantChannel<RetI>, then calls node_.set_output_channel with the raw ptr.

template<auto Func, typename Variant,
         typename InputTag  = in<>,
         typename OutputTag = out<>>
class VariantNodeWrapper;

template<auto Func, typename Variant,
         fixed_string... InNames, fixed_string... OutNames>
class VariantNodeWrapper<Func, Variant, in<InNames...>, out<OutNames...>>
    : public IVariantNode<Variant>
{
    using NodeT = Node<Func, in<InNames...>, out<OutNames...>>;

public:
    using args_tuple   = typename NodeT::args_tuple;
    using return_tuple = typename NodeT::return_tuple;

    static constexpr std::size_t n_in  = NodeT::input_count;
    static constexpr std::size_t n_out = NodeT::output_count;

    explicit VariantNodeWrapper(std::size_t fifo_capacity = 5)
        : node_(fifo_capacity)
        , in_channels_(n_in)
        , out_channels_(n_out)
        , out_type_indices_(n_out, std::type_index(typeid(void)))
    {
        init_inputs(std::make_index_sequence<n_in>{}, fifo_capacity);
        init_out_types(std::make_index_sequence<n_out>{});
    }

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

    void prepare() override { node_.prepare(); }
    void start() override { node_.start(); }
    void stop()  override { node_.stop(); }
    bool running() const override { return node_.running(); }
    const NodeStats& stats() const override { return node_.stats(); }
    void set_name(std::string name) override { node_.set_name(std::move(name)); }
    NodeSnapshot node_snapshot(const std::string& name, double elapsed_s) const override {
        return node_.node_snapshot(name, elapsed_s);
    }

    // ── IVariantNode ──────────────────────────────────────────────────────────

    std::size_t input_count()  const override { return n_in; }
    std::size_t output_count() const override { return n_out; }

    std::type_index input_type(std::size_t i) const override {
        return in_channels_[i]->type_index();
    }
    std::type_index output_type(std::size_t i) const override {
        return out_type_indices_[i];
    }

    std::shared_ptr<IVariantChannel<Variant>> input_channel(std::size_t i) override {
        return in_channels_[i];
    }

    void set_output_channel(std::size_t i,
                            std::shared_ptr<IVariantChannel<Variant>> ch) override {
        set_output_impl(i, std::move(ch), std::make_index_sequence<n_out>{});
    }

private:
    template<std::size_t... Is>
    void init_inputs(std::index_sequence<Is...>, std::size_t cap) {
        ((init_one_input<Is>(cap)), ...);
    }

    template<std::size_t I>
    void init_one_input(std::size_t cap) {
        using T = std::tuple_element_t<I, args_tuple>;
        auto shared_ch = std::make_shared<Channel<T>>(cap);
        // Install the shared Channel<T> in the node's input slot
        node_.template set_input_channel<I>(shared_ch);
        // Wrap it in a VariantChannel for PyNetwork to use
        in_channels_[I] = std::make_shared<VariantChannel<T, Variant>>(std::move(shared_ch));
    }

    template<std::size_t... Is>
    void init_out_types(std::index_sequence<Is...>) {
        ((out_type_indices_[Is] =
              std::type_index(typeid(std::tuple_element_t<Is, return_tuple>))), ...);
    }

    template<std::size_t... Is>
    void set_output_impl(std::size_t port,
                         std::shared_ptr<IVariantChannel<Variant>> ch,
                         std::index_sequence<Is...>) {
        bool matched = false;
        ((Is == port && (set_output_at<Is>(std::move(ch)), matched = true)), ...);
        if (!matched)
            throw std::out_of_range("set_output_channel: port index out of range");
    }

    template<std::size_t I>
    void set_output_at(std::shared_ptr<IVariantChannel<Variant>> ch) {
        using T = std::tuple_element_t<I, return_tuple>;
        auto* typed = dynamic_cast<VariantChannel<T, Variant>*>(ch.get());
        if (!typed)
            throw std::runtime_error(
                "set_output_channel: type mismatch at output port " + std::to_string(I));
        // The downstream VariantChannel holds a shared Channel<T>; point our output at it.
        // We need a raw ptr — extract it via a getter.
        node_.template set_output_channel<I>(typed->raw_ptr());
        out_channels_[I] = std::move(ch);
    }

    NodeT node_;
    std::vector<std::shared_ptr<IVariantChannel<Variant>>> in_channels_;
    std::vector<std::shared_ptr<IVariantChannel<Variant>>> out_channels_;
    std::vector<std::type_index>                           out_type_indices_;
};

// ── PythonConverter ───────────────────────────────────────────────────────────
// Specialise for each type you want to use across a PyNetwork.
// Required members: to_python(const T&), from_python(nb::object).
// Optional member:  static constexpr const char* type_name = "friendly_name";
// Built-in specialisations for int/float/double/bool/std::string are provided
// automatically when you include <kpn/python/auto_bind.hpp>.

template<typename T>
struct PythonConverter {};

} // namespace kpn