// This file is part of OpenCV project.
// It is subject to the license terms in the LICENSE file found in the top-level directory
// of this distribution and at http://opencv.org/license.html.
//
#include <iostream>
#include <vector>
#include <algorithm>
#include <iostream>
#include <fstream>
#include <chrono>
#include <filesystem>
#include <sys/stat.h>
#include <sys/types.h>
#include <stdlib.h>
#include <unistd.h>
#include "opencv2/objdetect.hpp"
#include <opencv2/dnn.hpp>

#include "mq_worker.h"
#if ENABLE_GRPC
#include "fileprocess.h"
#endif
#include "http.h"
#include "libqr.h"

static
int detect_cmd(char **argv, int argc)
{
    char *file = argv[0];
    auto orig = imread(file);
    QRCodeDetector detector;

    Mat points;
    Mat straight;
    auto r = detector.detectAndDecode(orig, points, straight);
    printf("r: %s\n", r.c_str());
    printf("points: %d %d\n", points.rows, points.cols);
    for (int i = 0; i < points.cols; i++) {
        printf("%f ", points.at<float>(0, i));
    }
    printf("\n");
    return 0;
}

static
int angle_cmd(char **argv, int argc)
{
    char *file = argv[0];
    printf("file: %s\n", file);
    Mat orig = imread(file);
    string qrcode, err;
    float angle;
    ProcessState ps;
    auto r = emblem_dot_angle(ps, orig, angle, qrcode, err);

    if (!r) {
        cerr << r << ":" << err << endl;
        return 1;
    }
    printf("angle: %.1f\n", angle);
    printf("qrcode: %s\n", qrcode.c_str());
    return 0;
}

static
int dot_cmd(char **argv, int argc)
{
    ProcessState ps;
    char *file = argv[0];
    printf("file: %s\n", file);
    Mat orig = imread(file);
    string qrcode, err;
    float angle;
    auto r = emblem_dot_angle(ps, orig, angle, qrcode, err);

    if (!r) {
        cerr << r << ":" << err << endl;
        return 1;
    }
    string outfile = string(file) + ".dot.jpg";
    printf("angle: %.1f\n", angle);
    printf("qrcode: %s\n", qrcode.c_str());
    printf("saving dot file: %s\n", outfile.c_str());
    imwrite(outfile, ps.dot_area);
    return 0;
}

static
int clarity_cmd(char **argv, int argc)
{
    string err;

    char *file = argv[0];
    printf("file: %s\n", file);
    Mat orig = imread(file);
    Mat gray;
    cvtColor(orig, gray, COLOR_BGR2GRAY);
    auto c = laplacian(gray, err);
    printf("clarity: %lf\n", c);

    return 0;
}

static
int rectify_cmd(char **argv, int argc)
{
    char *file = argv[0];
    string err;
    ProcessState ps;
    Mat orig = imread(file);

    ps.orig = &orig;
    preprocess(ps);

    if (!detect_qr(ps, 0.20, false, err)) {
        cerr << err << endl;
        return 1;
    }
    string outfile = string(file) + ".qr.jpg";
    imwrite(outfile, ps.straighten);
    return 0;
}

static
int topleft_cmd(char **argv, int argc)
{
    char *file = argv[0];
    string err;
    ProcessState ps;
    Mat orig = imread(file);

    ps.orig = &orig;
    preprocess(ps);

    if (!detect_qr(ps, 0.02, true, err)) {
        cerr << err << endl;
        return 1;
    }
    string outfile = string(file) + ".topleft.jpg";
    Mat &base = ps.straighten;
    auto crop = Rect(0, 0, base.cols / 2, base.rows / 2);
    Mat result = base(crop);
    imwrite(outfile, result);
    return 0;
}

