#!/usr/bin/env python3
"""
TTS语音合成ROS2节点

功能：
1. 订阅LLM回复文字话题 (llm/response)
2. 使用Xinference API将文字转换为语音
3. 播放语音输出
"""

import os
import time
import rclpy
from rclpy.node import Node
from std_msgs.msg import String
import threading
import sys
import traceback
import numpy as np
import sounddevice as sd
import requests
from xinference.client import Client

class XinferenceTTS:
    """Xinference语音合成"""
    
    def __init__(self, logger):
        self.model_uid = "CosyVoice2-0.5B"
        self.logger = logger
        
        # 音频参数
        self.sampling_rate = 16000
        
        # 初始化Xinference客户端
        try:
            self.client = Client("http://124.93.18.180:9997")
            self.model = self.client.get_model(self.model_uid)
            self.logger.info(f"Xinference客户端初始化成功，模型: {self.model_uid}")
        except Exception as e:
            self.logger.error(f"Xinference客户端初始化失败: {e}")
            self.client = None
            self.model = None
    
    def synthesize(self, text):
        """合成语音（使用Xinference API）"""
        try:
            if not self.model:
                self.logger.error("Xinference模型未初始化")
                return None
            
            self.logger.info(f"正在调用Xinference TTS API，文本: {text[:30]}...")
            
            # 调用Xinference TTS API - 使用官方speech方法
            audio_data = self.model.speech(
                input=text,
                response_format='mp3',
                speed=1.0,
                stream=False
            )
            
            if audio_data is not None:
                # 将音频数据转换为numpy数组
                if isinstance(audio_data, bytes):
                    audio_int16 = np.frombuffer(audio_data, dtype=np.int16)
                    audio_float32 = audio_int16.astype(np.float32) / 32767.0
                    return audio_float32
                elif isinstance(audio_data, np.ndarray):
                    return audio_data
                else:
                    self.logger.error(f"不支持的音频数据类型: {type(audio_data)}")
                    return None
            else:
                self.logger.error("Xinference TTS API返回空数据")
                return None
            
        except Exception as e:
            self.logger.error(f"合成错误: {e}")
            traceback.print_exc()
            return None
    
    def play_audio(self, audio_data):
        """播放音频"""
        try:
            # 确保是numpy数组
            if not isinstance(audio_data, np.ndarray):
                audio_data = np.array(audio_data)
            
            # 如果是双声道，转为单声道
            if len(audio_data.shape) > 1:
                audio_data = audio_data.mean(axis=1)
            
            # 使用OutputStream控制缓冲区大小和延迟，减少ALSA欠载错误
            with sd.OutputStream(
                samplerate=self.sampling_rate,
                channels=1,
                dtype=np.float32,
                latency=0.1,  # 增加延迟以减少ALSA欠载错误
                blocksize=int(self.sampling_rate * 0.1)  # 增加缓冲区大小
            ) as stream:
                # 写入音频数据
                stream.write(audio_data)
                # 等待缓冲区中的数据完全播放完成
                stream.stop()  # 确保所有数据都已播放
                time.sleep(0.1)  # 额外等待0.1秒，确保最后一个字播放完成
            
            # 播放完成后延迟足够时间，避免捕获回声和残留声音
            time.sleep(0.4)  # 增加到0.5秒，确保回声完全消失
            
            self.logger.info('TTS语音播放完成')
        except Exception as e:
            self.logger.error(f'播放音频时出错: {e}')
            traceback.print_exc()

class TTSNode(Node):
    """TTS语音合成ROS2节点"""
    
    def __init__(self):
        super().__init__('tts_node')
        
        # 声明参数
        self.model_uid = os.getenv('XINFERENCE_MODEL_UID', 'tts-model')
        
        # 获取参数
        
        # 创建订阅者
        self.response_subscriber = self.create_subscription(
            String,
            'llm/response',
            self.response_callback,
            10
        )
        
        # 创建播放状态发布者
        self.playback_status_publisher = self.create_publisher(
            String,
            'tts/playback_status',
            10
        )
        
        # 初始化XinferenceTTS
        self.tts_engine = XinferenceTTS(
            self.model_uid,
            self.get_logger()
        )
        
        # 语音播放线程锁
        self.playback_lock = threading.Lock()
        
        self.get_logger().info('TTS语音合成节点初始化完成')
    
    def response_callback(self, msg):
        """处理LLM回复回调"""
        text = msg.data
        self.get_logger().info(f'收到LLM回复: {text}')
        
        # 在单独的线程中播放语音，避免阻塞回调
        threading.Thread(target=self.speak, args=(text,)).start()
    
    def speak(self, text):
        """播放语音"""
        with self.playback_lock:
            try:
                # 发布开始播放状态
                status_msg = String()
                status_msg.data = "playing"
                self.playback_status_publisher.publish(status_msg)
                self.get_logger().info('发布播放状态: playing')
                
                # 合成语音
                self.get_logger().info(f'开始合成语音: {text[:30]}...')
                audio_data = self.tts_engine.synthesize(text)
                
                if audio_data is not None:
                    # 播放语音
                    self.tts_engine.play_audio(audio_data)
                else:
                    self.get_logger().error('语音合成失败，无法播放')
            except Exception as e:
                self.get_logger().error(f'播放语音时出错: {e}')
                traceback.print_exc()
            finally:
                # 发布播放完成状态
                status_msg = String()
                status_msg.data = "idle"
                self.playback_status_publisher.publish(status_msg)
                self.get_logger().info('发布播放状态: idle')
    
    def destroy_node(self):
        """销毁节点"""
        super().destroy_node()

def main(args=None):
    """主函数"""
    rclpy.init(args=args)
    
    # 创建节点
    node = TTSNode()
    
    try:
        # 运行节点
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        # 销毁节点
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
