#pragma once
#include "channel.hpp"
#include "diagnostics.hpp"
#include "inode.hpp"
#include "port.hpp"
#include "traits.hpp"

#include <array>
#include <atomic>
#include <chrono>
#include <iostream>
#include <memory>
#include <optional>
#include <thread>
#include <tuple>
#include <utility>

namespace kpn {

// ── FanoutNode ────────────────────────────────────────────────────────────────
//
// Reads one item from its single input channel and pushes a copy to each of
// N output channels. All N downstream nodes receive every item.
//
// Usage:
//   auto fan = make_fanout<Image, 2>(/*capacity=*/8);
//   net.connect("src",  src.output<0>(),   "fan",   fan.input<0>())
//      .connect("fan",  fan.output<0>(),   "nodeA", nodeA.input<0>())
//      .connect("fan",  fan.output<1>(),   "nodeB", nodeB.input<0>())

template<typename T, std::size_t N, std::size_t Id = 0>
class FanoutNode : public INode {
public:
    using args_tuple   = std::tuple<T>;
    using return_tuple = repeat_tuple_t<T, N>;
    using return_raw   = return_tuple;

    static constexpr std::size_t input_count    = 1;
    static constexpr std::size_t output_count   = N;
    static constexpr std::size_t unique_tag     = Id;
    static constexpr bool        is_fanout_node = true;

    explicit FanoutNode(std::size_t fifo_capacity = 5)
        : fifo_capacity_(fifo_capacity)
    {
        input_ch_ = std::make_shared<Channel<T>>(fifo_capacity);
    }

    ~FanoutNode() override { stop(); }

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

    void start() override {
        input_ch_->enable();
        stop_flag_.store(false, std::memory_order_relaxed);
        thread_ = std::jthread([this](std::stop_token) { run_loop(); });
    }

    void stop() override {
        stop_flag_.store(true, std::memory_order_relaxed);
        input_ch_->disable();
        if (thread_.joinable()) thread_.request_stop(), thread_.join();
    }

    bool running() const override {
        return thread_.joinable() && !stop_flag_.load(std::memory_order_relaxed);
    }

    void set_name(std::string name) override { name_ = std::move(name); }

    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,
                0.0,  // queue_wait_ms — fanout is not pool-scheduled
                stats_.total_exec_us.load(std::memory_order_relaxed) / 1000.0};
    }

    // ── Port access ───────────────────────────────────────────────────────────

    template<std::size_t I = 0>
    InputPort<FanoutNode, I> input() {
        static_assert(I == 0, "FanoutNode has exactly one input");
        return {*this};
    }

    template<std::size_t I>
    OutputPort<FanoutNode, I> output() {
        static_assert(I < N, "FanoutNode output index out of range");
        return {*this};
    }

    // ── Internal channel accessors (called by Network::connect) ───────────────

    template<std::size_t I>
    Channel<T>& input_channel() {
        static_assert(I == 0);
        return *input_ch_;
    }

    template<std::size_t I>
    void set_input_channel(std::shared_ptr<Channel<T>> ch) {
        static_assert(I == 0);
        input_ch_ = std::move(ch);
    }

    template<std::size_t I>
    void set_output_channel(Channel<T>* ch) {
        static_assert(I < N);
        out_channels_[I] = ch;
    }

