KPN/include/kpn/channel.hpp

#pragma once
#include "diagnostics.hpp"
#include <atomic>
#include <chrono>
#include <cstdint>
#include <functional>
#include <memory>
#include <stdexcept>
#include <string>
#include <thread>
#include <type_traits>

namespace kpn {

// ── Data size trait ───────────────────────────────────────────────────────────
// Returns the number of bytes of logical payload carried by a value.
// Defaults to sizeof(T), which is correct for PODs and fixed-size types.
// Specialize for heap-owning types (e.g. cv::Mat) to get accurate bandwidth:
//
//   template<> struct kpn::ChannelDataSize<cv::Mat> {
//       static std::size_t bytes(const cv::Mat& m) { return m.total() * m.elemSize(); }
//   };

template<typename T>
struct ChannelDataSize {
    static std::size_t bytes(const T&) { return sizeof(T); }
};

// ── Storage policy ────────────────────────────────────────────────────────────

template<typename T>
struct channel_storage_policy {
    static constexpr bool by_value =
        std::is_trivially_copyable_v<T> && sizeof(T) <= 8;
};

template<typename T>
using channel_storage_t = std::conditional_t<
    channel_storage_policy<T>::by_value,
    T,
    std::shared_ptr<const T>
>;

// ── Exceptions ────────────────────────────────────────────────────────────────

class ChannelOverflowError : public std::runtime_error {
public:
    explicit ChannelOverflowError(std::size_t capacity)
        : std::runtime_error("channel overflow: capacity " + std::to_string(capacity) +
                             " exceeded") {}
    ChannelOverflowError(std::size_t capacity, std::string context)
        : std::runtime_error(std::move(context) + ": capacity " + std::to_string(capacity) +
                             " exceeded") {}
};

class ChannelClosedError : public std::runtime_error {
public:
    ChannelClosedError() : std::runtime_error("channel closed") {}
};

// ── CPU pause hint ────────────────────────────────────────────────────────────
// Signals the CPU that this is a spin-wait loop, improving HT sibling throughput
// and preventing branch-predictor thrash on x86. Falls back to a compiler barrier.

[[maybe_unused]] static void spin_hint() noexcept {
#if defined(__x86_64__) || defined(__i386__)
    __asm__ volatile("pause" ::: "memory");
#elif defined(__aarch64__) || defined(__arm__)
    __asm__ volatile("yield" ::: "memory");
#else
    std::atomic_signal_fence(std::memory_order_seq_cst);
#endif
}

// ── Channel ───────────────────────────────────────────────────────────────────
// SPSC ring buffer with atomic wait/notify and configurable spin-before-sleep.
//
// `spin_count` (constructor arg, default 200): number of pause-hint iterations
// before falling back to atomic::wait (futex). At ~20 ns/pause on x86 this is
// ~4 µs. Set to 0 to disable spinning (useful for power-constrained or
// predominantly-idle pipelines).
//
// Memory ordering contract (SPSC):
//   push(): tail_.store(release) pairs with pop()'s tail_.load(acquire)
//           head_.load(acquire)  pairs with pop()'s head_.store(release)
//   pop():  head_.store(release) pairs with push()'s head_.load(acquire)
//           tail_.load(acquire)  pairs with push()'s tail_.store(release)

template<typename T>
class Channel {
public:
    using storage_type = channel_storage_t<T>;

    explicit Channel(std::size_t capacity = 5, std::size_t spin_count = 200)
        : capacity_(capacity), spin_count_(spin_count)
    {
        std::size_t rs = 1;
        while (rs <= capacity) rs <<= 1;   // smallest power-of-2 > capacity
        ring_mask_ = rs - 1;
        buf_ = std::make_unique<storage_type[]>(rs);
    }

    Channel(const Channel&)            = delete;
    Channel& operator=(const Channel&) = delete;

    // Push a value.
    // - If channel is disabled (accepting_ == false): silently drop.
    // - If channel is full (fill >= capacity_): throw ChannelOverflowError.
    void push(T value) {
        if (!accepting_.load(std::memory_order_relaxed)) {
            stats_.record_drop();
            return;
        }
        const std::size_t data_bytes = ChannelDataSize<T>::bytes(value);
        const std::size_t t = tail_.load(std::memory_order_relaxed);
        const std::size_t h = head_.load(std::memory_order_acquire);

        if (!accepting_.load(std::memory_order_acquire)) {
            stats_.record_drop();
            return;
        }
        if (t - h >= capacity_) {
            stats_.record_overflow();
            throw ChannelOverflowError(capacity_);
        }

        const bool was_empty = (t == h);
        buf_[t & ring_mask_] = make_storage(std::move(value));
        tail_.store(t + 1, std::memory_order_release);
        stats_.record_push(t - h + 1, data_bytes);

        wake_.fetch_add(1, std::memory_order_release);
        wake_.notify_one();

        if (was_empty && push_callback_)
            push_callback_();
    }

