#include "acoustic_analyzer/ros/acoustic_node.h"
#include "acoustic_analyzer/core/pipeline.h"
#include "acoustic_analyzer/io/wav_file_source.h"
#ifndef _WIN32
#include "acoustic_analyzer/io/mobile_phone_source.h"
#endif
#include <ros/ros.h>
#include <ros/package.h>
#include <std_msgs/Float32MultiArray.h>
#include <yaml-cpp/yaml.h>
#include <memory>
#include <vector>
#include <string>
#include <regex>

namespace acoustic {

class AcousticNode {
public:
    AcousticNode(ros::NodeHandle& nh, ros::NodeHandle& pnh, ThreatPublisher* publisher)
        : nh_(nh), pnh_(pnh), threat_publisher_(publisher) {
        load_params();
        init_pipeline();
        init_source();
    }

    void run() {
        ros::Rate rate(params_.publish_rate);
        while (ros::ok()) {
            ros::spinOnce();
            if (source_type_ == "wav_file" && wav_source_) {
                process_wav_source();
            }
            rate.sleep();
        }
    }

private:
    ros::NodeHandle nh_;
    ros::NodeHandle pnh_;
    ros::Subscriber audio_sub_;
    ThreatPublisher* threat_publisher_ = nullptr;

    std::string source_type_;
    std::unique_ptr<Pipeline> pipeline_;
    std::unique_ptr<WavFileSource> wav_source_;
#ifndef _WIN32
    std::unique_ptr<MobilePhoneSource> phone_source_;
#endif

    struct Params {
        int sample_rate = 16000;
        float chunk_duration = 2.0f;
        float hop_duration = 0.5f;
        int n_channels = 4;
        std::string source_type;
        std::string mobile_phone_topic;
        float mobile_phone_timeout = 10.0f;
        std::string wav_file_path;
        std::string model_path;
        std::string label_map_path;
        float publish_rate = 10.0f;
    } params_;

    /**
     * @brief 展开 YAML 中的 $(find package_name) 语法为绝对路径
     */
    static std::string expandFindPath(const std::string& path) {
        static const std::regex findRegex(R"(\$\(find\s+([^)]+)\))");
        std::smatch match;
        std::string result = path;
        if (std::regex_search(path, match, findRegex) && match.size() > 1) {
            std::string pkgPath = ros::package::getPath(match[1].str());
            if (!pkgPath.empty()) {
                result = match.prefix().str() + pkgPath + match.suffix().str();
            } else {
                ROS_WARN("ros::package::getPath(%s) returned empty, keeping raw path", match[1].str().c_str());
            }
        }
        return result;
    }

    void load_params() {
        std::string yaml_path;
        pnh_.param<std::string>("config_file", yaml_path, "");
        if (!yaml_path.empty()) {
            try {
                YAML::Node config = YAML::LoadFile(yaml_path);
                params_.sample_rate = config["audio"]["sample_rate"].as<int>(16000);
                params_.chunk_duration = config["audio"]["chunk_duration"].as<float>(2.0f);
                params_.hop_duration = config["audio"]["hop_duration"].as<float>(0.5f);
                params_.n_channels = config["audio"]["n_channels"].as<int>(4);
                params_.source_type = config["source"]["type"].as<std::string>("mic_array");
                params_.mobile_phone_topic = config["source"]["mobile_phone_topic"].as<std::string>("/mobile_phone/audio");
                params_.mobile_phone_timeout = config["source"]["mobile_phone_timeout"].as<float>(10.0f);
                params_.wav_file_path = config["source"]["wav_file_path"].as<std::string>("");
                params_.model_path = expandFindPath(config["classifier"]["model_path"].as<std::string>(""));
                params_.label_map_path = expandFindPath(config["classifier"]["label_map_path"].as<std::string>(""));
                params_.publish_rate = config["output"]["publish_rate"].as<float>(10.0f);
            } catch (const std::exception& e) {
                ROS_WARN("Failed to load YAML config: %s", e.what());
            }
        }
        source_type_ = params_.source_type;
    }

