#include "acoustic_analyzer/core/gunshot_classifier.h"
#include <cmath>
#include <fstream>
#include <sstream>
#include <iostream>
#include <algorithm>
#include <numeric>
#include <memory>

#ifdef _WIN32
#include <windows.h>
#endif

#include <onnxruntime_cxx_api.h>

namespace acoustic {

struct GunshotClassifier::Impl {
    ClassifierConfig config;
    std::vector<std::string> labels;
    bool loaded = false;

    Ort::Env env;
    std::unique_ptr<Ort::Session> session;
    std::unique_ptr<Ort::SessionOptions> session_options;

    explicit Impl(const ClassifierConfig& cfg)
        : config(cfg), env(ORT_LOGGING_LEVEL_WARNING, "acoustic") {
        session_options = std::make_unique<Ort::SessionOptions>();
        session_options->SetIntraOpNumThreads(1);
        session_options->SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
    }

    bool LoadLabelMap(const std::string& path) {
        std::ifstream f(path);
        if (!f) return false;
        std::string content((std::istreambuf_iterator<char>(f)),
                             std::istreambuf_iterator<char>());
        labels.clear();

        // Try simple JSON object format: {"0": "ambient", "1": "threat"}
        if (content.find('{') != std::string::npos) {
            std::size_t pos = 0;
            while ((pos = content.find('"', pos)) != std::string::npos) {
                std::size_t key_start = pos + 1;
                std::size_t key_end = content.find('"', key_start);
                if (key_end == std::string::npos) break;
                std::string key = content.substr(key_start, key_end - key_start);
                int idx = -1;
                try { idx = std::stoi(key); } catch (...) {}

                std::size_t colon = content.find(':', key_end);
                if (colon == std::string::npos) break;
                std::size_t val_quote = content.find('"', colon);
                if (val_quote == std::string::npos) break;
                std::size_t val_end = content.find('"', val_quote + 1);
                if (val_end == std::string::npos) break;
                std::string label = content.substr(val_quote + 1, val_end - val_quote - 1);

                if (idx >= 0) {
                    if (idx >= static_cast<int>(labels.size())) labels.resize(idx + 1);
                    labels[idx] = label;
                }
                pos = val_end + 1;
            }
        } else {
            // Fallback: "index: label" line format
            std::size_t line_start = 0;
            while (line_start < content.size()) {
                std::size_t line_end = content.find('\n', line_start);
                if (line_end == std::string::npos) line_end = content.size();
                std::string line = content.substr(line_start, line_end - line_start);
                line_start = line_end + 1;

                std::size_t pos = line.find(':');
                if (pos == std::string::npos) continue;
                int idx = 0;
                try {
                    idx = std::stoi(line.substr(0, pos));
                } catch (...) {
                    continue;
                }
                std::string label = line.substr(pos + 1);
                std::size_t start = label.find_first_not_of(" \t\r\n");
                std::size_t end = label.find_last_not_of(" \t\r\n");
                if (start != std::string::npos && end != std::string::npos) {
                    label = label.substr(start, end - start + 1);
                } else {
                    label.clear();
                }
                if (idx >= static_cast<int>(labels.size())) {
                    labels.resize(idx + 1);
                }
                labels[idx] = label;
            }
        }

        // Remove empty entries and compact
        std::vector<std::string> cleaned;
        for (const auto& l : labels) {
            if (!l.empty()) cleaned.push_back(l);
        }
        labels = std::move(cleaned);
        return !labels.empty();
    }

    bool LoadModel(const std::string& path) {
        // Verify file exists before attempting ONNX load
        std::ifstream check(path, std::ios::binary);
        if (!check) {
            std::cerr << "[ERROR] Model file not found: " << path << std::endl;
            return false;
        }
        check.close();

#ifdef _WIN32
        int wlen = MultiByteToWideChar(CP_UTF8, 0, path.c_str(), -1, nullptr, 0);
        std::wstring wpath;
        if (wlen > 0) {
            wpath.resize(static_cast<std::size_t>(wlen));
            MultiByteToWideChar(CP_UTF8, 0, path.c_str(), -1, &wpath[0], wlen);
        }
        try {
            session = std::make_unique<Ort::Session>(env, wpath.c_str(), *session_options);
        } catch (const Ort::Exception& e) {
            std::cerr << "[ERROR] ONNX Runtime failed to load model: " << e.what() << std::endl;
            return false;
        }
#else
        try {
            session = std::make_unique<Ort::Session>(env, path.c_str(), *session_options);
        } catch (const Ort::Exception& e) {
            std::cerr << "[ERROR] ONNX Runtime failed to load model: " << e.what() << std::endl;
            return false;
        }
#endif
        return true;
    }

