#!/usr/bin/env python3
import os
import signal
import sys
from datetime import datetime

import dashscope
import sounddevice as sd
from dashscope.audio.asr import *

# sd.default.device = None
sd.default.device = (0, 0)
sd.default.samplerate = 16000
sd.default.channels = 1

try:
    import rclpy
    from rclpy.node import Node
    from std_msgs.msg import String
    HAS_ROS2 = True
except ImportError:
    HAS_ROS2 = False

stream = None
recognition = None


class ASRCallback(RecognitionCallback):
    def __init__(self, publisher=None):
        super().__init__()
        self.publisher = publisher

    def on_open(self) -> None:
        global stream
        print('ASR 回调已打开。')
        try:
            stream = sd.InputStream(
                samplerate=16000,
                channels=1,
                dtype='int16',
                blocksize=3200
            )
            stream.start()
            print(f'使用 sounddevice 打开音频流成功，\n当前设备输入设备: {sd.default.device[0]}, 输出设备: {sd.default.device[1]}')

        except Exception as e:
            print(f'初始化音频设备时出错: {e}')
            raise

    def on_close(self) -> None:
        global stream
        print('ASR 回调已关闭。')
        if stream:
            stream.stop()
            stream.close()
            stream = None

    def on_complete(self) -> None:
        print('ASR 回调已完成。')

    def on_error(self, message) -> None:
        print('ASR 回调错误ID: ', message.request_id)
        print('ASR 回调错误信息: ', message.message)
        global stream
        if stream:
            stream.stop()
            stream.close()
        sys.exit(1)

    def on_event(self, result: RecognitionResult) -> None:
        sentence = result.get_sentence()
        if 'text' in sentence:
            print('ASR 识别文本: ', sentence['text'])

            if RecognitionResult.is_sentence_end(sentence):
                print('ASR 识别文本结束。')

                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}')


def signal_handler(sig, frame):
    recognition.stop()
    if HAS_ROS2:
        rclpy.shutdown()
    print('ASR 回调已停止。')
    if recognition:
        print(
            '[Metric] 请求ID: {}, 第一个包延迟ms: {}, 最后包延迟ms: {}'
            .format(
                recognition.get_last_request_id(),
                recognition.get_first_package_delay(),
                recognition.get_last_package_delay(),
            ))
    sys.exit(0)


def main(args=None):
    publisher = None
    node = None
    callback = None

    if HAS_ROS2:
        rclpy.init(args=args)
        node = rclpy.create_node('asr_node')
        publisher = node.create_publisher(String, 'voice_words', 10)
        print("ROS2 ASR 节点已初始化，正在发布到 'voice_words' 话题")

        callback = ASRCallback(publisher)
    else:
        callback = ASRCallback(publisher)

    dashscope.api_key = os.environ.get('DASHSCOPE_API_KEY')
    dashscope.base_websocket_api_url = 'wss://dashscope.aliyuncs.com/api-ws/v1/inference'

    global recognition
    recognition = Recognition(
        model='fun-asr-realtime',
        format='pcm',
        sample_rate=16000,
        semantic_punctuation_enabled=False,
        callback=callback)

    recognition.start()

    signal.signal(signal.SIGINT, signal_handler)
    print("ASR 节点运行中，Ctrl+C 停止 ...")

    if HAS_ROS2 and node:
        import threading
        spin_thread = threading.Thread(target=rclpy.spin, args=(node,), daemon=True)
        spin_thread.start()
        print("ROS2 spin线程已启动")

    try:
        while True:
            if stream:
                data, overflowed = stream.read(3200)
                data_bytes = data.tobytes()
                recognition.send_audio_frame(data_bytes)
            else:
                import time
                time.sleep(0.1)
                if not stream and not callback:
                     break
    except KeyboardInterrupt:
        pass
    finally:
        recognition.stop()
        if HAS_ROS2:
            rclpy.shutdown()


if __name__ == '__main__':
    main()
