#include <iostream>
#include <opencv2/opencv.hpp>
#include "../Yolo11_ONNX/Yolo_ONNX.h" // 包含DLL的头文件

int main()
{
    // --- 1. 准备输入和参数 ---
    const wchar_t* model_path = L"D:/dev/models/best.onnx";
    cv::Mat image = cv::imread("D:/dev/dataset/003.jpg");

    if (image.empty()) {
        std::cerr << "Error: Could not read the image." << std::endl;
        return -1;
    }

    // 【新增】将所有可配置参数定义在这里，方便修改
    const float conf_threshold = 0.9f;
    const float iou_threshold = 0.45f;
    const int input_width = 2016;
    const int input_height = 1536;

    const char* class_names[] = { /* ... 您的类别列表 ... */ "tiaojuan", "zhujiesi", "yulingwen" }; // 示例
    int class_count = sizeof(class_names) / sizeof(class_names[0]);

    // --- 2. 调用更新后的DLL函数 ---
    Detection* detections = nullptr;
    int detections_count = 0;

    std::cout << "Performing detection with conf=" << conf_threshold << ", iou=" << iou_threshold << ", size=" << input_width << "x" << input_height << std::endl;

    int result = perform_detection(
        model_path,
        image.data,
        image.cols,
        image.rows,
        &detections,
        &detections_count,
        class_names,
        class_count,
        conf_threshold,
        iou_threshold,
        input_width,
        input_height
    );

    if (result != 0) {
        std::cerr << "Detection failed with code: " << result << std::endl;
        free_memory(detections);
        return -1;
    }

    std::cout << "Detection successful. Found " << detections_count << " objects." << std::endl;

    // --- 3. 打印检测结果 ---
    for (int i = 0; i < detections_count; ++i) {
        const auto& d = detections[i];
        std::cout << " - Class: " << class_names[d.class_id]
            << ", Score: " << d.score
            << ", Box: [" << d.x << ", " << d.y << ", " << d.width << ", " << d.height << "]" << std::endl;
    }

    // --- 4. 调用DLL函数绘制结果并显示 ---
    unsigned char* output_image_bytes = nullptr;
    int output_image_size = 0;

    draw_and_encode_image(
        image.data,
        image.cols,
        image.rows,
        detections,
        detections_count,
        class_names,
        class_count,
        &output_image_bytes,
        &output_image_size
    );

    if (output_image_bytes && output_image_size > 0) {
        std::vector<unsigned char> buffer(output_image_bytes, output_image_bytes + output_image_size);
        cv::Mat result_image = cv::imdecode(buffer, cv::IMREAD_COLOR);

        // 定义要保存的文件名
        std::string output_filename = "detection_result.jpg";

        // 使用 OpenCV 的 imwrite 函数将图片保存到硬盘
        bool success = cv::imwrite(output_filename, result_image);

        // 检查是否保存成功并打印提示信息
        if (success) {
            std::cout << "Annotated image successfully saved to: " << output_filename << std::endl;
        }
        else {
            std::cerr << "Error: Failed to save the annotated image." << std::endl;
        }
    }


    // --- 5. 释放内存 ---
    std::cout << "Freeing memory..." << std::endl;
    free_memory(detections);
    free_image_memory(output_image_bytes);

    std::cout << "Done." << std::endl;

    return 0;
}