#include "thermal_preprocessor.h"
#include "thermal_detector.h"
#include <opencv2/imgcodecs.hpp>
#include <opencv2/highgui.hpp>
#include <iostream>
#include <chrono>
#include <cstring>

static void print_usage(const char* prog) {
    std::cout << "Usage: " << prog << " <image_path> [options]\n"
              << "       " << prog << " <directory>  [options]\n\n"
              << "Options:\n"
              << "  --no-show    Don't display results (save only)\n"
              << "  --json       Output JSON format result\n"
              << "  --help       Show this help\n\n"
              << "Examples:\n"
              << "  thermal_detector test_images/photo.jpg\n"
              << "  thermal_detector test_images/ --no-show\n";
}

static std::string threat_type_str(ThreatType t) {
    return ThreatTypeInfo::label(t);
}

static void print_assessment(const Assessment& a, const std::string& filename, double ms) {
    auto level = a.level();
    std::cout << "========================================\n";
    std::cout << " Image: " << filename << "\n";
    std::cout << "========================================\n";
    std::cout << " Time: " << ms << " ms\n";
    std::cout << " Ambient: " << a.ambient_temp << " C\n";
    std::cout << " Level: [" << ThreatLevelInfo::label(level) << "] "
              << a.detections.size() << " detection(s)\n";

    if (!a.persons.empty())
        std::cout << "   Person x" << a.persons.size() << "\n";
    if (!a.barrels.empty())
        std::cout << "   Barrel Overheat x" << a.barrels.size() << "\n";
    if (!a.cold_weapons.empty())
        std::cout << "   Cold Weapon x" << a.cold_weapons.size() << "\n";
    if (!a.flashes.empty())
        std::cout << "   Flash Residue x" << a.flashes.size() << "\n";
    if (!a.hotspots.empty())
        std::cout << "   Hotspot x" << a.hotspots.size() << "\n";

    for (size_t i = 0; i < std::min(a.detections.size(), (size_t)15); i++) {
        const auto& d = a.detections[i];
        std::cout << "    [" << threat_type_str(d.type) << "] conf:"
                  << (int)(d.confidence * 100) << "% temp:" << d.max_temp
                  << "C area:" << d.area << "px\n";
    }
    std::cout << "========================================\n\n";
}

static void print_json(const Assessment& a, const std::string& filename) {
    std::cout << "{\n";
    std::cout << "  \"file\": \"" << filename << "\",\n";
    std::cout << "  \"threat_level\": \"" << ThreatLevelInfo::label(a.level()) << "\",\n";
    std::cout << "  \"is_dangerous\": " << (a.is_dangerous() ? "true" : "false") << ",\n";
    std::cout << "  \"detection_count\": " << a.detections.size() << ",\n";
    std::cout << "  \"ambient_temp\": " << a.ambient_temp << ",\n";

    std::cout << "  \"detections\": [\n";
    for (size_t i = 0; i < a.detections.size(); i++) {
        const auto& d = a.detections[i];
        std::cout << "    {\"type\":\"" << threat_type_str(d.type)
                  << "\", \"confidence\":" << d.confidence
                  << ", \"max_temp\":" << d.max_temp
                  << ", \"bbox\":{\"x\":" << d.bbox.x
                  << ",\"y\":" << d.bbox.y
                  << ",\"w\":" << d.bbox.width
                  << ",\"h\":" << d.bbox.height << "}}";
        if (i + 1 < a.detections.size()) std::cout << ",";
        std::cout << "\n";
    }
    std::cout << "  ]\n}\n";
}

static std::string get_output_path(const std::string& input_path) {
    auto pos = input_path.find_last_of("/\\");
    std::string dir = (pos != std::string::npos) ? input_path.substr(0, pos + 1) : "";
    std::string fname = (pos != std::string::npos) ? input_path.substr(pos + 1) : input_path;
    return dir + "result_" + fname;
}

static bool has_image_ext(const std::string& path) {
    std::string ext;
    auto pos = path.rfind('.');
    if (pos != std::string::npos) ext = path.substr(pos);
    return ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp" || ext == ".tif";
}

int main(int argc, char* argv[]) {
    if (argc < 2) { print_usage(argv[0]); return 1; }

    std::string source = argv[1];
    bool show = true, json_out = false;

    for (int i = 2; i < argc; i++) {
        if (std::strcmp(argv[i], "--no-show") == 0) show = false;
        else if (std::strcmp(argv[i], "--json") == 0) json_out = true;
        else if (std::strcmp(argv[i], "--help") == 0) { print_usage(argv[0]); return 0; }
    }

    ThermalPreprocessor preprocessor;
    ThermalDetector detector;

    auto process_one = [&](const std::string& path) {
        cv::Mat img = cv::imread(path, cv::IMREAD_COLOR);
        if (img.empty())
            img = cv::imread(path, cv::IMREAD_GRAYSCALE);
        if (img.empty()) {
            std::cerr << "Cannot read: " << path << "\n";
            return;
        }

        auto t0 = std::chrono::high_resolution_clock::now();
        auto preprocessed = preprocessor.preprocess(img);
        auto assessment = detector.detect(preprocessed);
        auto t1 = std::chrono::high_resolution_clock::now();
        double ms = std::chrono::duration<double, std::milli>(t1 - t0).count();

        if (json_out) {
            print_json(assessment, path);
        } else {
            print_assessment(assessment, path, ms);
        }

        // 保存标注图
        cv::Mat vis = ThermalDetector::draw_detections(img, assessment);
        std::string out_path = get_output_path(path);
        cv::imwrite(out_path, vis);
        if (!json_out)
            std::cout << " Saved: " << out_path << "\n";

        if (show) {
            cv::imshow("Thermal Detection", vis);
            cv::waitKey(0);
            cv::destroyAllWindows();
        }
    };

    process_one(source);
    return 0;
}
