#include <kpn/kpn.hpp>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/videoio.hpp>
#include <iostream>
#include <tuple>
#include <thread>
#include <chrono>

// Teach KPN how many bytes a cv::Mat actually carries (header + pixel data).
template<>
struct kpn::ChannelDataSize<cv::Mat> {
    static std::size_t bytes(const cv::Mat& m) { return m.total() * m.elemSize(); }
};

// ── Cell-shading pipeline ─────────────────────────────────────────────────────
//
//  [capture] --"colour"--> [quantise] ──────────────────────────┐
//                                                               ├──> [composite] ──"result"──┐
//  [capture] --"grey"---> [to_gray] --> [edges] ──"edges"───────┘                            │
//                                            └──────────────────────────────"edges"──────────┴──> [display]
//
// DisplayNode derives from MainThreadNode<> — two inputs (composite + raw edges),
// two windows opened in the constructor. step() is called on the main thread.

// ── Gradient base for synthetic pattern ──────────────────────────────────────

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

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

// --8<-- [start:capture_fn]
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()};
}
// --8<-- [end:capture_fn]

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 both the composite frame and the original edge mask so downstream
// nodes (display) can receive both without 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};
}

// ── DisplayNode ───────────────────────────────────────────────────────────────
//
// Two inputs: the cell-shaded composite frame and the raw edge mask.
// MainThreadNode<> provides channel ownership, INode boilerplate, and stats.
// The constructor opens both windows on the main thread (Wayland requirement).
// operator() is called by step() whenever both channels have a frame ready.

// --8<-- [start:display_node]
class DisplayNode : public kpn::MainThreadNode<DisplayNode,
                                               kpn::in<"composite", "edges">,
                                               cv::Mat, cv::Mat> {
public:
    DisplayNode() : MainThreadNode(8) {
        cv::namedWindow("Cell Shade", cv::WINDOW_NORMAL);
        cv::namedWindow("Edge Mask",  cv::WINDOW_NORMAL);
        cv::resizeWindow("Cell Shade", 1280, 720);
        cv::resizeWindow("Edge Mask",   640, 360);
    }

    ~DisplayNode() { cv::destroyAllWindows(); }

    bool operator()(cv::Mat composite, cv::Mat edges) {
        cv::imshow("Cell Shade", composite);
        cv::Mat edges_bgr;
        cv::cvtColor(edges, edges_bgr, cv::COLOR_GRAY2BGR);
        cv::imshow("Edge Mask", edges_bgr);
        int key = cv::waitKey(1);
        if (key == 'q' || key == 27) return false;
        return window_open("Cell Shade") && window_open("Edge Mask");
    }

private:
    static bool window_open(const char* name) {
        try { return cv::getWindowProperty(name, cv::WND_PROP_VISIBLE) >= 1; }
        catch (const cv::Exception&) { return false; }
    }
};
// --8<-- [end:display_node]

// ─────────────────────────────────────────────────────────────────────────────

int main() {
    using namespace kpn;

    // --8<-- [start:opencv_network]
    auto src       = make_node<capture>  (out<"colour","grey">{},            8);
    auto gray_node = make_node<to_gray>  (in<"bgr">{},   out<"gray">{},      8);
    auto edge_node = make_node<edges_fn> (in<"gray">{},  out<"edges">{},     8);
    auto quant     = make_node<quantise> (in<"bgr">{},   out<"quantised">{}, 8);
    auto comp      = make_node<composite>(in<"edges","colour">{}, out<"result","edges">{}, 8);

    // DisplayNode: two windows opened in constructor, step() drives main thread.
    DisplayNode disp;

    Network net;
    net.add("src",     src)
       .add("gray",    gray_node)
       .add("edges",   edge_node)
       .add("quant",   quant)
       .add("comp",    comp)
       .add("display", disp)
       .connect("src",   src.template output<"colour">(),      "quant",   quant.template input<"bgr">())
       .connect("quant", quant.template output<"quantised">(), "comp",    comp.template input<"colour">())
       .connect("src",   src.template output<"grey">(),        "gray",    gray_node.template input<"bgr">())
       .connect("gray",  gray_node.template output<"gray">(),  "edges",   edge_node.template input<"gray">())
       .connect("edges", edge_node.template output<"edges">(), "comp",    comp.template input<"edges">())
       .connect("comp",  comp.template output<"result">(),     "display", disp.template input<"composite">())
       .connect("comp",  comp.template output<"edges">(),      "display", disp.template input<"edges">())
       .build();
    // --8<-- [end:opencv_network]

    net.set_watchdog_interval(std::chrono::milliseconds(5000));
#ifdef KPN_WEB_DEBUG
    net.set_web_debug_port(9090);
#endif

    std::cout << "Cell-shading pipeline running. Press 'q' to stop.\n";
    std::cout << "Web debug UI: http://localhost:9090\n";

    // --8<-- [start:main_thread_step]
    net.start();

    // Main thread drives display — imshow/waitKey stay on the GUI thread.
    // step() returns false when operator() returns false (q pressed / window closed).
    while (disp.step())
        cv::waitKey(8); // yield event loop when no frame ready

    net.stop();
    // --8<-- [end:main_thread_step]
    return 0;
}