#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_();
    }

    // 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) {
                if (!accepting_.load(std::memory_order_acquire))
                    throw ChannelClosedError{};

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

                if (h == t) {
                    // Still empty after spin — sleep until push() or disable() fires.
                    // Re-check tail after loading w to guard against a lost wakeup.
                    if (tail_.load(std::memory_order_acquire) != h) 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);
            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.
    bool try_pop_now(T& out) {
        const std::size_t h = head_.load(std::memory_order_relaxed);
        if (h == tail_.load(std::memory_order_acquire)) return false;
        out = extract(std::move(buf_[h & ring_mask_]));
        head_.store(h + 1, std::memory_order_release);
        stats_.record_pop();
        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);
    }

    // Size derived lazily from ring indices — no separate counter on the hot path.
    std::size_t size() const {
        return tail_.load(std::memory_order_relaxed)
             - head_.load(std::memory_order_relaxed);
    }
    std::size_t approx_size() const { return size(); }

    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;
    }

    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_;
    ChannelStats                    stats_;

    // Separate cache lines: head_ is written only by the consumer;
    // tail_ and wake_ are written only by the producer.
    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