static
int find_roi_start_point(Mat &bin, Point &p)
{
    int npoints = 4;

    for (int i = 0; i < bin.cols / 3; i++) {
        uchar sum = 0;
        for (int j = 0; j < npoints; j++) {
            int v = i + j;
            sum += bin.at<uchar>(v, v);
        }
        if (sum == 0) {
            p = Point(i, i);
            return 0;
        }
    }
    cerr << "find_roi_start_point" << endl;
    return -1;
}

static
vector<int> count_black(Mat &bin, bool count_rows, int size)
{
    vector<int> ret;
    for (int i = 0; i < size; i++) {
        int count = 0;
        for (int j = 0; j < size; j++) {
            int x = count_rows ? j : i;
            int y = count_rows ? i : j;
            if (bin.at<uchar>(y, x) == 0) count++;
        }
        ret.push_back(count);
    }
    return ret;
}

static int find_start_of_first_black_range(const vector<int> &data)
{
    size_t i = 0;
    int m = *std::max_element(data.begin(), data.end());
    int thres = m * 50 / 100;
    while (i < data.size() - 3) {
        if (data[i] >= thres && data[i + 1] >= thres && data[i + 2] >= thres) {
            break;
        }
        i++;
    }
    if (i >= data.size() - 3) return -1;
    return i;
}

static int find_end_of_first_black_range(const vector<int> &data)
{
    size_t i = 0;
    int m = *std::max_element(data.begin(), data.end());
    int thres = m * 50 / 100;
    while (i < data.size() - 3) {
        if (data[i] >= thres && data[i + 1] >= thres && data[i + 2] >= thres) {
            break;
        }
        i++;
    }
    if (i >= data.size() - 3) return -1;
    while (i < data.size() - 3) {
        if (data[i] < thres || data[i + 1] < thres || data[i + 2] < thres) {
            break;
        }
        i++;
    }
    if (i >= data.size() - 3) return -1;
    return i;
}

static
int find_roi_rect(Mat &bin, Point &start, Rect &rect, bool inner, string &err)
{
    Mat visited;
    bin.copyTo(visited);
    vector<Point> q;
    q.push_back(start);
    int min_x = bin.rows, min_y = bin.rows, max_x = 0, max_y = 0;
    int orig_size = max(bin.rows, bin.cols);

    while (q.size()) {
        Point p = q.back();
        q.pop_back();
        visited.at<uchar>(p.y, p.x) = 255;
        min_x = min(min_x, p.x);
        min_y = min(min_y, p.y);
        max_x = max(max_x, p.x);
        max_y = max(max_y, p.y);
        for (int xoff = -1; xoff <= 1; xoff++) {
            for (int yoff = -1; yoff <= 1; yoff++) {
                int x = p.x + xoff;
                int y = p.y + yoff;
                if (x >= 0 && x < visited.cols && y >= 0 && y < visited.rows) {
                    auto v = visited.at<uchar>(y, x);
                    if (v == 0) {
                        if (q.size() >= 20000) {
                            err = string("roi detected range too large: ") + to_string(q.size());
                            return -1;
                        }
                        q.push_back(Point(x, y));
                        visited.at<uchar>(y, x) = 255;
                    }
                }
            }
        }
    }
    if (max_x - min_x < 50 || max_y - min_y < 50) {
        err = "detected roi outer region too small";
        return -1;
    }

    auto size = std::max(max_x, max_y);

    auto row_sums = count_black(bin, true, size);
    auto col_sums = count_black(bin, false, size);

    min_x = inner ? find_end_of_first_black_range(col_sums) : find_start_of_first_black_range(col_sums);
    min_y = inner ? find_end_of_first_black_range(row_sums) : find_start_of_first_black_range(row_sums);

    if (min_x < 0 || min_y < 0) {
        err = "min_x or min_y is negative";
        return -1;
    }

    // find the max values, similarly
    std::reverse(col_sums.begin(), col_sums.end());
    std::reverse(row_sums.begin(), row_sums.end());
    max_x = size - (inner ? find_end_of_first_black_range(col_sums) : find_start_of_first_black_range(col_sums));
    max_y = size - (inner ? find_end_of_first_black_range(row_sums) : find_start_of_first_black_range(row_sums));

    if (max_x < 0 || max_y < 0) return -1;

    if (max_x - min_x < 50 || max_y - min_y < 50) {
        err = "detected roi region too small";
        return -1;
    }
    size = (max_x - min_x + max_y - min_y) / 2;

    if (size < orig_size / 5 || size > orig_size * 3 / 5) {
        err = "size of found region is out of valid range";
        return -1;
    }
    rect = Rect(min_x + 1, min_y + 1, size, size);

    return 0;
}

