#pragma once
// Nanobind binding helpers for KPN++ Python interface.
// Included only by python/kpn_python.cpp — do not include from core headers.

#include "../variant_node.hpp"
#include "../network.hpp"

#ifdef KPN_BUILD_PYTHON
#include <nanobind/nanobind.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>

#include <functional>
#include <map>
#include <stdexcept>
#include <string>
#include <typeindex>
#include <vector>

namespace kpn::python {

namespace nb = nanobind;

// ── PyNetwork<Variant> ────────────────────────────────────────────────────────
// Runtime graph builder for Python. Holds IVariantNode instances and connects
// them via IVariantChannel adapters. The variant only lives at the boundary;
// each node's internal Channel<T> stores raw T values.

template<typename Variant>
class PyNode; // forward declaration

template<typename Variant>
class PyNetwork {
public:
    using VNode    = IVariantNode<Variant>;
    using VChannel = IVariantChannel<Variant>;

    // ── GC support ────────────────────────────────────────────────────────────
    // Visit every Python object this network transitively holds (currently the
    // callable of each PyNode). Used by the Network type's tp_traverse slot so
    // Python's cyclic GC can discover instance → callable → globals() cycles.
    // Defined out-of-line below, once PyNode is a complete type.
    template<typename Fn>
    void visit_python_objects(Fn&& visit) const;
    // Drop all Python references held by nodes, breaking any cycle (tp_clear).
    void clear_python_objects();

    // ── Builder API ───────────────────────────────────────────────────────────

    void add(std::string name, std::shared_ptr<VNode> node) {
        if (nodes_.count(name))
            throw std::runtime_error("duplicate node name: " + name);
        node->set_name(name);
        nodes_.emplace(name, std::move(node));
        adj_[name];
    }

    // connect(src_name, out_idx, dst_name, in_idx)
    void connect(const std::string& src_name, std::size_t out_idx,
                 const std::string& dst_name, std::size_t in_idx)
    {
        auto& src = node_at(src_name);
        auto& dst = node_at(dst_name);

        if (out_idx >= src.output_count())
            throw std::out_of_range(src_name + ": output index " +
                                    std::to_string(out_idx) + " out of range");
        if (in_idx >= dst.input_count())
            throw std::out_of_range(dst_name + ": input index " +
                                    std::to_string(in_idx) + " out of range");

        if (src.output_type(out_idx) != dst.input_type(in_idx))
            throw std::runtime_error(
                "type mismatch: " + src_name + ".output[" + std::to_string(out_idx) +
                "] (" + src.output_type(out_idx).name() + ") → " +
                dst_name + ".input[" + std::to_string(in_idx) +
                "] (" + dst.input_type(in_idx).name() + ")");

        auto ch = dst.input_channel(in_idx);
        src.set_output_channel(out_idx, std::move(ch));
        adj_[src_name].push_back(dst_name);
    }

    void build() {
        topo_.clear();
        std::map<std::string, int> color;
        for (auto& [name, _] : nodes_)
            if (color[name] == 0) dfs(name, color);
    }

    // ── Lifecycle ─────────────────────────────────────────────────────────────

    void start() {
        for (auto& name : topo_)
            nodes_.at(name)->start();
    }

    void stop() {
        for (auto it = topo_.rbegin(); it != topo_.rend(); ++it)
            nodes_.at(*it)->stop();
    }

    // ── Python tap/inject ─────────────────────────────────────────────────────

    nb::object read(const std::string& node_name, std::size_t out_idx) {
        auto key = tap_key(node_name, out_idx);
        if (!taps_.count(key)) {
            auto& src = node_at(node_name);
            if (out_idx >= src.output_count())
                throw std::out_of_range(node_name + ": output index out of range");
            auto tap = make_tap_channel(src.output_type(out_idx));
            src.set_output_channel(out_idx, tap);
            taps_[key] = std::move(tap);
        }
        Variant v;
        {
            nb::gil_scoped_release release;
            v = taps_.at(key)->pop();
        }
        return variant_to_python(std::move(v));
    }

    void write(const std::string& node_name, std::size_t in_idx, nb::object value) {
        auto& dst = node_at(node_name);
        if (in_idx >= dst.input_count())
            throw std::out_of_range(node_name + ": input index out of range");

        auto ch = dst.input_channel(in_idx);
        Variant v = python_to_variant(ch->type_index(), std::move(value));
        {
            nb::gil_scoped_release release;
            ch->push(std::move(v));
        }
    }

    // ── Python-callable node creation ─────────────────────────────────────────
    // Creates a PyNode wrapping a Python callable and adds it to the graph.
    // Type names must have been registered via register_full_type<T>().