    void init_pipeline() {
        PipelineConfig config;
        config.sample_rate = static_cast<uint32_t>(params_.sample_rate);
        config.chunk_duration = params_.chunk_duration;
        config.hop_duration = params_.hop_duration;
        config.n_mels = 64;
        config.classifier.model_path = params_.model_path;
        config.classifier.label_map_path = params_.label_map_path;
        config.classifier.threshold = 0.7f;
        config.classifier.smoothing_window = 3;
        config.mic_array.num_mics = static_cast<uint32_t>(params_.n_channels);
        config.mic_array.layout = "cross";
        config.mic_array.spacing = 0.15f;
        pipeline_ = std::make_unique<Pipeline>(config);
    }

    void init_source() {
        if (source_type_ == "mic_array") {
            audio_sub_ = nh_.subscribe("/microphone_array/audio", 10,
                                       &AcousticNode::on_mic_array_audio, this);
        } else if (source_type_ == "mobile_phone") {
            audio_sub_ = nh_.subscribe(params_.mobile_phone_topic, 10,
                                       &AcousticNode::on_mobile_phone_audio, this);
        } else if (source_type_ == "wav_file") {
            if (!params_.wav_file_path.empty()) {
                wav_source_ = std::make_unique<WavFileSource>(params_.wav_file_path, params_.sample_rate);
                if (!wav_source_->open()) {
                    ROS_ERROR("Failed to open WAV file: %s", params_.wav_file_path.c_str());
                }
            }
        }
    }

    // Convert multi-channel vector-of-vectors to flat interleaved format
    std::vector<float> flatten_audio(const std::vector<std::vector<float>>& audio, int channels) {
        if (audio.empty() || channels == 0) return {};
        size_t samples = audio[0].size();
        std::vector<float> flat(samples * channels);
        for (size_t s = 0; s < samples; ++s) {
            for (int ch = 0; ch < channels; ++ch) {
                flat[s * channels + ch] = (ch < static_cast<int>(audio.size()) && s < audio[ch].size())
                                            ? audio[ch][s] : 0.0f;
            }
        }
        return flat;
    }

    void publish_if_available(const AcousticFrame& frame) {
        if (threat_publisher_ && !frame.is_clear) {
            threat_publisher_->Publish(frame);
        }
    }

    void on_mic_array_audio(const std_msgs::Float32MultiArray::ConstPtr& msg) {
        if (msg->layout.dim.size() < 2) return;
        int channels = static_cast<int>(msg->layout.dim[0].size);
        int samples = static_cast<int>(msg->layout.dim[1].size);
        if (channels == 0 || samples == 0) return;

        // Assuming data is interleaved or [channels x samples] row-major
        std::vector<float> flat(msg->data.begin(), msg->data.end());
        auto frame = pipeline_->Process(flat);
        publish_if_available(frame);
    }

    void on_mobile_phone_audio(const std_msgs::Float32MultiArray::ConstPtr& msg) {
        if (msg->data.empty()) return;
        std::vector<float> flat(msg->data.begin(), msg->data.end());
        auto frame = pipeline_->Process(flat);
        publish_if_available(frame);
    }

    void process_wav_source() {
        if (!wav_source_ || !wav_source_->is_open()) return;
        size_t chunk_samples = static_cast<size_t>(params_.chunk_duration * params_.sample_rate);
        std::vector<std::vector<float>> audio;
        size_t got = wav_source_->read(audio, chunk_samples);
        if (got == 0) {
            ROS_INFO("WAV playback finished");
            ros::shutdown();
            return;
        }
        auto flat = flatten_audio(audio, static_cast<int>(wav_source_->num_channels()));
        auto frame = pipeline_->Process(flat);
        publish_if_available(frame);
    }
};

std::unique_ptr<AcousticNode> create_acoustic_node(
    ros::NodeHandle& nh,
    ros::NodeHandle& pnh,
    ThreatPublisher* publisher) {
    return std::make_unique<AcousticNode>(nh, pnh, publisher);
}

} // namespace acoustic
