"""
CosyVoice API 服务类
负责与CosyVoice API的交互
"""
import os
import logging
from typing import Optional, Dict, Any, Tuple
from gradio_client import Client, handle_file

logger = logging.getLogger(__name__)

class CosyVoiceService:
    """CosyVoice API服务类"""
    
    def __init__(self, api_url: str = "http://127.0.0.1:8080/"):
        self.api_url = api_url
        self.client = None
        
    def connect(self) -> bool:
        """连接到CosyVoice服务"""
        try:
            self.client = Client(self.api_url)
            logger.info(f"成功连接到CosyVoice服务: {self.api_url}")
            return True
        except Exception as e:
            logger.error(f"连接CosyVoice服务失败: {str(e)}")
            return False
    
    def get_available_voices(self) -> list:
        """获取可用的音色列表"""
        try:
            if not self.client:
                if not self.connect():
                    return []
            
            voices = self.client.predict(api_name="/refresh_sft_spk")
            # 过滤掉不需要的音色
            filtered_voices = [voice for voice in voices if voice != '.ipynb_checkpoints']
            return filtered_voices
        except Exception as e:
            logger.error(f"获取音色列表失败: {str(e)}")
            return []
    
    def get_reference_audios(self) -> list:
        """获取参考音频列表"""
        try:
            if not self.client:
                if not self.connect():
                    return []
            
            audio_files = self.client.predict(api_name="/refresh_prompt_wav")
            return audio_files
        except Exception as e:
            logger.error(f"获取参考音频列表失败: {str(e)}")
            return []
    
    def recognize_audio(self, audio_file_path: str) -> str:
        """语音识别：将音频转换为文本"""
        try:
            if not self.client:
                if not self.connect():
                    return ""
            
            text = self.client.predict(
                prompt_wav=handle_file(audio_file_path),
                api_name="/prompt_wav_recognition"
            )
            return text
        except Exception as e:
            logger.error(f"语音识别失败: {str(e)}")
            return ""
    
    def generate_speech_with_preset_voice(
        self, 
        text: str, 
        voice: str = "中文女", 
        seed: int = 42,
        speed: float = 1.0,
        stream: bool = False
    ) -> Tuple[Optional[str], Optional[str]]:
        """使用预训练音色生成语音"""
        try:
            if not self.client:
                if not self.connect():
                    return None, None
            
            result = self.client.predict(
                tts_text=text,
                mode_checkbox_group="预训练音色",
                sft_dropdown=voice,
                seed=seed,
                speed=speed,
                stream="true" if stream else "false",
                api_name="/generate_audio"
            )
            
            # result是一个元组 [流式音频路径, 完整音频路径]
            if isinstance(result, (list, tuple)) and len(result) >= 2:
                return result[0], result[1]
            else:
                return result, result
                
        except Exception as e:
            logger.error(f"预训练音色语音生成失败: {str(e)}")
            return None, None
    
    def generate_speech_with_voice_cloning(
        self, 
        text: str, 
        reference_audio_path: str,
        reference_text: str = "",
        seed: int = 42
    ) -> Tuple[Optional[str], Optional[str]]:
        """使用语音克隆生成语音"""
        try:
            if not self.client:
                if not self.connect():
                    return None, None
            
            # 如果没有提供参考文本，先进行语音识别
            if not reference_text:
                reference_text = self.recognize_audio(reference_audio_path)
                if not reference_text:
                    logger.warning("参考音频识别失败，使用空文本")
                    reference_text = ""
            
            result = self.client.predict(
                tts_text=text,
                mode_checkbox_group="3s极速复刻",
                prompt_text=reference_text,
                prompt_wav_upload=handle_file(reference_audio_path),
                seed=seed,
                api_name="/generate_audio"
            )
            
            if isinstance(result, (list, tuple)) and len(result) >= 2:
                return result[0], result[1]
            else:
                return result, result
                
        except Exception as e:
            logger.error(f"语音克隆生成失败: {str(e)}")
            return None, None
    
    def generate_speech_with_natural_control(
        self, 
        text: str, 
        instruction: str = "请用温柔甜美的女声朗读",
        seed: int = 42
    ) -> Tuple[Optional[str], Optional[str]]:
        """使用自然语言控制生成语音"""
        try:
            if not self.client:
                if not self.connect():
                    return None, None
            
            result = self.client.predict(
                tts_text=text,
                mode_checkbox_group="自然语言控制",
                instruct_text=instruction,
                seed=seed,
                api_name="/generate_audio"
            )
            
            if isinstance(result, (list, tuple)) and len(result) >= 2:
                return result[0], result[1]
            else:
                return result, result
                
        except Exception as e:
            logger.error(f"自然语言控制语音生成失败: {str(e)}")
            return None, None
    
    def generate_random_seed(self) -> int:
        """生成随机种子"""
        try:
            if not self.client:
                if not self.connect():
                    return 42
            
            seed = self.client.predict(api_name="/generate_random_seed")
            return int(seed) if seed else 42
        except Exception as e:
            logger.error(f"生成随机种子失败: {str(e)}")
            return 42
    
    def test_connection(self) -> Dict[str, Any]:
        """测试与CosyVoice服务的连接"""
        try:
            if not self.connect():
                return {
                    "success": False,
                    "message": "无法连接到CosyVoice服务",
                    "api_url": self.api_url
                }
            
            # 尝试获取音色列表来测试连接
            voices = self.get_available_voices()
            
            return {
                "success": True,
                "message": "CosyVoice服务连接成功",
                "api_url": self.api_url,
                "available_voices": voices
            }
            
        except Exception as e:
            return {
                "success": False,
                "message": f"连接测试失败: {str(e)}",
                "api_url": self.api_url
            }

# 全局服务实例
cosyvoice_service = CosyVoiceService()