#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 PyNetwork {
public:
    using VNode    = IVariantNode<Variant>;
    using VChannel = IVariantChannel<Variant>;

    // ── 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)
    // Wires src's output port out_idx to dst's input port in_idx.
    // Type check: both sides must carry the same T.
    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() + ")");

        // The destination node owns the input channel — get it, then tell src to use it.
        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 ─────────────────────────────────────────────────────

    // Read one value from node's output port. Releases GIL while blocking.
    nb::object read(const std::string& node_name, std::size_t out_idx) {
        // We need a channel that sits on the output of this node.
        // read() installs a tap channel if not already present.
        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");
            // Create a tap channel matching the output type and wire it
            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));
    }

    // Write a Python value into node's input port. Releases GIL while blocking.
    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));
        }
    }

    // ── Converter registration ────────────────────────────────────────────────
    // Called once per type at module init time to register to/from Python converters.

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

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) {
        // Create the right VariantChannel<T> based on the registered type index.
        // We need a factory registered per type — stored in tap_factories_.
        auto it = tap_factories_.find(type);
        if (it == tap_factories_.end())
            throw std::runtime_error(
                "no tap factory for type: " + std::string(type.name()) +
                " — was register_type() called for this type?");
        return it->second();
    }

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

public:
    // Called by register_type to also register a tap channel factory.
    template<typename T>
    void register_tap_factory(std::size_t capacity = 5) {
        auto idx = std::type_index(typeid(T));
        tap_factories_[idx] = [capacity]() -> std::shared_ptr<VChannel> {
            auto ch = std::make_shared<Channel<T>>(capacity);
            return std::make_shared<VariantChannel<T, Variant>>(std::move(ch));
        };
    }

private:
    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_;
    std::map<std::type_index, std::function<std::shared_ptr<VChannel>()>> tap_factories_;
};

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

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();
            // Release GIL while joining — run_loop may be waiting to acquire it.
            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);
    }

private:
    void run_loop() {
        // This thread does not hold the GIL. It acquires it only for Python calls.
        while (!stop_flag_.load(std::memory_order_relaxed)) {
            try {
                auto t0 = clock_t::now();

                // Pop all inputs — no GIL needed, these are pure C++ channel ops
                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();

                // Acquire GIL only for the Python call and type conversion
                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);

                // Push outputs — no GIL needed
                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_;
};

// ── register_py_network ───────────────────────────────────────────────────────
// Registers PyNetwork<Variant> and PyNode<Variant> with the given nanobind module.
// Call once per module, passing the Variant type derived from your registered node types.

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)
        .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") = 0)
        .def("write",  &Net::write,
             nb::arg("node"), nb::arg("in_idx"), nb::arg("value"));
}

} // namespace kpn::python

#endif // KPN_BUILD_PYTHON