#!/home/ubuntu/miniconda3/envs/act_env/bin/python3
import os
import signal  # for keyboard events handling (press "Ctrl+C" to terminate recording)
import sys
from datetime import datetime
import subprocess

import dashscope
from dashscope.audio.asr import *


# ROS2 导入
try:
    import rclpy
    from rclpy.node import Node
    from std_msgs.msg import String, Int32, Bool
    HAS_ROS2 = True
except ImportError:
    HAS_ROS2 = False

recognition = None
audio_process = None

# 设置录制参数
sample_rate = 16000  # sampling rate (Hz)
channels = 1  # mono channel
dtype = 'int16'  # data type
format_pcm = 'pcm'  # the format of the audio data
block_size = 3200  # number of frames per buffer

# 用 PulseAudio 录音，完全跳过 ALSA / PyAudio
def audio_generator():
    # 命令：纯 PulseAudio 录音，无 ALSA
    cmd = [
        "parec",
        "--rate=16000",
        "--channels=1",
        "--format=s16le",
        "--latency=1024",
        "-"
    ]
    
    # 启动子进程
    p = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        bufsize=1024
    )
    
    # 持续读取音频
    while True:
        data = p.stdout.read(1024 * 2)
        if not data:
            break
        yield data


# 实时语音识别回调
class Callback(RecognitionCallback):
    def __init__(self, publisher=None, state_publisher=None):
        super().__init__()
        self.publisher = publisher
        self.state_publisher = state_publisher
        self.is_llm_activated = False

    def update_llm_activation(self, is_activated):
        """更新LLM激活状态"""
        self.is_llm_activated = is_activated
        print(f'LLM激活状态已更新: {is_activated}')

    def on_open(self) -> None:
        print('RecognitionCallback open.')
        # 不再使用PyAudio，改为在main函数中使用audio_generator

    def on_close(self) -> None:
        print('RecognitionCallback close.')
        # 不再使用PyAudio，资源释放由main函数处理

    def on_complete(self) -> None:
        print('RecognitionCallback completed.')  # recognition completed

    def on_error(self, message) -> None:
        print('RecognitionCallback task_id: ', message.request_id)
        print('RecognitionCallback error: ', message.message)
        # Forcefully exit the program
        sys.exit(1)

    def on_event(self, result: RecognitionResult) -> None:
        sentence = result.get_sentence()
        if 'text' in sentence:
            print('RecognitionCallback text: ', sentence['text'])

            if RecognitionResult.is_sentence_end(sentence):
                print(
                    'RecognitionCallback sentence end, request_id:%s, usage:%s'
                    % (result.get_request_id(), result.get_usage(sentence)))
                
                # 发布到 ROS2 话题
                if self.publisher:
                    msg = String()
                    # 获取当前时间戳并格式化
                    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
                    # 将时间戳和识别内容拼接
                    msg.data = f'[{timestamp}] {sentence["text"]}'
                    
                    self.publisher.publish(msg)
                    print(f'已发布到 ROS2 话题: {msg.data}')

                # 检查LLM激活状态和命令
                text = sentence['text']
                if self.is_llm_activated and self.state_publisher:
                    state_msg = Int32()
                    publish = False
                    
                    # 检查命令逻辑
                    commands = ['起来', '起立', '向前走', '前进', '左转', '右转', '向左走', '向右走', '坐下']
                    matched_command = None
                    for cmd in commands:
                        if cmd in text:
                            matched_command = cmd
                            break
                    
                    if matched_command:
                        # 使用包含匹配，只要语音中包含关键词就执行
                        if '起来' in matched_command or '起立' in matched_command:
                            state_msg.data = 1
                            publish = True
                        elif '向前走' in matched_command or '前进' in matched_command:
                            state_msg.data = 2
                            publish = True
                        elif '左转' in matched_command:
                            state_msg.data = 3
                            publish = True
                        elif '右转' in matched_command:
                            state_msg.data = 4
                            publish = True
                        elif '向左走' in matched_command:
                            state_msg.data = 5
                            publish = True
                        elif '向右走' in matched_command:
                            state_msg.data = 6
                            publish = True
                        elif '坐下' in matched_command:
                            state_msg.data = 7
                            publish = True
                        
                        if publish:
                            self.state_publisher.publish(state_msg)
                            print(f'检测到命令，发布 state_translate: {state_msg.data}')