    /// Called when a pop frees a slot in a previously-full ring.
    ///
    /// The mirror of `set_push_callback`, and it exists for the same reason:
    /// a producer must be able to *park* rather than spin. Without it the only
    /// lossless option is `push_blocking`, which sleeps inside the caller's
    /// thread — and when that thread is a scheduler worker, parking it starves
    /// every node pinned to it (see the hold-and-wait note on push_sentinel).
    void set_space_callback(std::function<void()> cb) { space_callback_ = std::move(cb); }

    /// True when a push would currently succeed. Used to close the lost-wakeup
    /// race: a producer that parks must re-check after clearing its queued flag,
    /// because a space_callback fired in between would otherwise be swallowed.
    bool has_space() const {
        return tail_.load(std::memory_order_relaxed) -
               head_.load(std::memory_order_acquire) < capacity_;
    }

    /// Non-blocking, lossless push. Returns false when the ring is full, having
    /// changed nothing — the caller keeps the value and retries when woken.
    bool try_push(T& value) {
        if (!accepting_.load(std::memory_order_acquire)) { stats_.record_drop(); return true; }
        const std::size_t t = tail_.load(std::memory_order_relaxed);
        const std::size_t h = head_.load(std::memory_order_acquire);
        if (t - h >= capacity_) return false;

        const std::size_t data_bytes = ChannelDataSize<T>::bytes(value);
        const bool was_empty = (t == h);
        buf_[t & ring_mask_] = make_storage(std::move(value));
        tail_.store(t + 1, std::memory_order_release);
        stats_.record_push(t - h + 1, data_bytes);
        wake_.fetch_add(1, std::memory_order_release);
        wake_.notify_one();
        if (was_empty && push_callback_) push_callback_();
        return true;
    }

    // Lossless push with BACKPRESSURE: if the ring is full, wait for the consumer to
    // drain instead of dropping (the throwing push()) — the producer just runs slower.
    // Use when every value must be delivered (e.g. replaying a dump for scoring, where
    // a dropped frame silently corrupts the result). SPSC: only the sole producer may
    // call it. Returns false if the channel was disabled while waiting.
    bool push_blocking(T value) {
        for (;;) {
            if (!accepting_.load(std::memory_order_acquire)) {
                stats_.record_drop();
                return false;
            }
            const std::size_t t = tail_.load(std::memory_order_relaxed);
            const std::size_t h = head_.load(std::memory_order_acquire);
            if (t - h < capacity_) {          // space available → normal push
                const std::size_t data_bytes = ChannelDataSize<T>::bytes(value);
                const bool was_empty = (t == h);
                buf_[t & ring_mask_] = make_storage(std::move(value));
                tail_.store(t + 1, std::memory_order_release);
                stats_.record_push(t - h + 1, data_bytes);
                wake_.fetch_add(1, std::memory_order_release);
                wake_.notify_one();
                if (was_empty && push_callback_) push_callback_();
                return true;
            }
            // full: yield briefly and retry (consumer will drain)
            std::this_thread::sleep_for(std::chrono::microseconds(50));
        }
    }

    // Lossless, non-blocking delivery for a must-deliver control token (EOF).
    //
    // A sentinel is stored out-of-band — in a dedicated slot that does NOT
    // consume ring capacity — so this can never overflow and never blocks the
    // caller. That distinction is essential: each KPN node has a single worker
    // thread, so a *blocking* push would park that thread and stop it draining
    // its own input, cascading into a hold-and-wait deadlock under backpressure.
    // Setting a flag and returning keeps the worker free to keep popping.
    //
    // The consumer's pop() drains the ring first, then delivers this sentinel,
    // preserving ordering (EOF arrives after all data pushed before it).
    //
    // Only the sole producer may call it (SPSC contract, same as push()).
    // Returns false if the channel is already disabled (token discarded —
    // teardown is in progress, so the sentinel is moot).
    bool push_sentinel(T value) {
        if (!accepting_.load(std::memory_order_acquire)) {
            stats_.record_drop();
            return false;
        }
        eof_value_ = make_storage(std::move(value));
        has_eof_.store(true, std::memory_order_release);
        // Wake a consumer blocked in pop(): the sentinel is now deliverable even
        // though the ring may be empty.
        wake_.fetch_add(1, std::memory_order_release);
        wake_.notify_one();
        if (push_callback_) push_callback_();
        return true;
    }

