#define KPN_BUILD_PYTHON
#include <kpn/python/bindings.hpp>
#include <nanobind/nanobind.h>
#include <nanobind/stl/shared_ptr.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>

#include <iostream>
#include <map>

namespace nb = nanobind;
using namespace kpn;
using namespace kpn::python;

// ── Node functions for the hello-pipeline examples ────────────────────────────

static int  produce()        { return 42; }
static int  double_it(int x) { return x * 2; }
static void print_it(int x)  { std::cout << "result: " << x << '\n'; }

// ── Variant type ──────────────────────────────────────────────────────────────
// Deduplicated port types across all registered node functions.
// produce: () → int  |  double_it: int → int  |  print_it: int → void
// Unique types: int

using KpnVariant  = std::variant<int>;
using Net         = PyNetwork<KpnVariant>;
using ProduceNode = VariantNodeWrapper<produce,   KpnVariant>;
using DoubleItNode= VariantNodeWrapper<double_it,  KpnVariant>;
using PrintItNode = VariantNodeWrapper<print_it,   KpnVariant>;

// ── Converter helpers for int ─────────────────────────────────────────────────

static nb::object int_to_py(const KpnVariant& v) {
    return nb::int_(std::get<int>(v));
}

static KpnVariant int_from_py(nb::object o) {
    return KpnVariant{ nb::cast<int>(o) };
}

// ── Type name resolver (extensible) ──────────────────────────────────────────

static std::type_index resolve_type(const std::string& name) {
    if (name == "int") return std::type_index(typeid(int));
    throw std::runtime_error("unknown type name '" + name +
                             "' — only 'int' is registered in this module");
}

// ── Ensure converters are registered on a Net instance ───────────────────────

static void ensure_converters(Net& net) {
    net.register_type<int>(
        [](const int& v) -> nb::object { return nb::int_(v); },
        [](nb::object o) -> int        { return nb::cast<int>(o); }
    );
    net.register_tap_factory<int>();
}

NB_MODULE(kpn_python, m) {
    m.doc() = "KPN++ Python bindings — Kahn Process Network library";

    // ── IVariantNode base ─────────────────────────────────────────────────────
    nb::class_<IVariantNode<KpnVariant>>(m, "INode");

    // ── Concrete C++ node wrappers ─────────────────────────────────────────────
    // Factories return shared_ptr so Python can pass them to net.add().
    nb::class_<ProduceNode,  IVariantNode<KpnVariant>>(m, "ProduceNode")
        .def("__init__", [](ProduceNode* self, std::size_t cap) {
            new (self) ProduceNode(cap);
        }, nb::arg("capacity") = 5);

    nb::class_<DoubleItNode, IVariantNode<KpnVariant>>(m, "DoubleItNode")
        .def("__init__", [](DoubleItNode* self, std::size_t cap) {
            new (self) DoubleItNode(cap);
        }, nb::arg("capacity") = 5);

    nb::class_<PrintItNode,  IVariantNode<KpnVariant>>(m, "PrintItNode")
        .def("__init__", [](PrintItNode* self, std::size_t cap) {
            new (self) PrintItNode(cap);
        }, nb::arg("capacity") = 5);

    // Expose factory functions that return shared_ptr — these are what net.add() accepts.
    m.def("make_produce",   [](std::size_t cap) -> std::shared_ptr<IVariantNode<KpnVariant>> {
        return std::make_shared<ProduceNode>(cap);
    }, nb::arg("capacity") = 5);
    m.def("make_double_it", [](std::size_t cap) -> std::shared_ptr<IVariantNode<KpnVariant>> {
        return std::make_shared<DoubleItNode>(cap);
    }, nb::arg("capacity") = 5);
    m.def("make_print_it",  [](std::size_t cap) -> std::shared_ptr<IVariantNode<KpnVariant>> {
        return std::make_shared<PrintItNode>(cap);
    }, nb::arg("capacity") = 5);

    // ── Network ───────────────────────────────────────────────────────────────
    nb::class_<Net>(m, "Network")
        .def(nb::init<>())

        // add(name, c++_node) — for pre-constructed C++ nodes
        .def("add", [](Net& self, std::string name,
                       std::shared_ptr<IVariantNode<KpnVariant>> node) {
            self.add(std::move(name), std::move(node));
        }, nb::arg("name"), nb::arg("node"))

        // add_node(name, callable, inputs=[type_names], outputs=[type_names])
        // Creates a pure Python processing node.
        .def("add_node", [](Net& self,
                            std::string name,
                            nb::object callable,
                            std::vector<std::string> in_names,
                            std::vector<std::string> out_names,
                            std::size_t capacity)
        {
            std::vector<std::type_index> in_types, out_types;
            for (auto& s : in_names)  in_types.push_back(resolve_type(s));
            for (auto& s : out_names) out_types.push_back(resolve_type(s));

            std::map<std::type_index, std::function<nb::object(const KpnVariant&)>> to_py;
            std::map<std::type_index, std::function<KpnVariant(nb::object)>>        from_py;
            std::map<std::type_index, PyNode<KpnVariant>::ChannelFactory>            ch_factories;

            auto int_idx = std::type_index(typeid(int));
            to_py[int_idx]      = int_to_py;
            from_py[int_idx]    = int_from_py;
            ch_factories[int_idx] = [](std::size_t cap) -> std::shared_ptr<IVariantChannel<KpnVariant>> {
                auto ch = std::make_shared<Channel<int>>(cap);
                return std::make_shared<VariantChannel<int, KpnVariant>>(std::move(ch));
            };

            auto node = std::make_shared<PyNode<KpnVariant>>(
                std::move(callable),
                std::move(in_types),
                std::move(out_types),
                std::move(to_py),
                std::move(from_py),
                std::move(ch_factories),
                capacity
            );
            self.add(std::move(name), std::move(node));
        },
        nb::arg("name"), nb::arg("callable"),
        nb::arg("inputs")   = std::vector<std::string>{},
        nb::arg("outputs")  = std::vector<std::string>{},
        nb::arg("capacity") = 5)

        .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)

        // read(node, out_idx=0) — blocking pop from a C++ node's output port into Python
        .def("read", [](Net& self, const std::string& node, std::size_t out_idx) {
            ensure_converters(self);
            return self.read(node, out_idx);
        }, nb::arg("node"), nb::arg("out_idx") = 0)

        // write(node, in_idx, value) — push a Python value into a node's input port
        .def("write", [](Net& self, const std::string& node,
                         std::size_t in_idx, nb::object value) {
            ensure_converters(self);
            self.write(node, in_idx, std::move(value));
        }, nb::arg("node"), nb::arg("in_idx"), nb::arg("value"));
}