#include <pybind11/functional.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>

namespace py = pybind11;

namespace pipeline {

struct Data {
    // Placeholder for future C++ payload; expose a generic Python dict to share state.
    py::dict payload;
};

using OpFn = void (*)(Data &);

enum class ImplKind { Cpp, Python };

struct OpEntry {
    ImplKind impl;
    OpFn cpp_fn = nullptr;
    py::object py_fn;  // holds ref for Python ops
};

static std::unordered_map<std::string, OpEntry> &registry() {
    static std::unordered_map<std::string, OpEntry> reg;
    return reg;
}

static std::mutex &registry_mutex() {
    static std::mutex m;
    return m;
}

void register_cpp_op(const std::string &name, OpFn fn) {
    std::lock_guard<std::mutex> lock(registry_mutex());
    auto &reg = registry();
    if (reg.count(name)) {
        throw std::runtime_error("Operation already registered: " + name);
    }
    reg[name] = OpEntry{ImplKind::Cpp, fn, py::none()};
}

void register_py_op(const std::string &name, py::object fn) {
    if (!fn) {
        throw std::runtime_error("register_py_op: callable is null");
    }
    std::lock_guard<std::mutex> lock(registry_mutex());
    auto &reg = registry();
    if (reg.count(name)) {
        throw std::runtime_error("Operation already registered: " + name);
    }
    reg[name] = OpEntry{ImplKind::Python, nullptr, std::move(fn)};
}

const OpEntry &get_op(const std::string &name) {
    auto &reg = registry();
    auto it = reg.find(name);
    if (it == reg.end()) {
        throw std::runtime_error("Unknown operation: " + name);
    }
    return it->second;
}

std::vector<std::string> list_ops() {
    std::vector<std::string> names;
    names.reserve(registry().size());
    for (auto const &kv : registry()) {
        names.push_back(kv.first);
    }
    return names;
}

std::vector<py::dict> describe_ops() {
    std::vector<py::dict> out;
    out.reserve(registry().size());
    for (auto const &kv : registry()) {
        const auto &name = kv.first;
        const auto &entry = kv.second;
        py::dict d;
        d["name"] = name;
        if (entry.impl == ImplKind::Cpp) {
            d["impl"] = "cpp";
            d["cpp_fn_ptr"] = reinterpret_cast<uintptr_t>(entry.cpp_fn);
        } else {
            d["impl"] = "python";
            d["callable"] = entry.py_fn;
        }
        out.push_back(std::move(d));
    }
    return out;
}

class Pipeline {
  public:
    void add_step(const std::string &name) { steps_.push_back(name); }

    void run(Data &d) const {
        for (auto const &name : steps_) {
            const auto &entry = get_op(name);
            if (entry.impl == ImplKind::Cpp) {
                if (!entry.cpp_fn) {
                    throw std::runtime_error("Null cpp function pointer for op: " + name);
                }
                entry.cpp_fn(d);
            } else {
                py::gil_scoped_acquire gil;
                entry.py_fn(d);
            }
        }
    }

  private:
    std::vector<std::string> steps_;
};

// C API for C++ plugins (PyCapsule)
using register_op_c_api_t = void (*)(const char *name, OpFn fn);

static void register_op_c_api(const char *name, OpFn fn) {
    register_cpp_op(std::string{name}, fn);
}

}  // namespace pipeline

PYBIND11_MODULE(_backend, m) {
    using namespace pipeline;

    py::class_<Data>(m, "Data")
        .def(py::init<>())
        .def_readwrite("payload", &Data::payload, "Generic payload dict shared across ops");

    py::class_<Pipeline>(m, "Pipeline")
        .def(py::init<>())
        .def("add_step", &Pipeline::add_step, py::arg("name"))
        .def("run", &Pipeline::run, py::arg("data"));

    m.def("register_cpp_op", &register_cpp_op, py::arg("name"), py::arg("fn"));
    m.def("register_py_op", &register_py_op, py::arg("name"), py::arg("fn"));
    m.def("list_ops", &list_ops);
    m.def("describe_ops", &describe_ops);

    // Expose capsule for C++ plugins
    auto capsule = py::capsule(reinterpret_cast<void *>(register_op_c_api), "pipeline_backend.register_op_v1");
    m.attr("_register_op_capsule") = capsule;
}