static void get_bin(Mat &orig, Mat &out)
{
    Mat gray;
    Mat filtered;
    Point start;
    Rect roi_rect;
    cvtColor(orig, gray, COLOR_BGR2GRAY);
    convertScaleAbs(gray, gray, 2);
    // bilateralFilter(gray, filtered, 9, 150, 150, BORDER_DEFAULT);
    medianBlur(gray, gray, 9);
    adaptiveThreshold(gray, out, 255, ADAPTIVE_THRESH_GAUSSIAN_C, THRESH_BINARY, 11, 2);
}

static
int find_roi(Mat &qr_straighten, Mat &roi, bool inner, string &err)
{
    Mat bin;
    Rect topleft_r(0, 0, qr_straighten.cols / 2, qr_straighten.rows / 2);
    Mat topleft = qr_straighten(topleft_r);
    get_bin(topleft, bin);

    Point start;
    Rect roi_rect;
    auto r = find_roi_start_point(bin, start);
    if (r) {
        err = "failed to find roi start point";
        return r;
    }
    r = find_roi_rect(bin, start, roi_rect, inner, err);
    if (r) {
        return r;
    }
    roi = qr_straighten(roi_rect);
    return 0;
}

int roi_process_one(const char *file, bool inner, string &err, bool warp, Mat *roi_out)
{
    Mat roi;
    Mat orig = imread(file, IMREAD_COLOR);
    Mat qr_with_margin;

    if (warp) {
        ProcessState ps;
        ps.orig = &orig;
        preprocess(ps);
        if (!detect_qr(ps, 0.02, true, err)) {
            cerr << err << ":" << file << endl;
            return 1;
        }
        qr_with_margin = ps.straighten;
    } else {
        qr_with_margin = orig;
    }
    if (qr_with_margin.cols <= 0 || qr_with_margin.rows <= 0) return -1;
    auto r = find_roi(qr_with_margin, roi, inner, err);
    if (r) return r;
    if (roi_out) {
        *roi_out = roi;
    }
    return 0;
}

static bool is_dir(char *path)
{
    struct stat st;

    if (lstat(path, &st) == 0) {
        if (S_ISREG(st.st_mode)) {
            return false;
        } else if (S_ISDIR(st.st_mode)) {
            return true;
        } else {
            return false;
        }
    } else {
        perror("Error in lstat");
    }
    return false;
}

static void save_roi(string orig_file, Mat &roi)
{
    string outfile = orig_file + ".roi.jpg";
    cout << "save: " << outfile << endl;
    imwrite(outfile, roi);
}

static
int frame_roi_cmd(char **argv, int argc)
{
    char *file = argv[0];
    string err;
    int ret = 0;
    Mat roi;

    cout << "frame roi processing: " << file << endl;
    ret = roi_process_one(file, false, err, true, &roi);
    if (ret) {
        cerr << "failed to process: " << file << ":" << err <<endl;
        return ret;
    }
    save_roi(file, roi);
    return ret;
}

