#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 <std_msgs/Float32MultiArray.h>
#include <yaml-cpp/yaml.h>
#include <memory>
#include <vector>
#include <string>

namespace acoustic {

class AcousticNode {
public:
    AcousticNode(ros::NodeHandle& nh, ros::NodeHandle& pnh) : nh_(nh), pnh_(pnh) {
        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_;

    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_;

    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 = config["classifier"]["model_path"].as<std::string>("");
                params_.label_map_path = 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() {
        PipelineParams p;
        p.feature.sample_rate = params_.sample_rate;
        p.feature.n_mels = 64;
        p.feature.n_fft = 2048;
        p.feature.hop_length = 512;
        p.feature.f_min = 0.0f;
        p.feature.f_max = 8000.0f;
        p.feature.preemphasis = 0.97f;
        p.model_path = params_.model_path;
        p.label_map_path = params_.label_map_path;
        p.sample_rate = params_.sample_rate;
        p.mic_geometry.num_mics = params_.n_channels;
        p.mic_geometry.layout = "cross";
        p.mic_geometry.spacing = 0.15f;
        pipeline_ = std::make_unique<Pipeline>(p);
        if (!pipeline_->init()) {
            ROS_ERROR("Pipeline initialization failed");
        }
    }

    void init_source() {
        if (source_type_ == "mic_array") {
            audio_sub_ = nh_.subscribe(params_.mobile_phone_topic.empty()
                                        ? "/microphone_array/audio"
                                        : "/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());
                }
            }
        }
    }

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

        std::vector<std::vector<float>> audio(channels, std::vector<float>(samples));
        for (int ch = 0; ch < channels; ++ch) {
            for (int s = 0; s < samples; ++s) {
                audio[ch][s] = msg->data[ch * samples + s];
            }
        }

        std::vector<TrackedThreat> threats;
        if (channels == 1) {
            pipeline_->process_single_channel(audio[0], threats);
        } else {
            pipeline_->process(audio, threats);
        }
        // Publish threats via threat_publisher (would be called by main loop)
    }

    void on_mobile_phone_audio(const std_msgs::Float32MultiArray::ConstPtr& msg) {
        if (msg->data.empty()) return;
        std::vector<TrackedThreat> threats;
        pipeline_->process_single_channel(msg->data, threats);
    }

    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;
        }
        std::vector<TrackedThreat> threats;
        if (wav_source_->num_channels() == 1) {
            pipeline_->process_single_channel(audio[0], threats);
        } else {
            pipeline_->process(audio, threats);
        }
    }
};

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

} // namespace acoustic
