#define KPN_BUILD_PYTHON
#include <kpn/python/auto_bind.hpp>

#include <nanobind/ndarray.h>

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/videoio.hpp>

#include <chrono>
#include <cmath>
#include <iostream>
#include <thread>
#include <tuple>

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

// ── PythonConverter<cv::Mat> ──────────────────────────────────────────────────
// Converts cv::Mat ↔ numpy array (uint8, HxW or HxWxC shape).
//
// to_python: clones the mat onto the heap; the numpy array owns it via a
//   capsule deleter — no shared cv::Mat refcount dangling after the Variant dies.
// from_python: calls numpy.ascontiguousarray, then clones into an owned cv::Mat.

namespace kpn {

template<> struct PythonConverter<cv::Mat> {
    static constexpr const char* type_name = "mat";

    static nb::object to_python(const cv::Mat& m) {
        // Must be called with the GIL held (always true: called from read() or
        // from within the gil_scoped_acquire block in PyNode::run_loop).
        auto np = nb::module_::import_("numpy");
        cv::Mat c = m.clone();  // ensure contiguous, independently owned
        nb::bytes raw(reinterpret_cast<const char*>(c.data),
                       c.total() * c.elemSize());
        nb::object arr = np.attr("frombuffer")(raw, "uint8");
        int H = c.rows, W = c.cols, C = c.channels();
        arr = arr.attr("reshape")(
            C > 1 ? nb::make_tuple(H, W, C) : nb::make_tuple(H, W));
        return arr.attr("copy")();  // writable, lifetime-independent copy
    }

    static cv::Mat from_python(nb::object o) {
        auto np = nb::module_::import_("numpy");
        // Ensure contiguous uint8 layout (in-place if already compatible)
        nb::object arr = np.attr("ascontiguousarray")(o, "uint8");
        auto shape = nb::cast<std::vector<int>>(arr.attr("shape"));
        if (shape.size() < 2 || shape.size() > 3)
            throw std::runtime_error(
                "cv::Mat from_python: expected 2D (H×W) or 3D (H×W×C) uint8 array");
        int H = shape[0], W = shape[1];
        int C = (shape.size() == 3) ? shape[2] : 1;
        int type = C > 1 ? CV_8UC(C) : CV_8UC1;

        // Cast to ndarray to get the raw data pointer
        auto binfo = nb::cast<nb::ndarray<nb::numpy, uint8_t>>(arr);
        cv::Mat wrap(H, W, type, binfo.data());
        return wrap.clone();  // own the pixel data
    }
};

} // namespace kpn

// ── Pipeline functions ────────────────────────────────────────────────────────

static cv::Mat make_gradient(int W, int H) {
    cv::Mat xr(H, W, CV_8UC1), yg(H, W, CV_8UC1), b(H, W, CV_8UC1, cv::Scalar(128));
    for (int x = 0; x < W; ++x) xr.col(x).setTo(x * 255 / W);
    for (int y = 0; y < H; ++y) yg.row(y).setTo(y * 255 / H);
    cv::Mat channels[3] = {b, yg, xr};
    cv::Mat grad;
    cv::merge(channels, 3, grad);
    return grad;
}

static std::tuple<cv::Mat, cv::Mat> capture() {
    constexpr int W = 640, H = 480;
    static cv::VideoCapture cap;
    static bool opened = false;
    if (!opened) {
        opened = true;
        cap.open(0, cv::CAP_V4L2);
        if (cap.isOpened()) {
            cap.set(cv::CAP_PROP_FRAME_WIDTH,  W);
            cap.set(cv::CAP_PROP_FRAME_HEIGHT, H);
        } else {
            std::cerr << "[capture] no webcam — using synthetic animated pattern\n";
        }
    }

    cv::Mat frame;
    if (cap.isOpened()) {
        auto t0 = std::chrono::steady_clock::now();
        cap >> frame;
        auto elapsed = std::chrono::steady_clock::now() - t0;
        if (elapsed < std::chrono::milliseconds(20))
            std::this_thread::sleep_for(std::chrono::milliseconds(33) - elapsed);
        if (frame.empty()) frame = cv::Mat::zeros(H, W, CV_8UC3);
    } else {
        static int tick = 0;
        static cv::Mat grad = make_gradient(W, H);
        ++tick;
        frame = grad.clone();
        int r = 150 + (tick % 80) * 4;
        cv::circle(frame, {W/2,   H/2},   r,     {255, 200,   0}, -1);
        cv::circle(frame, {W/2,   H/2},   r / 2, {  0, 128, 255}, -1);
        cv::circle(frame, {W*2/5, H*2/5}, r / 3, {200,   0, 200}, -1);
        std::this_thread::sleep_for(std::chrono::milliseconds(33));
    }
    return {frame.clone(), frame.clone()};
}

static cv::Mat to_gray(cv::Mat bgr) {
    cv::Mat gray;
    cv::cvtColor(bgr, gray, cv::COLOR_BGR2GRAY);
    return gray;
}

static cv::Mat edges_fn(cv::Mat gray) {
    cv::Mat blurred, mask;
    cv::GaussianBlur(gray, blurred, {5, 5}, 0);
    cv::Canny(blurred, mask, 50, 150);
    return mask;
}

static cv::Mat quantise(cv::Mat bgr) {
    constexpr int levels = 4;
    constexpr double step = 256.0 / levels;
    static const cv::Mat lut = []() {
        cv::Mat l(1, 256, CV_8UC1);
        for (int i = 0; i < 256; ++i)
            l.at<uchar>(i) = cv::saturate_cast<uchar>(
                std::floor(i / step) * step + step / 2.0);
        return l;
    }();
    cv::Mat out;
    cv::LUT(bgr, lut, out);
    return out;
}

// Returns composite frame AND edge mask so the display node can show both
// without needing a fan-out on the edges channel.
static std::tuple<cv::Mat, cv::Mat> composite(cv::Mat edge_mask, cv::Mat colour) {
    cv::Mat result = colour.clone();
    result.setTo(cv::Scalar(0, 0, 0), edge_mask);
    return {result, edge_mask};
}

// ── Registry ──────────────────────────────────────────────────────────────────
// Variant deduced as std::variant<cv::Mat> — every node uses only cv::Mat.

using CvNodes = NodeRegistry<
    Entry<capture,   "capture">,
    Entry<to_gray,   "to_gray">,
    Entry<edges_fn,  "edges">,
    Entry<quantise,  "quantise">,
    Entry<composite, "composite">
>;

// ── Module ────────────────────────────────────────────────────────────────────

NB_MODULE(kpn_opencv, m) {
    m.doc() = "KPN++ OpenCV bindings for the cell-shading pipeline";

    // Registers: Network, INode, CaptureNode, ToGrayNode, EdgesNode,
    //            QuantiseNode, CompositeNode, and make_<name>() factories.
    // Network.add_node(name, callable, inputs=["mat"], outputs=["mat"])
    // accepts Python callables that receive/return numpy uint8 arrays.
    bind_network<CvNodes>(m);

    // Note: bind_debug is omitted here — cv::Mat functions cannot be called
    // directly from Python without the variant/network machinery. Use
    // net.write() + net.read() to inject/inspect individual nodes instead.
}