    bool Initialize() {
        if (!config.label_map_path.empty()) {
            if (!LoadLabelMap(config.label_map_path)) {
                // Continue even if label map fails; may have embedded labels
            }
        }
        if (!config.model_path.empty()) {
            if (!LoadModel(config.model_path)) {
                return false;
            }
        }
        loaded = session != nullptr;
        return loaded;
    }

    std::pair<std::string, float> Predict(const Eigen::MatrixXf& mel_spectrogram) {
        if (!loaded || session == nullptr) return {"unknown", 0.0f};

        int n_mels = static_cast<int>(mel_spectrogram.rows());
        int n_frames = static_cast<int>(mel_spectrogram.cols());

        // Flatten to std::vector<float> with layout [batch, time_frames, n_mels]
        std::vector<float> input_data(static_cast<std::size_t>(n_mels) * n_frames);
        for (int t = 0; t < n_frames; ++t) {
            for (int m = 0; m < n_mels; ++m) {
                input_data[t * n_mels + m] = mel_spectrogram(m, t);
            }
        }

        // Cached memory info and I/O names (created once on first use)
        static Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
        static const char* input_names[] = {"input"};
        static const char* output_names[] = {"output"};

        std::vector<int64_t> input_shape = {1, n_frames, n_mels};
        Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
            memory_info, input_data.data(), input_data.size(),
            input_shape.data(), input_shape.size());

        // Run inference
        auto output_tensors = session->Run(
            Ort::RunOptions{nullptr},
            input_names, &input_tensor, 1,
            output_names, 1);

        // Get output
        float* output_data = output_tensors.front().GetTensorMutableData<float>();
        auto output_shape = output_tensors.front().GetTensorTypeAndShapeInfo().GetShape();
        std::size_t num_classes = labels.empty() ? static_cast<std::size_t>(output_shape.back()) : labels.size();
        if (num_classes == 0) num_classes = 2;

        // Softmax
        std::vector<float> probs(num_classes);
        float max_val = output_data[0];
        for (std::size_t i = 1; i < num_classes; ++i) {
            if (output_data[i] > max_val) max_val = output_data[i];
        }
        float sum = 0.0f;
        for (std::size_t i = 0; i < num_classes; ++i) {
            probs[i] = std::exp(output_data[i] - max_val);
            sum += probs[i];
        }
        for (std::size_t i = 0; i < num_classes; ++i) {
            probs[i] /= sum;
        }

        // Find best
        std::size_t best_idx = 0;
        for (std::size_t i = 1; i < num_classes; ++i) {
            if (probs[i] > probs[best_idx]) best_idx = i;
        }

        if (best_idx >= labels.size()) return {"unknown", probs[best_idx]};
        return {labels[best_idx], probs[best_idx]};
    }
};

GunshotClassifier::GunshotClassifier(const ClassifierConfig& config)
    : impl_(std::make_unique<Impl>(config)) {
    impl_->Initialize();
}

GunshotClassifier::~GunshotClassifier() = default;
GunshotClassifier::GunshotClassifier(GunshotClassifier&&) noexcept = default;
GunshotClassifier& GunshotClassifier::operator=(GunshotClassifier&&) noexcept = default;

std::pair<std::string, float> GunshotClassifier::Predict(const Eigen::MatrixXf& mel_spectrogram) {
    return impl_->Predict(mel_spectrogram);
}

const std::vector<std::string>& GunshotClassifier::Labels() const noexcept {
    return impl_->labels;
}

bool GunshotClassifier::IsLoaded() const noexcept {
    return impl_->loaded;
}

}  // namespace acoustic