    void add_node_python(std::string name, nb::object callable,
                         std::vector<std::string> in_names,
                         std::vector<std::string> out_names,
                         std::size_t capacity = 5)
    {
        std::vector<std::type_index> in_types, out_types;
        for (auto& s : in_names)  in_types.push_back(resolve_type_name(s));
        for (auto& s : out_names) out_types.push_back(resolve_type_name(s));

        add(std::move(name),
            std::make_shared<PyNode<Variant>>(
                std::move(callable),
                std::move(in_types),
                std::move(out_types),
                to_python_,
                from_python_,
                ch_factories_,
                capacity));
    }

    // ── Type converter registration ───────────────────────────────────────────

    template<typename T>
    void register_type(
        std::function<nb::object(const T&)> to_py,
        std::function<T(nb::object)>        from_py)
    {
        auto idx = std::type_index(typeid(T));
        to_python_[idx]   = [to_py](const Variant& v) { return to_py(std::get<T>(v)); };
        from_python_[idx] = [from_py](nb::object o) -> Variant {
            return Variant{ from_py(std::move(o)) };
        };
    }

    // register_channel_factory<T>: registers factory for creating input channels.
    template<typename T>
    void register_channel_factory() {
        ch_factories_[std::type_index(typeid(T))] =
            [](std::size_t cap) -> std::shared_ptr<VChannel> {
                return std::make_shared<VariantChannel<T, Variant>>(
                    std::make_shared<Channel<T>>(cap));
            };
    }

    // Backward-compatible alias.
    template<typename T>
    void register_tap_factory(std::size_t = 5) {
        register_channel_factory<T>();
    }

    // register_full_type<T>: registers converters + channel factory + type name.
    // This is what auto_bind.hpp calls; manual bindings can call register_type +
    // register_tap_factory separately for backward compatibility.
    template<typename T>
    void register_full_type(
        std::function<nb::object(const T&)> to_py,
        std::function<T(nb::object)>        from_py,
        const char* friendly_name = nullptr)
    {
        register_type<T>(std::move(to_py), std::move(from_py));
        register_channel_factory<T>();
        auto idx = std::type_index(typeid(T));
        type_names_.insert_or_assign(typeid(T).name(), idx);
        if (friendly_name) type_names_.insert_or_assign(friendly_name, idx);
    }

    // ── Type name lookup ──────────────────────────────────────────────────────

    void register_type_name(const std::string& name, std::type_index idx) {
        type_names_.insert_or_assign(name, idx);
    }

    std::type_index resolve_type_name(const std::string& name) const {
        auto it = type_names_.find(name);
        if (it == type_names_.end())
            throw std::runtime_error(
                "type '" + name + "' not registered — call register_full_type<T>() first");
        return it->second;
    }

private:
    VNode& node_at(const std::string& name) {
        auto it = nodes_.find(name);
        if (it == nodes_.end())
            throw std::runtime_error("unknown node: " + name);
        return *it->second;
    }

    void dfs(const std::string& name, std::map<std::string, int>& color) {
        color[name] = 1;
        for (auto& nbr : adj_[name]) {
            if (color[nbr] == 1)
                throw std::runtime_error("cycle detected in graph");
            if (color[nbr] == 0) dfs(nbr, color);
        }
        color[name] = 2;
        topo_.insert(topo_.begin(), name);
    }

    std::string tap_key(const std::string& node, std::size_t idx) {
        return node + ":" + std::to_string(idx);
    }

    std::shared_ptr<VChannel> make_tap_channel(std::type_index type,
                                                std::size_t cap = 5) {
        auto it = ch_factories_.find(type);
        if (it == ch_factories_.end())
            throw std::runtime_error(
                "no channel factory for type: " + std::string(type.name()) +
                " — call register_full_type<T>() or register_tap_factory<T>()");
        return it->second(cap);
    }

    nb::object variant_to_python(Variant v) {
        auto idx = std::visit([](auto& x) {
            return std::type_index(typeid(x));
        }, v);
        auto it = to_python_.find(idx);
        if (it == to_python_.end())
            throw std::runtime_error("no to_python converter for type");
        return it->second(v);
    }

    Variant python_to_variant(std::type_index idx, nb::object obj) {
        auto it = from_python_.find(idx);
        if (it == from_python_.end())
            throw std::runtime_error("no from_python converter for type");
        return it->second(std::move(obj));
    }

    std::map<std::string, std::shared_ptr<VNode>>    nodes_;
    std::map<std::string, std::vector<std::string>>  adj_;
    std::vector<std::string>                         topo_;
    std::map<std::string, std::shared_ptr<VChannel>> taps_;