def signal_handler(sig, frame):
    print('Ctrl+C pressed, stop recognition ...')
    # Stop recognition
    recognition.stop()
    if HAS_ROS2:
        rclpy.shutdown()
    print('Recognition stopped.')
    print(
        '[Metric] requestId: {}, first package delay ms: {}, last package delay ms: {}'
        .format(
            recognition.get_last_request_id(),
            recognition.get_first_package_delay(),
            recognition.get_last_package_delay(),
        ))
    # Forcefully exit the program
    sys.exit(0)


# main function
def main(args=None):
    # ROS2 初始化
    publisher = None
    state_publisher = None
    activation_subscriber = None
    node = None
    callback = None
    
    if HAS_ROS2:
        rclpy.init(args=args)
        node = rclpy.create_node('voice_recognition_node')
        publisher = node.create_publisher(String, 'voice_words', 10) #asr
        state_publisher = node.create_publisher(Int32, '/state_translate', 10)
        print("ROS2 节点已初始化，正在发布到 'voice_words' 和 '/state_translate' 话题")
        
        # 创建识别回调
        callback = Callback(publisher, state_publisher)
        
        # 创建激活状态订阅者
        def activation_callback(msg):
            if callback:
                callback.update_llm_activation(msg.data)
        
        activation_subscriber = node.create_subscription(Bool, 'llm_activation', activation_callback, 10)
        print("已订阅 'llm_activation' 话题")
    else:
        callback = Callback(publisher, state_publisher)

    # 新加坡和北京地域的API Key不同。获取API Key：https://help.aliyun.com/zh/model-studio/get-api-key
    # 若没有配置环境变量，请用百炼API Key将下行替换为：dashscope.api_key = "sk-xxx"
    dashscope.api_key = os.environ.get('DASHSCOPE_API_KEY')

    # 以下为北京地域url，若使用新加坡地域的模型，需将url替换为：wss://dashscope-intl.aliyuncs.com/api-ws/v1/inference
    dashscope.base_websocket_api_url='wss://dashscope.aliyuncs.com/api-ws/v1/inference'

    # Call recognition service by async mode, you can customize the recognition parameters, like model, format,
    # sample_rate
    global recognition
    recognition = Recognition(
        # model='fun-asr-realtime',
        model='paraformer-realtime-v2',
        format=format_pcm,
        # 'pcm'、'wav'、'opus'、'speex'、'aac'、'amr', you can check the supported formats in the document
        sample_rate=sample_rate,
        # support 8000, 16000
        semantic_punctuation_enabled=False,
        callback=callback)

    # Start recognition
    recognition.start()

    signal.signal(signal.SIGINT, signal_handler)
    print("Press 'Ctrl+C' to stop recording and recognition...")
    
    # 如果使用ROS2，在单独的线程中运行spin
    if HAS_ROS2 and node:
        import threading
        spin_thread = threading.Thread(target=rclpy.spin, args=(node,), daemon=True)
        spin_thread.start()
        print("ROS2 spin线程已启动")
    
    # Create a keyboard listener until "Ctrl+C" is pressed

    try:
        # 使用audio_generator获取音频数据
        for data in audio_generator():
            if len(data) > 0:
                recognition.send_audio_frame(data)
    except KeyboardInterrupt:
        pass
    finally:
        recognition.stop()
        if HAS_ROS2:
            rclpy.shutdown()

if __name__ == '__main__':
    main()