static
int roi_cmd(char **argv, int argc)
{
    char *file = argv[0];
    string err;
    int ret = 0;

    cout << "roi processing: " << file << endl;
    if (is_dir(file)) {
        for (auto const& dir_entry : filesystem::directory_iterator{file}) {
            auto path = dir_entry.path();
            Mat roi;
            if (roi_process_one(path.c_str(), false, err, false, &roi) != 0) {
                cerr << "failed: " << path << ":" << err << endl;
                ret = 1;
            }
            save_roi(file, roi);
        }
    } else {
        Mat roi;
        ret = roi_process_one(file, false, err, false, &roi);
        if (ret) {
            cerr << "failed to process: " << file << ":" << err <<endl;
        } else {
            save_roi(file, roi);
        }
    }

    return ret;
}

static
int roi_bench_cmd(char **argv, int argc)
{
    char *file = argv[0];
    int n = 0;
    string err;
    auto begin = chrono::system_clock::now();
    for (auto const& dir_entry : filesystem::directory_iterator{file}) {
        auto path = dir_entry.path();
        if (roi_process_one(path.c_str(), false, err, false, NULL) == 0) {
            n += 1;
        }
    }
    auto end = chrono::system_clock::now();
    std::chrono::duration<float> duration = end - begin;
    float seconds = duration.count();
    printf("qps: %.1f\n", n / seconds);
    return 0;
}

#if USE_PULSAR
static vector<string> split_path(const string &path)
{
    vector<string> ret;
    string cur = "";
    for (auto x: path) {
        if (x == '/') {
            if (cur.size()) {
                ret.push_back(cur);
                cur = "";
            }
        } else {
            cur += x;
        }
    }
    if (cur.size()) {
        ret.push_back(cur);
        cur = "";
    }
    return ret;
}

static string join_path(const vector<string> fs)
{
    string ret;
    const string sep = "/";

    for (auto x: fs) {
        ret += sep + x;
    }
    if (ret.size() == 0) ret = "/";
    return ret;
}

static string get_output_path(const string &path)
{
    auto ret = split_path(path);
    ret[3] = "roi";
    return join_path(ret);
}

static
int roi_worker_handle_image(const string &input_path,
                            const vector<uint8_t> &input,
                            string &output_path,
                            vector<uint8_t> &output)
{
    Mat roi;
    Mat orig = imdecode(input, IMREAD_COLOR);

    auto r = find_roi(orig, roi, true, true);
    if (r) return r;
    imencode(".jpg", roi, output);
    output_path = get_output_path(input_path);
    return 0;
}

static
int roi_worker_handle_image_nop(const string &input_path,
                                const vector<uint8_t> &input,
                                string &output_path,
                                vector<uint8_t> &output)
{
    output.push_back('f');
    output.push_back('o');
    output.push_back('o');
    output_path = get_output_path(input_path);
    return 0;
}

static
int roi_worker_cmd(char *topic)
{
    string worker_name = "roi-worker-";
    worker_name += to_string(rand());
    return mq_worker(topic, "roi-worker", roi_worker_handle_image);
}

static
int roi_worker_nop_cmd(char *topic)
{
    string worker_name = "roi-worker-";
    worker_name += to_string(rand());
    return mq_worker(topic, "roi-worker", roi_worker_handle_image_nop);
}

#endif

#if ENABLE_GRPC
static
int grpc_server_cmd(char *addr)
{
    return run_server(addr, roi_worker_handle_image);
}
#endif

static
int http_server_handle_image(const vector<uint8_t> &input,
                             vector<uint8_t> &output)
{
    string err;
    Mat roi;
    Mat orig = imdecode(input, IMREAD_COLOR);

    if (orig.empty()) {
        return -EINVAL;
    }
    auto r = find_roi(orig, roi, false, err);
    if (r) return r;
    imencode(".jpg", roi, output);
    return 0;
}

static
int http_server_cmd(char **argv, int argc)
{
    char *port = argv[0];
    return start_http_server(atoi(port), http_server_handle_image);
}

