#!/usr/bin/env python3
"""
ASR语音转文字ROS2节点

功能：
1. 使用阿里云ASR API进行语音识别
2. 接收麦克风音频输入并转换为文字
3. 处理识别结果（检查退出词、休眠词等）
4. 发布文字话题
"""

import os
import time
import rclpy
from rclpy.node import Node
from std_msgs.msg import String
import sys
import traceback
import threading
import dashscope
from dashscope.audio.asr import Recognition, RecognitionCallback, RecognitionResult
import pyaudio

class WakeWordDetector:
    """语音唤醒检测器"""
    
    def __init__(self):
        self.wake_words = ["优宝同学","又把同学","有保同学","你好优宝"]
        self.sleep_words = ["拜拜", "再见", "休息吧", "睡觉吧"]
        self.exit_words = ["关闭程序", "退出程序", "彻底退出"]
        
        try:
            from pypinyin import lazy_pinyin
            self.wake_words_pinyin = {
                word: ''.join(lazy_pinyin(word)) 
                for word in self.wake_words
            }
        except ImportError:
            self.wake_words_pinyin = None
    
    def check_wake_word(self, text):
        """检查是否包含唤醒词（支持拼音模糊匹配）"""
        if not text:
            return False
        
        for wake_word in self.wake_words:
            if wake_word in text:
                return True
        
        if self.wake_words_pinyin:
            try:
                from pypinyin import lazy_pinyin
                text_pinyin = ''.join(lazy_pinyin(text))
                
                for wake_word, wake_pinyin in self.wake_words_pinyin.items():
                    if wake_pinyin in text_pinyin:
                        return True
                    
                    similarity = self._calculate_similarity(wake_pinyin, text_pinyin)
                    if similarity > 0.7:
                        return True
                        
            except Exception:
                pass
        
        return False
    
    def check_sleep_word(self, text):
        """检查是否包含休眠词（支持拼音匹配）"""
        if not text:
            return False
        
        for sleep_word in self.sleep_words:
            if sleep_word in text:
                return True
        
        if self.wake_words_pinyin:
            try:
                from pypinyin import lazy_pinyin
                text_pinyin = ''.join(lazy_pinyin(text))
                for sleep_word in self.sleep_words:
                    sleep_pinyin = ''.join(lazy_pinyin(sleep_word))
                    if sleep_pinyin in text_pinyin:
                        return True
            except:
                pass
        
        return False
    
    def check_exit_word(self, text):
        """检查是否包含退出程序词（支持拼音匹配）"""
        if not text:
            return False
        
        for exit_word in self.exit_words:
            if exit_word in text:
                return True
        
        if self.wake_words_pinyin:
            try:
                from pypinyin import lazy_pinyin
                text_pinyin = ''.join(lazy_pinyin(text))
                for exit_word in self.exit_words:
                    exit_pinyin = ''.join(lazy_pinyin(exit_word))
                    if exit_pinyin in text_pinyin:
                        return True
            except:
                pass
        
        return False
    
    def _calculate_similarity(self, s1, s2):
        """计算两个拼音字符串的相似度（编辑距离）"""
        if not s1 or not s2:
            return 0.0
        
        len1, len2 = len(s1), len(s2)
        dp = [[0] * (len2 + 1) for _ in range(len1 + 1)]
        
        for i in range(len1 + 1):
            dp[i][0] = i
        for j in range(len2 + 1):
            dp[0][j] = j
        
        for i in range(1, len1 + 1):
            for j in range(1, len2 + 1):
                if s1[i-1] == s2[j-1]:
                    dp[i][j] = dp[i-1][j-1]
                else:
                    dp[i][j] = min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1]) + 1
        
        distance = dp[len1][len2]
        max_len = max(len1, len2)
        similarity = 1 - (distance / max_len) if max_len > 0 else 0
        
        return similarity

class ASRRecognitionCallback(RecognitionCallback):
    """阿里云ASR识别回调类"""
    
    def __init__(self, node):
        super().__init__()
        self.node = node
        self.current_text = ""
    
    def on_open(self) -> None:
        self.node.get_logger().info('ASR连接已建立')
    
    def on_close(self) -> None:
        self.node.get_logger().info('ASR连接已关闭')
    
    def on_complete(self) -> None:
        self.node.get_logger().info('ASR识别完成')
    
    def on_error(self, message) -> None:
        self.node.get_logger().error(f'ASR识别错误: {message.message}')
    
    def on_event(self, result) -> None:
        sentence = result.get_sentence()
        if 'text' in sentence:
            text = sentence['text']
            if text:
                self.node.get_logger().info(f'识别文本: {text}')
                
                if RecognitionResult.is_sentence_end(sentence):
                    self.node._process_recognition_result(text)