    std::map<std::type_index, std::function<nb::object(const Variant&)>> to_python_;
    std::map<std::type_index, std::function<Variant(nb::object)>>        from_python_;

    // Channel factory: type → function(capacity) → VChannel.
    // Used both for tap channels (read()) and PyNode input channel creation.
    std::map<std::type_index,
             std::function<std::shared_ptr<VChannel>(std::size_t)>> ch_factories_;

    // Friendly name → type_index (e.g. "int" → typeid(int)).
    std::map<std::string, std::type_index> type_names_;
};

// ── PyNode<Variant> ───────────────────────────────────────────────────────────
// A pure-Python processing node. Holds a nanobind callable.
// run_loop: pop inputs (release GIL), call Python (acquire GIL), push outputs.

template<typename Variant>
class PyNode : public IVariantNode<Variant> {
public:
    using VChannel = IVariantChannel<Variant>;

    using ChannelFactory =
        std::function<std::shared_ptr<VChannel>(std::size_t capacity)>;

    PyNode(nb::object callable,
           std::vector<std::type_index> in_types,
           std::vector<std::type_index> out_types,
           std::map<std::type_index, std::function<nb::object(const Variant&)>> to_py,
           std::map<std::type_index, std::function<Variant(nb::object)>>        from_py,
           std::map<std::type_index, ChannelFactory>                            ch_factories,
           std::size_t capacity = 5)
        : callable_(std::move(callable))
        , in_types_(std::move(in_types))
        , out_types_(std::move(out_types))
        , to_python_(std::move(to_py))
        , from_python_(std::move(from_py))
        , ch_factories_(std::move(ch_factories))
        , in_channels_(in_types_.size())
        , out_channels_(out_types_.size())
    {
        for (std::size_t i = 0; i < in_types_.size(); ++i) {
            auto it = ch_factories_.find(in_types_[i]);
            if (it == ch_factories_.end())
                throw std::runtime_error("PyNode: no channel factory for input type");
            in_channels_[i] = it->second(capacity);
        }
    }

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

    void start() override {
        for (auto& ch : in_channels_) 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);
        for (auto& ch : in_channels_) ch->disable();
        if (thread_.joinable()) {
            thread_.request_stop();
            nb::gil_scoped_release release;
            thread_.join();
        }
    }

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

    void set_name(std::string name) override {
        IVariantNode<Variant>::set_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 blk_ms   = stats_.total_blocked_us.load(std::memory_order_relaxed) / 1000.0;
        double total_ms = exec_ms + blk_ms;
        return { name, frames, exec_ms,
                 stats_.max_exec_us.load(std::memory_order_relaxed) / 1000.0,
                 blk_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 };
    }

    // ── IVariantNode ──────────────────────────────────────────────────────────

    std::size_t input_count()  const override { return in_types_.size(); }
    std::size_t output_count() const override { return out_types_.size(); }

    std::type_index input_type(std::size_t i)  const override { return in_types_[i]; }
    std::type_index output_type(std::size_t i) const override { return out_types_[i]; }

    std::shared_ptr<VChannel> input_channel(std::size_t i) override {
        return in_channels_[i];
    }

    void set_output_channel(std::size_t i,
                            std::shared_ptr<VChannel> ch) override {
        out_channels_[i] = std::move(ch);
    }

    // ── GC support (tp_traverse / tp_clear on the owning Network) ──────────────
    // The node holds a Python callable, which typically forms an
    // instance → callable → globals() → instance cycle. Expose the callable so
    // the Network's GC slots can traverse and clear it. See bindings.hpp's
    // network_tp_traverse/network_tp_clear.
    const nb::object& python_callable() const { return callable_; }
    void clear_python_callable() { callable_ = nb::object(); }

private:
    void run_loop() {
        while (!stop_flag_.load(std::memory_order_relaxed)) {
            try {
                auto t0 = clock_t::now();

                std::vector<Variant> inputs(in_channels_.size());
                for (std::size_t i = 0; i < in_channels_.size(); ++i)
                    inputs[i] = in_channels_[i]->pop();

                auto t1   = clock_t::now();
                auto cpu0 = NodeStats::cpu_now();

                std::vector<Variant> outputs;
                {
                    nb::gil_scoped_acquire acquire;
                    nb::list py_args;
                    for (auto& v : inputs)
                        py_args.append(variant_to_python(v));

                    nb::object result = callable_(*py_args);

                    if (out_channels_.size() == 1) {
                        outputs.push_back(python_to_variant(out_types_[0], result));
                    } else {
                        nb::tuple tup = nb::cast<nb::tuple>(result);
                        for (std::size_t i = 0; i < out_channels_.size(); ++i)
                            outputs.push_back(python_to_variant(out_types_[i], tup[i]));
                    }
                }

                auto cpu1 = NodeStats::cpu_now();
                auto t2   = clock_t::now();
                stats_.record_exec(duration_t(t2 - t1), duration_t(t1 - t0), cpu0, cpu1);

                for (std::size_t i = 0; i < out_channels_.size(); ++i) {
                    if (out_channels_[i])
                        out_channels_[i]->push(std::move(outputs[i]));
                }

            } catch (const ChannelClosedError&) {
                break;
            } catch (const ChannelOverflowError&) {
                // drop and continue
            }
        }
    }