static
int verify_cmd(char **args, int nargs)
{
    char *std_file = args[0];
    char *frame_file = args[1];
    Mat std = imread(std_file);
    Mat frame = imread(frame_file);
    Mat roi;
    int r;

    string err;
    r = roi_process_one(frame_file, false, err, true, &roi);
    if (r) {
        printf("failed to find roi: %s\n", err.c_str());
        return r;
    }
    double s = emblem_roi_similarity(FuzzyPixelCmp, std, roi, err);

    if (err.size()) {
        printf("err: %s\n", err.c_str());
        return 1;
    } else {
        printf("similarity: %f\n", s); return 0;
        return 0;
    }
}

static
void usage(const char *name, vector<string> &cmds)
{
    printf("usage: %s <cmd> <arg0>\n", name);
    printf("or for 2 args: %s <cmd> <arg0> <arg1>\n", name);
    printf("possible commands:\n");
    for (auto cmd: cmds) {
        printf("  %s\n", cmd.c_str());
    }
}

static
int roi_verify_cmd(char **argv, int argc)
{
    char *model_path = argv[0];
    char *input_file = argv[1];
    Mat input_img = imread(input_file);
    cv::dnn::Net net = cv::dnn::readNetFromONNX(model_path);
    if (net.empty()) {
        std::cerr << "Failed to load ONNX model!" << std::endl;
        return -1;
    }
    cv::Mat blob;
    cv::resize(input_img, input_img, cv::Size(128, 64));  // 调整图像大小
    cv::dnn::blobFromImage(input_img, blob, 1.0 / 255.0, cv::Size(64, 128), cv::Scalar(0.485, 0.456, 0.406), true, false);
    blob = (blob - cv::Scalar(0.485, 0.456, 0.406)) / cv::Scalar(0.229, 0.224, 0.225);  // 归一化
    net.setInput(blob);
    cv::Mat output = net.forward();
    float* data = (float*)output.data;
    int class_id = std::max_element(data, data + output.total()) - data;  // 找到最大概率的类别
    float confidence = data[class_id];
    std::cout << "Predicted class: " << class_id << ", Confidence: " << confidence << std::endl;

    return 0;
}

static
int side_by_side_cmd(char **argv, int argc)
{
    Mat frame_roi;
    string err;
    cout << "side by side processing: " << argv[0] << " " << argv[1] << endl;
    int r = roi_process_one(argv[0], false, err, true, &frame_roi);
    if (r) {
        cerr << "failed to process: " << argv[0] << ":" << err << endl;
        return r;
    }
    Mat std_roi = imread(argv[1]);
    Mat side_by_side;
    auto roi_size = Size(128, 128);
    resize(frame_roi, frame_roi, roi_size);
    resize(std_roi, std_roi, roi_size);
    hconcat(frame_roi, std_roi, side_by_side);
    auto fn = string(argv[0]) + ".side_by_side.jpg";
    imwrite(fn, side_by_side);
    return 0;
}

#ifdef QRTOOL_MAIN
int main(int argc, char *argv[])
{
    string cmd = "help";
    if (argc > 1) {
        cmd = argv[1];
    }
    vector<string> cmds;

#define add_cmd(c, nargs) \
    do { \
        cmds.push_back(#c); \
        if (cmd == #c && argc >= 2 + nargs) return c##_cmd(&argv[2], argc - 2); \
    } while (0)

    add_cmd(detect, 1);
    add_cmd(angle, 1);
    add_cmd(dot, 1);
    add_cmd(clarity, 1);
    add_cmd(rectify, 1);
    add_cmd(topleft, 1);
    add_cmd(frame_roi, 1);
    add_cmd(roi, 1);
    add_cmd(roi_bench, 1);
#if USE_PULSAR
    add_cmd(roi_worker, 1);
    add_cmd(roi_worker_nop, 1);
#endif
#if ENABLE_GRPC
    add_cmd(grpc_server, 1);
#endif
    add_cmd(http_server, 1);
    add_cmd(verify, 2);
    add_cmd(roi_verify, 2);
    add_cmd(side_by_side, 2);
    usage(argv[0], cmds);
    return 1;
}
#endif