#include <iostream>
#include <vector>
#include <string>
#include <cmath>
#include <cstring>
#include "acoustic_analyzer/core/gunshot_classifier.h"
#include "acoustic_analyzer/core/feature_extractor.h"
#include "acoustic_analyzer/io/wav_file_source.h"

using namespace acoustic;

void print_usage(const char* prog) {
    std::cerr << "Usage: " << prog << " --model <onnx> --wav <wav_file> [--label_map <json>]" << std::endl;
}

int main(int argc, char** argv) {
    std::string model_path;
    std::string wav_path;
    std::string label_map_path;

    for (int i = 1; i < argc; ++i) {
        if (std::strcmp(argv[i], "--model") == 0 && i + 1 < argc) {
            model_path = argv[++i];
        } else if (std::strcmp(argv[i], "--wav") == 0 && i + 1 < argc) {
            wav_path = argv[++i];
        } else if (std::strcmp(argv[i], "--label_map") == 0 && i + 1 < argc) {
            label_map_path = argv[++i];
        }
    }

    if (model_path.empty() || wav_path.empty()) {
        print_usage(argv[0]);
        return 1;
    }

    WavFileSource wav(wav_path);
    if (!wav.open()) {
        std::cerr << "Failed to open WAV: " << wav_path << std::endl;
        return 1;
    }

    std::vector<std::vector<float>> audio;
    size_t n = wav.read(audio, wav.total_samples());
    if (n == 0 || audio.empty()) {
        std::cerr << "Empty WAV file" << std::endl;
        return 1;
    }

    FeatureExtractorParams feat_p;
    feat_p.sample_rate = wav.sample_rate();
    FeatureExtractor extractor(feat_p);
    std::vector<float> features;
    if (!extractor.extract(audio, features)) {
        std::cerr << "Feature extraction failed" << std::endl;
        return 1;
    }

    GunshotClassifier classifier(model_path, label_map_path, 0.7f, 3);
    if (!classifier.init()) {
        std::cerr << "Classifier init failed" << std::endl;
        return 1;
    }

    ClassificationResult result;
    if (!classifier.classify(features, result)) {
        std::cerr << "Classification failed" << std::endl;
        return 1;
    }

    std::cout << "Label: " << result.label << std::endl;
    std::cout << "Confidence: " << result.confidence << std::endl;
    for (const auto& kv : result.probs) {
        std::cout << "  " << kv.first << ": " << kv.second << std::endl;
    }
    return 0;
}
