KPN/include/kpn/main_thread_node.hpp

#pragma once
#include "channel.hpp"
#include "diagnostics.hpp"
#include "fixed_string.hpp"
#include "node.hpp"       // INode
#include "port.hpp"

#include <atomic>
#include <chrono>
#include <cstddef>
#include <memory>
#include <optional>
#include <tuple>

namespace kpn {

// ── MainThreadNode ────────────────────────────────────────────────────────────
//
// Base class for nodes that must run on the main thread (e.g. OpenCV display
// on Wayland/Qt). Registered as a normal INode in the Network so it appears
// in diagnostics and the web UI, but spawns no thread.
//
// Usage:
//   class MyDisplay : public kpn::MainThreadNode<MyDisplay, in<"a","b">, TypeA, TypeB> {
//   public:
//       MyDisplay(...) { /* constructor runs on main thread */ }
//       bool operator()(TypeA a, TypeB b) { ...; return true; /* false = stop */ }
//   };
//
//   MyDisplay disp(...);
//   net.add("display", disp).connect(...).build();
//   net.start();
//   while (disp.step()) ;   // drives the event loop on the main thread
//   net.stop();
//
// step() behaviour:
//   - try_pop on every input channel with zero timeout
//   - if all inputs have data: calls operator(), records stats, returns its result
//   - if any input is missing: returns true immediately (caller should yield/waitKey)

template<typename Derived, typename InputTag, typename... Args>
class MainThreadNode;

template<typename Derived, fixed_string... InNames, typename... Args>
class MainThreadNode<Derived, in<InNames...>, Args...> : public INode {
public:
    static constexpr std::size_t input_count = sizeof...(Args);

    static_assert(
        sizeof...(InNames) == 0 || sizeof...(InNames) == input_count,
        "MainThreadNode: name count must match input type count, or provide none"
    );

    using args_tuple = std::tuple<Args...>;  // required by Network::connect type check

    explicit MainThreadNode(std::size_t fifo_capacity = 8) {
        init_channels(std::make_index_sequence<input_count>{}, fifo_capacity);
    }

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

    void start() override {
        enable_channels(std::make_index_sequence<input_count>{});
        running_.store(true, std::memory_order_relaxed);
    }

    void stop() override {
        running_.store(false, std::memory_order_relaxed);
        disable_channels(std::make_index_sequence<input_count>{});
    }

    bool running() const override {
        return running_.load(std::memory_order_relaxed);
    }

    void set_name(std::string) override {}

    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 (for Network::connect) ───────────────────────────────────

    template<std::size_t I>
    Channel<std::tuple_element_t<I, std::tuple<Args...>>>& input_channel() {
        return *std::get<I>(channels_);
    }

    template<std::size_t I>
    InputPort<Derived, I> input() {
        static_assert(I < input_count, "input index out of range");
        return {static_cast<Derived&>(*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>();
    }

    // ── Main-thread driver ────────────────────────────────────────────────────

    // Call this in a loop on the main thread instead of net.start()'s thread.
    // Returns false when operator() returns false or all channels are closed.
    bool step() {
        if (!running_.load(std::memory_order_relaxed)) return false;

        auto t0 = clock_t::now();
        auto inputs = try_pop_all(std::make_index_sequence<input_count>{});
        auto t1 = clock_t::now();

        if (!inputs.has_value()) return true; // not all inputs ready — yield

        auto cpu0 = NodeStats::cpu_now();
        bool cont = std::apply(
            [this](Args&&... a) {
                return static_cast<Derived*>(this)->operator()(std::forward<Args>(a)...);
            },
            std::move(*inputs));
        auto cpu1 = NodeStats::cpu_now();
        auto t2 = clock_t::now();

        stats_.record_exec(duration_t(t2 - t1), duration_t(t1 - t0), cpu0, cpu1);
        return cont;
    }

private:
    template<std::size_t... Is>
    void init_channels(std::index_sequence<Is...>, std::size_t cap) {
        ((std::get<Is>(channels_) =
              std::make_unique<Channel<std::tuple_element_t<Is, args_tuple>>>(cap)), ...);
    }

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

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

    // Try to pop one item from every channel with zero timeout.
    // Returns nullopt if any channel has no data ready.
    template<std::size_t... Is>
    std::optional<args_tuple> try_pop_all(std::index_sequence<Is...>) {
        args_tuple result;
        bool all_ready = true;
        // Use a fold that short-circuits on first missing item
        ((all_ready = all_ready &&
            std::get<Is>(channels_)->try_pop(
                std::get<Is>(result), std::chrono::milliseconds(0))), ...);
        if (!all_ready) return std::nullopt;
        return result;
    }

    // Build the channel tuple type
    template<typename Tup, std::size_t... Is>
    static auto make_channel_tuple(std::index_sequence<Is...>)
        -> std::tuple<std::unique_ptr<Channel<std::tuple_element_t<Is, Tup>>>...>;

    using channels_t = decltype(make_channel_tuple<args_tuple>(
        std::make_index_sequence<input_count>{}));

    channels_t           channels_;
    std::atomic<bool>    running_{false};
    NodeStats            stats_;
};

} // namespace kpn