#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") {}
};

// Nothing available *right now* on a channel that is still open. Distinct from
// ChannelClosedError, which means upstream is finished and never coming back.
//
// Conflating the two is expensive in one direction only: a consumer that reads
// "empty" as "closed" stops a live node permanently, and because a stopping
// node disables its own inputs and outputs, one benign empty read takes the
// rest of the pipeline with it. The reverse costs nothing.
class ChannelEmptyError : public std::runtime_error {
public:
    ChannelEmptyError() : std::runtime_error("channel empty") {}
};

// ── 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_);
        }

        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();

        // Level-triggered, not edge-triggered — see set_push_callback.
        if (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_;
    }

    /// Outcome of a non-blocking push.
    ///
    /// try_push used to return bool, and returned *true* for a closed channel —
    /// so "delivered" and "discarded because nobody is listening" were the same
    /// answer. Both mean "stop trying", which is why the callers were correct,
    /// but neither they nor the producer's own accounting could tell a value
    /// that arrived from one that was thrown away. Only the channel's drop
    /// counter knew.
    enum class PushResult { Taken, Full, Closed };

    /// Non-blocking, lossless push. Returns Full when the ring is full, having
    /// changed nothing — the caller keeps the value and retries when woken.
    PushResult try_push(T& value) {
        if (!accepting_.load(std::memory_order_acquire)) {
            stats_.record_drop();
            return PushResult::Closed;
        }
        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 PushResult::Full;

        const std::size_t data_bytes = ChannelDataSize<T>::bytes(value);
        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();
        // Level-triggered, not edge-triggered — see set_push_callback.
        if (push_callback_) push_callback_();
        return PushResult::Taken;
    }

    // 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);
                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();
                // Level-triggered, not edge-triggered — see set_push_callback.
                if (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()).
    //
    // The slot holds exactly one undelivered token. A second offered before the
    // first is taken is refused, not queued and not overwritten: two control
    // tokens on one channel means the stream ended twice, which is a caller
    // protocol error rather than backpressure, and silently coalescing them
    // would hide it.
    /// Outcome of offering a sentinel. SlotBusy is a protocol error, not
    /// backpressure: it means a second control token was offered while the
    /// first was still undelivered, and a channel carries at most one.
    enum class SentinelResult { Taken, Closed, SlotBusy };

    /// Non-consuming form. `value` is left untouched unless the result is
    /// Taken, so a refused token is still the caller's to report.
    SentinelResult try_push_sentinel(T& value) {
        if (!accepting_.load(std::memory_order_acquire)) {
            stats_.record_drop();
            return SentinelResult::Closed;
        }
        // Refuse rather than overwrite. Overwriting lost the first token
        // silently, and worse, wrote eof_value_ while the consumer could be
        // moving the previous one out of it — a data race on the storage, which
        // for a shared_ptr payload is a torn refcount rather than a stale read.
        //
        // Checking here is what makes the slot a correct SPSC handshake: the
        // producer is the only writer of eof_value_ and the only one that sets
        // has_eof_, the consumer is the only one that clears it, so observing
        // false here means the consumer has finished with the storage and will
        // not touch it again until this store publishes the next token.
        if (has_eof_.load(std::memory_order_acquire)) {
            stats_.record_drop();
            return SentinelResult::SlotBusy;
        }
        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 SentinelResult::Taken;
    }

    /// Consuming convenience form. Returns false when the token was not stored,
    /// whether because the channel is closed or because one is already pending.
    bool push_sentinel(T value) {
        return try_push_sentinel(value) == SentinelResult::Taken;
    }

    // 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 after every successful push.
    //
    // It fires on every push, not on the empty→non-empty transition, and that
    // is a correctness requirement rather than a simplification.
    //
    // The edge version tested `was_empty = (t == h)` using an `h` sampled
    // *before* the item was published. A PoolNode consumer decides whether to
    // run again from the level (count_ready → approx_size), so the two sides
    // could each read the other as stale and both stand down:
    //
    //   producer (push)                    consumer (PoolNode firing)
    //   ------------------------           ----------------------------
    //   samples t=782, h=781
    //     -> was_empty = false, no wake
    //                                      pops idx 781, head_ = 782
    //                                      count_ready(): head_==tail_==782
    //                                      -> not ready, gate released to Idle
    //   tail_.store(783)
    //
    // The item is in the ring, the node is idle, and no wake is outstanding.
    // Worse, the failure is absorbing: every later push now sees a non-empty
    // ring, so `was_empty` is false forever and the callback never fires again.
    // The node sleeps while its backlog grows and its consumer waits on it.
    //
    // Re-reading head_ after the tail_ store does not fix it. That is the
    // store-buffer pattern, and under acquire/release both sides may legally
    // read stale; forbidding it needs seq_cst on the producer's tail_ store and
    // head_ load *and* on the consumer's head_ store and tail_ load — a fence
    // on both hot paths. Firing unconditionally is correct by construction:
    // the callback runs after the publishing store, so a consumer that observes
    // the level at all observes the item.
    //
    // The redundant wakes are cheap. on_input_ready re-checks the level, and
    // SubmitGate::claim() collapses a wake arriving during a firing into the
    // firing already in flight, so the cost is one CAS, not one extra run.
    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).
    // head_ is loaded first, deliberately. Both indices only ever increase, so
    // reading head_ before tail_ can at worst under-report a concurrent push;
    // the other order can read a head_ that has advanced past the tail_ already
    // sampled, and the unsigned difference then wraps to ~2^64. A caller
    // polling "is this channel empty yet" against that value never terminates.
    std::size_t size() const {
        const std::size_t h = head_.load(std::memory_order_relaxed);
        const std::size_t t = tail_.load(std::memory_order_acquire);
        return t - h;
    }

    // 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 {
        // head_ before tail_, for the reason given on size().
        const std::size_t h = head_.load(std::memory_order_relaxed);
        const std::size_t t = tail_.load(std::memory_order_acquire);
        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;
        // Re-check emptiness *after* observing has_eof_, not before.
        //
        // Callers check the ring is empty and then call this, but the producer
        // can push a value and publish the sentinel in the window between those
        // two steps — so the sentinel would be delivered with a real value still
        // queued behind it, breaking the "sentinel is strictly last" contract
        // that downstream teardown depends on. a0c4bf5 closed the variant where
        // the caller's emptiness check used a stale tail_ snapshot; this is the
        // one where the check is fresh but simply too early.
        //
        // Checking here is what makes it sound: the producer publishes the
        // sentinel with a release store *after* its ring pushes, so a consumer
        // that has observed has_eof_ has also observed every tail_ advance
        // before it. If the ring is non-empty now, those values genuinely
        // precede the sentinel and must be delivered first.
        if (head_.load(std::memory_order_relaxed)
                != tail_.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