    nb::object variant_to_python(const Variant& v) {
        auto idx = std::visit([](const auto& x) {
            return std::type_index(typeid(x));
        }, v);
        return to_python_.at(idx)(v);
    }

    Variant python_to_variant(std::type_index idx, nb::object obj) {
        return from_python_.at(idx)(std::move(obj));
    }

    nb::object callable_;
    std::vector<std::type_index>                                          in_types_;
    std::vector<std::type_index>                                          out_types_;
    std::map<std::type_index, std::function<nb::object(const Variant&)>>  to_python_;
    std::map<std::type_index, std::function<Variant(nb::object)>>         from_python_;
    std::map<std::type_index, ChannelFactory>                             ch_factories_;
    std::vector<std::shared_ptr<VChannel>>                                in_channels_;
    std::vector<std::shared_ptr<VChannel>>                                out_channels_;
    std::atomic<bool>  stop_flag_{false};
    std::jthread       thread_;
    NodeStats          stats_;
};

// ── PyNetwork GC helpers (defined here: PyNode is now complete) ────────────────

template<typename Variant>
template<typename Fn>
void PyNetwork<Variant>::visit_python_objects(Fn&& visit) const {
    for (const auto& [name, node] : nodes_)
        if (auto* py = dynamic_cast<const PyNode<Variant>*>(node.get()))
            visit(py->python_callable());
}

template<typename Variant>
void PyNetwork<Variant>::clear_python_objects() {
    for (auto& [name, node] : nodes_)
        if (auto* py = dynamic_cast<PyNode<Variant>*>(node.get()))
            py->clear_python_callable();
}

// ── GC type slots for the Network binding ─────────────────────────────────────
// The Network holds Python callables (via PyNode), forming uncollectable
// instance → callable → globals() → instance cycles at interpreter shutdown.
// These slots let Python's cyclic collector traverse and break them, silencing
// nanobind's leak warnings. See the nanobind "Reference leaks" documentation.

template<typename Variant>
int network_tp_traverse(PyObject* self, visitproc visit, void* arg) {
    Py_VISIT(Py_TYPE(self));
    if (!nb::inst_ready(self))
        return 0;
    auto* net = nb::inst_ptr<PyNetwork<Variant>>(self);
    int rv = 0;
    net->visit_python_objects([&](const nb::object& obj) {
        if (rv == 0 && obj.is_valid())
            rv = visit(obj.ptr(), arg);
    });
    return rv;
}

template<typename Variant>
int network_tp_clear(PyObject* self) {
    auto* net = nb::inst_ptr<PyNetwork<Variant>>(self);
    net->clear_python_objects();
    return 0;
}

template<typename Variant>
PyType_Slot* network_type_slots() {
    static PyType_Slot slots[] = {
        { Py_tp_traverse, reinterpret_cast<void*>(&network_tp_traverse<Variant>) },
        { Py_tp_clear,    reinterpret_cast<void*>(&network_tp_clear<Variant>) },
        { 0, nullptr }
    };
    return slots;
}

// ── register_py_network (legacy helper) ───────────────────────────────────────
// Registers PyNetwork<Variant> with the given nanobind module.
// Prefer bind_network<Registry> from auto_bind.hpp for new code.

template<typename Variant>
void register_py_network(nb::module_& m, const char* class_name = "Network") {
    using Net = PyNetwork<Variant>;

    nb::class_<Net>(m, class_name, nb::type_slots(network_type_slots<Variant>()))
        .def(nb::init<>())
        .def("connect", &Net::connect,
             nb::arg("src"), nb::arg("out_idx"),
             nb::arg("dst"), nb::arg("in_idx"))
        .def("build",  &Net::build)
        .def("start",  &Net::start)
        .def("stop",   &Net::stop)
        .def("read",   &Net::read,
             nb::arg("node"), nb::arg("out_idx") = std::size_t(0))
        .def("write",  &Net::write,
             nb::arg("node"), nb::arg("in_idx"), nb::arg("value"));
}

} // namespace kpn::python

#endif // KPN_BUILD_PYTHON