    // Blocking pop. Returns when an item is available.
    // Throws ChannelClosedError if the channel is disabled (regardless of fill).
    T pop() {
        for (;;) {
            // Snapshot wake_ BEFORE reading tail_ to prevent lost wakeups.
            const uint32_t    w = wake_.load(std::memory_order_relaxed);
            const std::size_t h = head_.load(std::memory_order_relaxed);
                  std::size_t t = tail_.load(std::memory_order_acquire);

            // If empty, spin before sleeping: avoids the futex when the next item
            // arrives within the spin window (~4 µs at default spin_count=200 on x86).
            if (h == t) {
                // Ring drained — deliver any pending out-of-band sentinel (EOF)
                // now, so it always arrives after the data pushed before it.
                //
                // Re-confirm emptiness against a fresh tail_ first: the snapshot
                // at the top of the loop may be stale (the producer can push more
                // values *and* the sentinel in the window since), and the sentinel
                // must never jump ahead of ring values pushed before it. The spin
                // and post-spin takes below already reload tail_ on the line above
                // them; this is the one take that used the loop-top snapshot.
                if (h == tail_.load(std::memory_order_acquire)) {
                    T s; if (take_sentinel(s)) return s;
                }

                if (!accepting_.load(std::memory_order_acquire))
                    throw ChannelClosedError{};

                for (std::size_t si = 0; si < spin_count_; ++si) {
                    spin_hint();
                    t = tail_.load(std::memory_order_acquire);
                    if (t != h) break;
                    { T s; if (take_sentinel(s)) return s; }
                    if (!accepting_.load(std::memory_order_relaxed))
                        throw ChannelClosedError{};
                }

                if (h == t) {
                    // Still empty after spin — sleep until push()/push_sentinel()
                    // or disable() fires. Re-check tail and the sentinel after
                    // loading w to guard against a lost wakeup.
                    if (tail_.load(std::memory_order_acquire) != h) continue;
                    if (has_eof_.load(std::memory_order_acquire)) continue;
                    wake_.wait(w, std::memory_order_relaxed);
                    continue;
                }
            }

            // Item available (found immediately or during spin).
            if (!accepting_.load(std::memory_order_acquire))
                throw ChannelClosedError{};
            T value = extract(std::move(buf_[h & ring_mask_]));
            head_.store(h + 1, std::memory_order_release);
            // A slot just freed: wake any producer parked on this channel.
            if (t - h >= capacity_ && space_callback_) space_callback_();
            stats_.record_pop();
            return value;
        }
    }

    // Non-blocking pop with timeout. For watchdog/display use only.
    bool try_pop(T& out, std::chrono::milliseconds timeout) {
        const auto deadline = std::chrono::steady_clock::now() + timeout;
        for (;;) {
            if (try_pop_now(out)) return true;
            if (!accepting_.load(std::memory_order_relaxed)) return false;
            if (std::chrono::steady_clock::now() >= deadline) return false;
            std::this_thread::sleep_for(std::chrono::microseconds(50));
        }
    }

    // Immediate non-blocking pop. Returns false if the ring is empty.
    // Once the ring is drained, delivers any pending out-of-band sentinel (EOF)
    // so pool nodes — which pop only via this path — still receive the token.
    bool try_pop_now(T& out) {
        const std::size_t h = head_.load(std::memory_order_relaxed);
        const std::size_t t = tail_.load(std::memory_order_acquire);
        if (h == t)
            return take_sentinel(out);
        out = extract(std::move(buf_[h & ring_mask_]));
        head_.store(h + 1, std::memory_order_release);
        stats_.record_pop();
        // Pool nodes pop only through here, so this is where a parked producer
        // gets woken: the ring was full, and it no longer is.
        if (t - h >= capacity_ && space_callback_) space_callback_();
        return true;
    }

    // Enable the channel (called by consumer node on start()).
    void enable() {
        accepting_.store(true, std::memory_order_relaxed);
    }

