#include <fstream>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <vector>
#include <string>
#include <cstring>

using namespace cv;
using namespace cv::dnn;

// 图像结构体定义
typedef struct {
    unsigned char* data;
    int width;
    int height;
    int channels;
} Image;

static std::string last_error;

// YOLO 检测器
class YOLODetector {
public:
    Net net;
    std::vector<std::string> classNames;
    float confThreshold;
    float nmsThreshold;
    
    YOLODetector() : confThreshold(0.5f), nmsThreshold(0.4f) {}
};

static std::map<std::string, YOLODetector*> yolo_detectors;

extern "C" {

// 图像解码
Image* cv_imdecode(const unsigned char* buf, size_t size) {
    try {
        std::vector<unsigned char> data(buf, buf + size);
        cv::Mat mat = cv::imdecode(data, cv::IMREAD_COLOR);
        if (mat.empty()) {
            last_error = "无法解码图像";
            return nullptr;
        }

        Image* img = new Image();
        img->width = mat.cols;
        img->height = mat.rows;
        img->channels = mat.channels();
        
        size_t data_size = mat.total() * mat.elemSize();
        img->data = (unsigned char*)malloc(data_size);
        memcpy(img->data, mat.data, data_size);

        return img;
    } catch (const std::exception& e) {
        last_error = e.what();
        return nullptr;
    }
}

// 释放图像
void cv_image_free(Image* img) {
    if (img) {
        if (img->data) {
            free(img->data);
        }
        delete img;
    }
}

// 加载 YOLO ONNX 模型
int cv_yolo_load(const char* name, const char* model_path, const char* classes_path) {
    try {
        YOLODetector* detector = new YOLODetector();
        
        // 加载 ONNX 模型
        detector->net = readNetFromONNX(model_path);
        if (detector->net.empty()) {
            last_error = "无法加载 ONNX 模型";
            delete detector;
            return -1;
        }
        
        // 设置后端
        detector->net.setPreferableBackend(DNN_BACKEND_OPENCV);
        detector->net.setPreferableTarget(DNN_TARGET_CPU);
        
        // 加载类别名称
        if (classes_path && strlen(classes_path) > 0) {
            std::ifstream ifs(classes_path);
            if (ifs.is_open()) {
                std::string line;
                while (std::getline(ifs, line)) {
                    if (!line.empty()) {
                        detector->classNames.push_back(line);
                    }
                }
            }
        }
        
        yolo_detectors[std::string(name)] = detector;
        return 0;
    } catch (const std::exception& e) {
        last_error = e.what();
        return -1;
    }
}

// YOLO 检测结果
typedef struct {
    float x1, y1, x2, y2;
    float confidence;
    int class_id;
    char class_name[64];
} YOLODetection;

// YOLO 检测
int cv_yolo_detect(const char* name, const unsigned char* img_data, int width, int height, int channels,
                   YOLODetection** detections, int* count) {
    try {
        auto it = yolo_detectors.find(std::string(name));
        if (it == yolo_detectors.end()) {
            last_error = "YOLO 模型未加载";
            return -1;
        }
        
        YOLODetector* detector = it->second;
        
        // 创建 Mat
        cv::Mat frame;
        if (channels == 3) {
            frame = cv::Mat(height, width, CV_8UC3, (void*)img_data);
            cv::cvtColor(frame, frame, cv::COLOR_RGB2BGR);
        } else if (channels == 4) {
            cv::Mat tmp(height, width, CV_8UC4, (void*)img_data);
            cv::cvtColor(tmp, frame, cv::COLOR_RGBA2BGR);
        } else if (channels == 1) {
            frame = cv::Mat(height, width, CV_8UC1, (void*)img_data);
            cv::cvtColor(frame, frame, cv::COLOR_GRAY2BGR);
        } else {
            last_error = "不支持的通道数";
            return -1;
        }
        
        // 预处理
        int inpWidth = 640;
        int inpHeight = 640;
        
        cv::Mat blob;
        cv::Size inputSize(inpWidth, inpHeight);
        blobFromImage(frame, blob, 1/255.0, inputSize, Scalar(0,0,0), true, false);
        
        // 推理
        detector->net.setInput(blob);
        std::vector<Mat> outs;
        detector->net.forward(outs, detector->net.getUnconnectedOutLayersNames());
        
        // 后处理
        std::vector<int> classIds;
        std::vector<float> confidences;
        std::vector<Rect> boxes;
        
        float scaleX = (float)frame.cols / inpWidth;
        float scaleY = (float)frame.rows / inpHeight;
        
        for (size_t i = 0; i < outs.size(); ++i) {
            // YOLO11 输出格式: [1, 5+num_classes, 8400] (通道在前)
            // 需要转置为 [8400, 5+num_classes]
            Mat out = outs[i];
            int num_detections = out.size[2];  // 8400
            int num_features = out.size[1];     // 5 + num_classes
            
            // 转置输出
            Mat out_transposed;
            if (out.size[1] < out.size[2]) {
                // [1, features, detections] -> [detections, features]
                Mat out_2d(out.size[1], out.size[2], CV_32F, out.data);
                transpose(out_2d, out_transposed);
            } else {
                out_transposed = out;
            }
            
            for (int j = 0; j < num_detections; ++j) {
                float* data = out_transposed.ptr<float>(j);
                
                float obj_conf = data[4];
                if (obj_conf < detector->confThreshold) continue;
                
                // 找最大类别分数
                int classId = 0;
                float maxScore = 0;
                for (int k = 5; k < num_features; ++k) {
                    if (data[k] > maxScore) {
                        maxScore = data[k];
                        classId = k - 5;
                    }
                }
                float confidence = obj_conf * maxScore;
                
                if (confidence > detector->confThreshold) {
                    float cx = data[0];
                    float cy = data[1];
                    float w = data[2];
                    float h = data[3];
                    
                    int centerX = (int)(cx * scaleX);
                    int centerY = (int)(cy * scaleY);
                    int width = (int)(w * scaleX);
                    int height = (int)(h * scaleY);
                    
                    int left = centerX - width / 2;
                    int top = centerY - height / 2;
                    
                    classIds.push_back(classId);
                    confidences.push_back(confidence);
                    boxes.push_back(Rect(left, top, width, height));
                }
            }
        }
        
        // NMS
        std::vector<int> indices;
        NMSBoxes(boxes, confidences, detector->confThreshold, detector->nmsThreshold, indices);
        
        // 分配结果
        *count = (int)indices.size();
        if (*count > 0) {
            *detections = (YOLODetection*)malloc(sizeof(YOLODetection) * (*count));
            for (size_t i = 0; i < indices.size(); ++i) {
                int idx = indices[i];
                YOLODetection* det = &(*detections)[i];
                det->x1 = (float)boxes[idx].x;
                det->y1 = (float)boxes[idx].y;
                det->x2 = (float)(boxes[idx].x + boxes[idx].width);
                det->y2 = (float)(boxes[idx].y + boxes[idx].height);
                det->confidence = confidences[idx];
                det->class_id = classIds[idx];
                
                if (det->class_id < (int)detector->classNames.size()) {
                    strncpy(det->class_name, detector->classNames[det->class_id].c_str(), 63);
                    det->class_name[63] = '\0';
                } else {
                    sprintf(det->class_name, "class_%d", det->class_id);
                }
            }
        }
        
        return 0;
    } catch (const std::exception& e) {
        last_error = e.what();
        return -1;
    }
}

// 释放 YOLO 检测结果
void cv_yolo_detections_free(YOLODetection* detections) {
    if (detections) {
        free(detections);
    }
}

// 释放 YOLO 检测器
void cv_yolo_unload(const char* name) {
    auto it = yolo_detectors.find(std::string(name));
    if (it != yolo_detectors.end()) {
        delete it->second;
        yolo_detectors.erase(it);
    }
}

// 滑块缺口匹配
int cv_slider_match(const Image* target, const Image* background, int* out_x) {
    try {
        cv::Mat target_mat(target->height, target->width, CV_8UC3, target->data);
        cv::Mat bg_mat(background->height, background->width, CV_8UC3, background->data);

        cv::Mat target_gray, bg_gray;
        cv::cvtColor(target_mat, target_gray, cv::COLOR_BGR2GRAY);
        cv::cvtColor(bg_mat, bg_gray, cv::COLOR_BGR2GRAY);

        // Canny 边缘检测
        cv::Mat target_edges, bg_edges;
        cv::Canny(target_gray, target_edges, 100, 200);
        cv::Canny(bg_gray, bg_edges, 100, 200);

        // 模板匹配
        cv::Mat result;
        cv::matchTemplate(bg_edges, target_edges, result, cv::TM_CCOEFF_NORMED);

        double min_val, max_val;
        cv::Point min_loc, max_loc;
        cv::minMaxLoc(result, &min_val, &max_val, &min_loc, &max_loc);

        *out_x = max_loc.x;
        return 0;
    } catch (const std::exception& e) {
        last_error = e.what();
        return -1;
    }
}

// 阴影滑块匹配
int cv_slider_comparison(const Image* target, const Image* background, int* out_x, int* out_y) {
    try {
        cv::Mat target_mat(target->height, target->width, CV_8UC3, target->data);
        cv::Mat bg_mat(background->height, background->width, CV_8UC3, background->data);

        // 计算差异
        cv::Mat diff;
        cv::absdiff(bg_mat, target_mat, diff);
        
        // 阈值处理
        cv::Mat thresh;
        cv::threshold(diff, thresh, 30, 255, cv::THRESH_BINARY);
        
        // 找到差异区域
        std::vector<std::vector<cv::Point>> contours;
        cv::findContours(thresh, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
        
        if (contours.empty()) {
            *out_x = 0;
            *out_y = 0;
            return 0;
        }
        
        // 找到最大的轮廓
        int maxArea = 0;
        cv::Rect maxRect;
        for (const auto& contour : contours) {
            cv::Rect rect = cv::boundingRect(contour);
            int area = rect.width * rect.height;
            if (area > maxArea) {
                maxArea = area;
                maxRect = rect;
            }
        }
        
        *out_x = maxRect.x;
        *out_y = maxRect.y;
        return 0;
    } catch (const std::exception& e) {
        last_error = e.what();
        return -1;
    }
}

// 检测旋转角度（简化版，使用特征点）
float cv_detect_rotation(const Image* img) {
    try {
        cv::Mat mat(img->height, img->width, CV_8UC3, img->data);
        
        // 转灰度
        cv::Mat gray;
        cv::cvtColor(mat, gray, cv::COLOR_BGR2GRAY);

        // 简化处理：返回 0 度
        // 实际实现需要更复杂的特征点匹配
        return 0.0f;
    } catch (const std::exception& e) {
        last_error = e.what();
        return 0.0f;
    }
}

// 图像相似度比较
float cv_compare_similarity(const Image* img1, const Image* img2) {
    try {
        cv::Mat mat1(img1->height, img1->width, CV_8UC3, img1->data);
        cv::Mat mat2(img2->height, img2->width, CV_8UC3, img2->data);

        // 确保 same size
        if (mat1.size() != mat2.size()) {
            cv::resize(mat2, mat2, mat1.size());
        }

        // 计算 histogram
        cv::Mat hsv1, hsv2;
        cv::cvtColor(mat1, hsv1, cv::COLOR_BGR2HSV);
        cv::cvtColor(mat2, hsv2, cv::COLOR_BGR2HSV);

        int h_bins = 50, s_bins = 60;
        int histSize[] = {h_bins, s_bins};
        float h_ranges[] = {0, 180};
        float s_ranges[] = {0, 256};
        const float* ranges[] = {h_ranges, s_ranges};
        int channels[] = {0, 1};

        cv::Mat hist1, hist2;
        cv::calcHist(&hsv1, 1, channels, cv::Mat(), hist1, 2, histSize, ranges);
        cv::calcHist(&hsv2, 1, channels, cv::Mat(), hist2, 2, histSize, ranges);

        cv::normalize(hist1, hist1, 0, 1, cv::NORM_MINMAX);
        cv::normalize(hist2, hist2, 0, 1, cv::NORM_MINMAX);

        double similarity = cv::compareHist(hist1, hist2, cv::HISTCMP_CORREL);
        return (float)similarity;
    } catch (const std::exception& e) {
        last_error = e.what();
        return 0.0f;
    }
}

const char* cv_get_last_error() {
    return last_error.c_str();
}

} // extern "C"