class ASRNode(Node):
    """ASR语音转文字ROS2节点"""
    
    def __init__(self):
        super().__init__('asr_node')
        
        self.declare_parameter('api_key', 'sk-c37ab757225d4891968c9e9963b61d69')
        self.declare_parameter('sample_rate', 16000)
        self.declare_parameter('channels', 1)
        self.declare_parameter('chunk_size', 3200)
        
        self.api_key = self.get_parameter('api_key').value
        self.sample_rate = self.get_parameter('sample_rate').value
        self.channels = self.get_parameter('channels').value
        self.chunk_size = self.get_parameter('chunk_size').value
        
        dashscope.api_key = self.api_key
        
        self.text_publisher = self.create_publisher(String, 'asr/text', 10)
        self.command_publisher = self.create_publisher(String, 'asr/command', 10)
        
        self.wake_detector = WakeWordDetector()
        
        self.recognition = None
        self.callback = None
        self.pyaudio = None
        self.stream = None
        self.is_running = False
        self.recognition_thread = None
        
        self._init_asr()
        
        self.get_logger().info('ASR语音转文字节点初始化完成')
    
    def _init_asr(self):
        """初始化ASR识别"""
        try:
            self.callback = ASRRecognitionCallback(self)
            
            self.recognition = Recognition(
                model='paraformer-realtime-v2',
                format='pcm',
                sample_rate=self.sample_rate,
                language_hints=['zh', 'en'],
                callback=self.callback
            )
            
            self.recognition.start()
            
            self.pyaudio = pyaudio.PyAudio()
            self.stream = self.pyaudio.open(
                format=pyaudio.paInt16,
                channels=self.channels,
                rate=self.sample_rate,
                input=True,
                frames_per_buffer=self.chunk_size
            )
            
            self.is_running = True
            self.recognition_thread = threading.Thread(target=self._recognition_loop, daemon=True)
            self.recognition_thread.start()
            
            self.get_logger().info('ASR识别服务已启动')
        except Exception as e:
            self.get_logger().error(f'初始化ASR失败: {e}')
            traceback.print_exc()
    
    def _recognition_loop(self):
        """识别循环，持续读取音频并发送给ASR服务"""
        while self.is_running and self.stream and self.stream.is_active():
            try:
                data = self.stream.read(self.chunk_size, exception_on_overflow=False)
                if self.recognition:
                    self.recognition.send_audio_frame(data)
            except Exception as e:
                self.get_logger().error(f'读取音频数据失败: {e}')
                break
    
    def _process_recognition_result(self, text):
        """处理识别结果"""
        if not text:
            return
        
        if self.wake_detector.check_exit_word(text):
            self.get_logger().info('检测到退出程序指令')
            msg = String()
            msg.data = 'exit'
            self.command_publisher.publish(msg)
            return
        
        if self.wake_detector.check_sleep_word(text):
            self.get_logger().info('检测到休眠指令')
            msg = String()
            msg.data = 'sleep'
            self.command_publisher.publish(msg)
            return
        
        if self.wake_detector.check_wake_word(text):
            self.get_logger().info('检测到唤醒词')
            original_text = text
            for wake_word in self.wake_detector.wake_words:
                text = text.replace(wake_word, "").strip()
            
            text = text.strip('，。！？、；：,.!?;: ')
            
            self.get_logger().info(f'去除唤醒词后: {text}')
        
        if text:
            msg = String()
            msg.data = text
            self.text_publisher.publish(msg)
            self.get_logger().info(f'发布ASR识别结果: {text}')
    
    def destroy_node(self):
        """销毁节点，释放资源"""
        self.is_running = False
        
        if self.recognition_thread and self.recognition_thread.is_alive():
            self.recognition_thread.join(timeout=2)
        
        if self.stream:
            try:
                self.stream.stop_stream()
                self.stream.close()
            except:
                pass
        
        if self.pyaudio:
            try:
                self.pyaudio.terminate()
            except:
                pass
        
        if self.recognition:
            try:
                self.recognition.stop()
            except:
                pass
        
        super().destroy_node()

def main(args=None):
    """主函数"""
    rclpy.init(args=args)
    
    node = ASRNode()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