    // Disable the channel: stop accepting new pushes, unblock any waiting pop().
    // Items already in the ring are abandoned and freed when the Channel is destroyed.
    void disable() {
        accepting_.store(false, std::memory_order_release);
        wake_.fetch_add(1, std::memory_order_release);
        wake_.notify_all();
    }

    // Register a callback fired when the queue transitions empty→non-empty.
    void set_push_callback(std::function<void()> cb) {
        push_callback_ = std::move(cb);
    }

    // Ring occupancy, derived lazily from indices — no separate counter on the
    // hot path. Excludes any out-of-band sentinel (that lives outside the ring).
    std::size_t size() const {
        return tail_.load(std::memory_order_relaxed)
             - head_.load(std::memory_order_relaxed);
    }

    // A pending out-of-band sentinel (EOF) counts as consumable work here even
    // though it holds no ring slot. This is what node readiness checks call, so
    // a channel carrying only a sentinel still schedules its consumer's next
    // fire — without this the sentinel would never be popped and the pipeline
    // would deadlock at teardown.
    std::size_t approx_size() const {
        return size() + (has_eof_.load(std::memory_order_acquire) ? 1u : 0u);
    }

    std::size_t capacity()     const { return capacity_; }
    bool        is_accepting() const { return accepting_.load(std::memory_order_relaxed); }
    const ChannelStats& stats() const { return stats_; }

    ChannelSnapshot snapshot(const std::string& name) const {
        const std::size_t t = tail_.load(std::memory_order_relaxed);
        const std::size_t h = head_.load(std::memory_order_relaxed);
        return {
            name,
            capacity_,
            t - h,
            stats_.peak_fill.load(std::memory_order_relaxed),
            stats_.pushes.load(std::memory_order_relaxed),
            stats_.bytes_pushed.load(std::memory_order_relaxed),
            stats_.drops.load(std::memory_order_relaxed),
            stats_.overflows.load(std::memory_order_relaxed),
            stats_.pops.load(std::memory_order_relaxed),
            sizeof(T),
        };
    }

private:
    static storage_type make_storage(T&& v) {
        if constexpr (channel_storage_policy<T>::by_value)
            return std::move(v);
        else
            return std::make_shared<const T>(std::move(v));
    }

    static T extract(storage_type&& s) {
        if constexpr (channel_storage_policy<T>::by_value)
            return std::move(s);
        else
            return *s;
    }

    // Consume the out-of-band sentinel if one is pending. Consumer-only.
    // Called only when the ring is observed empty, so the sentinel is always
    // delivered after every value pushed before it.
    bool take_sentinel(T& out) {
        if (!has_eof_.load(std::memory_order_acquire)) return false;
        out = extract(std::move(eof_value_));
        has_eof_.store(false, std::memory_order_release);
        stats_.record_pop();
        return true;
    }

    const std::size_t               capacity_;
    const std::size_t               spin_count_;
    std::size_t                     ring_mask_;
    std::unique_ptr<storage_type[]> buf_;
    std::function<void()>           push_callback_;
    std::function<void()> space_callback_;
    ChannelStats                    stats_;

    // Out-of-band sentinel (EOF): stored outside the ring so its delivery never
    // depends on ring capacity and never blocks the producer. Written by the
    // producer (push_sentinel), read+cleared by the consumer (take_sentinel);
    // has_eof_ is the publish/consume handshake.
    storage_type              eof_value_{};
    std::atomic<bool>         has_eof_{false};

    // Separate cache lines: head_ is written only by the consumer;
    // tail_ and wake_ are written only by the producer.
    //   wake_  wakes a blocked pop() on enqueue or on a pending sentinel.
    alignas(64) std::atomic<std::size_t> head_{0};
    alignas(64) std::atomic<std::size_t> tail_{0};
                std::atomic<uint32_t>    wake_{0};
                std::atomic<bool>        accepting_{true};
};

// ── Channel probe — type-erased snapshot accessor ─────────────────────────────
// Used by both Network and StaticNetwork for diagnostics.

struct IChannelProbe {
    virtual ~IChannelProbe() = default;
    virtual ChannelSnapshot snapshot() const = 0;
};

template<typename T>
struct ChannelProbe : IChannelProbe {
    const Channel<T>& ch;
    std::string       name;
    ChannelProbe(const Channel<T>& c, std::string n) : ch(c), name(std::move(n)) {}
    ChannelSnapshot snapshot() const override { return ch.snapshot(name); }
};

} // namespace kpn