#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <iostream>
#include <vector>
#include <string>
#include <filesystem>
#include "../common.h"
#include "extract_seal_main.h"

using namespace std;
namespace fs = std::filesystem;

struct ImageTask {
    string img_path;
    vector<vector<double>> det;
};

int main() {
    // TODO: 修改为你的图片目录路径
    // 示例: "c:\\Users\\用户名\\Desktop\\Final project\\image" 或使用相对路径 "..\\..\\image"
    string image_dir = "..\\..\\image";
    vector<ImageTask> tasks;
    
    cout << "Scanning image directory: " << image_dir << endl;
    
    try {
        for (const auto& entry : fs::directory_iterator(image_dir)) {
            if (entry.is_regular_file()) {
                string ext = entry.path().extension().string();
                if (ext == ".png" || ext == ".jpg" || ext == ".jpeg" || ext == ".bmp") {
                    ImageTask task;
                    task.img_path = entry.path().string();
                    
                    string filename = entry.path().filename().string();
                    if (filename == "2.png") {
                        task.det = {
                            {256.0, 384.0, 471.0, 592.0, 0.980, 2},
                            {511.0, 423.0, 665.0, 581.0, 0.969, 2}
                        };
                    } else if (filename == "1.png") {
                        task.det = {
                            {408.0, 32.0, 559.0, 186.0, 0.974, 0},
                            {267.0, 61.0, 348.0, 143.0, 0.973, 2},
                            {54.0, 36.0, 212.0, 147.0, 0.972, 1}
                        };
                    }
                    
                    tasks.push_back(task);
                    cout << "Found image: " << filename << endl;
                }
            }
        }
    } catch (const fs::filesystem_error& e) {
        cout << "Error scanning directory: " << e.what() << endl;
        return 1;
    }
    
    if (tasks.empty()) {
        cout << "No image files found in directory." << endl;
        return 1;
    }
    
    cout << "\nTotal images found: " << tasks.size() << endl;
    
    for (size_t task_idx = 0; task_idx < tasks.size(); ++task_idx) {
        const ImageTask& task = tasks[task_idx];
        string img_path = task.img_path;
        vector<vector<double>> det = task.det;
        
        cout << "\n========================================" << endl;
        cout << "Processing image " << task_idx + 1 << "/" << tasks.size() << endl;
        cout << "File: " << fs::path(img_path).filename().string() << endl;
        cout << "========================================" << endl;
        
        cv::Mat img = cv::imread(img_path);
        if (img.empty()) {
            cout << "Error: Cannot read image " << img_path << endl;
            if (task_idx < tasks.size() - 1) {
                cout << "Press Enter to continue to next image..." << endl;
                cin.ignore(numeric_limits<streamsize>::max(), '\n');
            }
            continue;
        }
        
        cout << "Image dimensions: " << img.cols << "x" << img.rows << endl;
        
        cout << "Number of detection boxes: " << det.size() << endl;
        for (size_t i = 0; i < det.size(); ++i) {
            cout << "Detection box " << i+1 << ": x1=" << det[i][0] << ", y1=" << det[i][1] 
                 << ", x2=" << det[i][2] << ", y2=" << det[i][3] 
                 << ", conf=" << det[i][4] << ", shape=" << det[i][5] << endl;
        }
        
        cout << "\nStart extracting seals..." << endl;
        vector<SealInfo> seal_set = extract_seal_main(img, det);
        
        cout << "Total extracted seals: " << seal_set.size() << endl << endl;
        
        // TODO: 修改为你的输出目录路径
        // 示例: "c:\\Users\\用户名\\Desktop\\Final project\\c\\extract_seal_main\\test_image" 或使用相对路径 "..\\extract_seal_main\\test_image"
        string output_dir = "..\\extract_seal_main\\test_image";
        fs::create_directories(output_dir);
        
        for (size_t i = 0; i < seal_set.size(); ++i) {
            cout << "========== Seal " << i << " ==========" << endl;
            
            string img_name = fs::path(img_path).stem().string();
            string cropped_filename = output_dir + "\\seal_" + img_name + "_" + to_string(i) + "_cropped.png";
            string ex_filename = output_dir + "\\seal_" + img_name + "_" + to_string(i) + "_ex_img.png";
            string adj_filename = output_dir + "\\seal_" + img_name + "_" + to_string(i) + "_adj_img.png";
            
            cv::imwrite(cropped_filename, seal_set[i].cropped_img);
            cv::imwrite(ex_filename, seal_set[i].extracted_img);
            cv::imwrite(adj_filename, seal_set[i].adj_img);
            
            cout << "Saved files:" << endl;
            cout << "  " << cropped_filename << " (cropped_img)" << endl;
            cout << "  " << ex_filename << " (extracted_img)" << endl;
            cout << "  " << adj_filename << " (adj_img)" << endl;
            
            cout << endl << "region: [" << seal_set[i].x1 << ", " << seal_set[i].y1 << ", " 
                 << seal_set[i].x2 << ", " << seal_set[i].y2 << "]" << endl;
            cout << "shape: " << seal_set[i].shape << endl;
            cout << "color_name: " << seal_set[i].color_name << endl;
            cout << "angle: " << seal_set[i].angle << endl;
            cout << "shape_info.type: " << seal_set[i].shape_info.type << endl;
            cout << "shape_info.minAreaRect: center=(" 
                 << seal_set[i].shape_info.minAreaRect.center.x << ", " 
                 << seal_set[i].shape_info.minAreaRect.center.y << "), size=(" 
                 << seal_set[i].shape_info.minAreaRect.size.width << ", " 
                 << seal_set[i].shape_info.minAreaRect.size.height << ")" << endl;
            cout << endl;
        }
        
        if (task_idx < tasks.size() - 1) {
            cout << "Press Enter to process next image..." << endl;
            cin.ignore(numeric_limits<streamsize>::max(), '\n');
        }
    }
    
    cout << "\nAll images processed!" << endl;
    return 0;
}