private:
    // Deliver `val` to every connected output, losslessly.
    //
    // Previously a full output cost the value: push() threw and the exception was
    // swallowed per output. A dropped item does not degrade a downstream result,
    // it silently changes one, and the consumer cannot tell it happened — so the
    // fanout waits instead, and the producer upstream runs slower.
    //
    // Unlike a pool node, a fanout owns a private thread, so waiting here costs
    // no scheduler worker and needs no space-callback park; a bounded retry is
    // enough. `stop_flag_` is re-checked every pass so teardown cannot hang on a
    // full output regardless of the order the network stops its nodes in.
    //
    // Outputs are retried independently, so a full output never delays delivery
    // to one with room. Note what that does *not* buy: the next input is not
    // popped until every output has accepted the current item, so one branch can
    // never run ahead of another by more than the slower branch's buffering.
    //
    // **That bound is a precondition on any topology where the branches rejoin.**
    // If a consumer on branch B blocks waiting for something branch A computes,
    // B's buffering must exceed the lead A needs, or the two wedge — B waiting on
    // A, A starved because the fanout is holding an item B will not take. Making
    // the fanout lossless is what puts that precondition on the topology; while
    // it dropped, the question could not arise.
    //
    // `parked` receives the time spent waiting on a full output, which the caller
    // charges to blocked rather than exec.
    //
    // Returns false if stopped with the value undelivered.
    bool deliver(const T& val, duration_t& parked) {
        std::array<std::optional<T>, N> pending;
        std::size_t outstanding = 0;
        for (std::size_t i = 0; i < N; ++i)
            if (out_channels_[i]) { pending[i].emplace(val); ++outstanding; }

        bool first_pass = true;
        auto park_from  = clock_t::now();

        for (;;) {
            for (std::size_t i = 0; i < N; ++i) {
                if (!pending[i]) continue;
                // Taken or Closed both mean "stop trying" — delivered, or gone
                // with the drop recorded. Only Full is worth another pass.
                if (out_channels_[i]->try_push(*pending[i])
                        != Channel<T>::PushResult::Full) {
                    pending[i].reset();
                    --outstanding;
                }
            }
            if (first_pass) { park_from = clock_t::now(); first_pass = false; }

            if (outstanding == 0) {
                parked = duration_t(clock_t::now() - park_from);
                return true;
            }
            if (stop_flag_.load(std::memory_order_relaxed)) {
                // Teardown with work in hand. One last throwing push per
                // outstanding output, purely so the channel's own stats record
                // the loss (drop if it is disabled, overflow if it is merely
                // full). The whole point of the lossless path is that a loss is
                // never invisible, and a silent `return` here would reintroduce
                // exactly the hole this function exists to close.
                for (std::size_t i = 0; i < N; ++i) {
                    if (!pending[i]) continue;
                    try { out_channels_[i]->push(std::move(*pending[i])); }
                    catch (const ChannelOverflowError&) {}
                }
                parked = duration_t(clock_t::now() - park_from);
                return false;
            }
            std::this_thread::sleep_for(std::chrono::microseconds(50));
        }
    }

    void run_loop() {
        while (!stop_flag_.load(std::memory_order_relaxed)) {
            try {
                auto t0  = clock_t::now();
                T    val = input_ch_->pop();
                auto t1  = clock_t::now();
                auto cpu0 = NodeStats::cpu_now();

                duration_t parked{0};
                const bool delivered = deliver(val, parked);

                auto cpu1 = NodeStats::cpu_now();
                auto t2   = clock_t::now();
                // Time spent waiting on a full output is *blocked*, not exec: a
                // parked fanout is idle, and charging it to exec would report the
                // node as busy exactly when it is the one being held up.
                stats_.record_exec(duration_t(t2 - t1) - parked,
                                   duration_t(t1 - t0) + parked, cpu0, cpu1);
                if (!delivered) break;
            } catch (const ChannelClosedError&) {
                break;
            }
        }
    }

    std::string                 name_;
    std::size_t                 fifo_capacity_;
    std::shared_ptr<Channel<T>> input_ch_;
    std::array<Channel<T>*, N>  out_channels_{};
    std::atomic<bool>           stop_flag_{false};
    std::jthread                thread_;
    NodeStats                   stats_;
};

// ── Factory ───────────────────────────────────────────────────────────────────

template<typename T, std::size_t N>
FanoutNode<T, N> make_fanout(std::size_t fifo_capacity = 5) {
    return FanoutNode<T, N, 0>(fifo_capacity);
}

} // namespace kpn