#include <iostream>
#include <opencv2/opencv.hpp>
#include <vector> // 
#include "../Yolo11_ONNX/Yolo_ONNX.h" // 

int main()
{
    const wchar_t* model_path = L"D:/dev/models/best_fixed_dim.onnx";

    // 
    // 
    std::vector<std::string> image_paths = {

        "D:/dev/dataset/qd/1.jpg",
        "D:/dev/dataset/qd/2.jpg",
        "D:/dev/dataset/qd/3.jpg",
        "D:/dev/dataset/qd/4.jpg",
        "D:/dev/dataset/qd/5.jpg"
    };

    const float conf_threshold = 0.15f;
    const float iou_threshold = 0.45f;
    const int input_width = 640;
    const int input_height = 640;

    const char* class_names[] = { "qd", "fl", "zw" };
    int class_count = sizeof(class_names) / sizeof(class_names[0]);


    std::cout << "Loading model... (This happens only once)" << std::endl;
    void* detector_handle = create_detector(model_path, input_width, input_height);

    if (detector_handle == nullptr) {
        std::cerr << "Error: Failed to create detector." << std::endl;
        return -1;
    }
    std::cout << "Model loaded successfully." << std::endl;

    for (const auto& image_path : image_paths)
    {
        std::cout << "\n--- Processing image: " << image_path << " ---" << std::endl;
        cv::Mat image = cv::imread(image_path);

        if (image.empty()) {
            std::cerr << "Error: Could not read image " << image_path << std::endl;
            continue; // 
        }

        // 
        Detection* detections = nullptr;
        int detections_count = 0;

        // 
        int result = perform_detection_on_session(
            detector_handle,    // 
            image.data,
            image.cols,
            image.rows,
            &detections,
            &detections_count,
            conf_threshold,
            iou_threshold
        );

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

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

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

        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 = "result_" + image_path.substr(image_path.find_last_of('/') + 1);
            cv::imwrite(output_filename, result_image);
            std::cout << "Annotated image saved to: " << output_filename << std::endl;
        }

        // --- 5. [修改] 
        // 
        free_memory(detections);
        free_image_memory(output_image_bytes);

    } // 

    // --- 6. [新增] 
    std::cout << "\nFreeing detector model..." << std::endl;
    free_detector(detector_handle);

    std::cout << "Done." << std::endl;
    return 0;
}