import io
import time
from ntpath import join

import requests
from fastapi import Request, Response, status, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
from .blackbox import Blackbox
from ..tts.tts_service import TTService
from ..configuration import MeloConf
from ..configuration import CosyVoiceConf
from ..configuration import SovitsConf
from injector import inject
from injector import singleton

import sys,os
sys.path.append('/Workspace/CosyVoice')
sys.path.append('/Workspace/CosyVoice/third_party/Matcha-TTS')
from cosyvoice.cli.cosyvoice import CosyVoice, CosyVoice2
# from cosyvoice.utils.file_utils import load_wav, speed_change

import soundfile as sf
import pyloudnorm as pyln
from melo.api import TTS as MELOTTS

from ..log.logging_time import logging_time
import logging
logger = logging.getLogger(__name__)

import random
import torch
import numpy as np

from pydub import AudioSegment
import subprocess

def set_all_random_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def convert_wav_to_mp3(wav_filename: str):
    # 检查文件是否是 .wav 格式
    if not wav_filename.lower().endswith(".wav"):
        raise ValueError("The input file must be a .wav file")

    # 构建对应的 .mp3 文件路径
    mp3_filename = wav_filename.replace(".wav", ".mp3")

    # 如果原 MP3 文件已存在，先删除
    if os.path.exists(mp3_filename):
        os.remove(mp3_filename)

    # 使用 FFmpeg 进行转换，直接覆盖原有的 MP3 文件
    command = ['ffmpeg', '-i', wav_filename, '-y', mp3_filename]  # `-y` 参数会自动覆盖目标文件
    result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    # 检查转换是否成功
    if result.returncode != 0:
        raise Exception(f"Error converting {wav_filename} to MP3: {result.stderr.decode()}")

    return mp3_filename

    
@singleton
class TTS(Blackbox):
    melo_mode: str
    melo_url: str
    melo_speed: int
    melo_device: str
    melo_language: str
    melo_speaker: str

    cosyvoice_mode: str
    cosyvoice_url: str
    cosyvoice_speed: int
    cosyvoice_device: str
    cosyvoice_language: str
    cosyvoice_speaker: str

    @logging_time(logger=logger)
    def melo_model_init(self, melo_config: MeloConf) -> None:
        self.melo_speed = melo_config.speed
        self.melo_device = melo_config.device
        self.melo_language = melo_config.language
        self.melo_speaker = melo_config.speaker
        self.melo_url = ''
        self.melo_mode = melo_config.mode
        self.melotts = None
        # self.speaker_ids = None
        self.melo_speaker = None        
        if self.melo_mode == 'local':
            self.melotts = MELOTTS(language=self.melo_language, device=self.melo_device)
            self.melo_speaker = self.melotts.hps.data.spk2id[self.melo_language]
        else:
            self.melo_url = melo_config.url
        logging.info('#### Initializing MeloTTS Service in ' + self.melo_device + ' mode...')
        print('1.#### Initializing MeloTTS Service in ' + self.melo_device + ' mode...')
    
    @logging_time(logger=logger)
    def cosyvoice_model_init(self, cosyvoice_config: CosyVoiceConf) -> None:
        self.cosyvoice_speed = cosyvoice_config.speed
        self.cosyvoice_device = cosyvoice_config.device
        self.cosyvoice_language = cosyvoice_config.language
        self.cosyvoice_speaker = cosyvoice_config.speaker
        self.cosyvoice_url = ''
        self.cosyvoice_mode = cosyvoice_config.mode
        self.cosyvoicetts = None
        # os.environ['CUDA_VISIBLE_DEVICES'] = str(cosyvoice_config.device)
        if self.cosyvoice_mode == 'local':
            # self.cosyvoicetts = CosyVoice('/Workspace/Models/CosyVoice/pretrained_models/CosyVoice-300M')
            self.cosyvoicetts = CosyVoice('/model/Voice/CosyVoice/pretrained_models/CosyVoice-300M')
            # self.cosyvoicetts = CosyVoice2('/model/Voice/CosyVoice/pretrained_models/CosyVoice2-0.5B', load_jit=True, load_onnx=False, load_trt=False)

            
        else:
            self.cosyvoice_url = cosyvoice_config.url
        logging.info('#### Initializing CosyVoiceTTS Service in ' + self.cosyvoice_device + ' mode...')
        print('1.#### Initializing CosyVoiceTTS Service in ' + self.cosyvoice_device + ' mode...')

    @logging_time(logger=logger)
    def sovits_model_init(self, sovits_config: SovitsConf) -> None:
        self.sovits_speed = sovits_config.speed
        self.sovits_device = sovits_config.device
        self.sovits_language = sovits_config.language
        self.sovits_speaker = sovits_config.speaker
        self.sovits_url = sovits_config.url
        self.sovits_mode = sovits_config.mode
        self.sovitstts = None
        self.speaker_ids = None

        self.sovits_text_lang = sovits_config.text_lang
        self.sovits_ref_audio_path = sovits_config.ref_audio_path
        self.sovits_prompt_lang = sovits_config.prompt_lang
        self.sovits_prompt_text = sovits_config.prompt_text
        self.sovits_text_split_method = sovits_config.text_split_method
        self.sovits_batch_size = sovits_config.batch_size
        self.sovits_media_type = sovits_config.media_type
        self.sovits_streaming_mode = sovits_config.streaming_mode
          
        if self.sovits_mode == 'local':
            # self.sovitsts = MELOTTS(language=self.melo_language, device=self.melo_device)
            self.sovits_speaker = self.sovitstts.hps.data.spk2id
        else:
            self.sovits_url = sovits_config.url
        logging.info('#### Initializing SoVITS Service in ' + self.sovits_device + ' mode...')
        print('1.#### Initializing SoVITS Service in ' + self.sovits_device + ' mode...')


    @inject   
    def __init__(self, melo_config: MeloConf, cosyvoice_config: CosyVoiceConf, sovits_config: SovitsConf, settings: dict) -> None:
        self.tts_service = TTService("yunfeineo")
        self.melo_model_init(melo_config)
        self.cosyvoice_model_init(cosyvoice_config)
        self.sovits_model_init(sovits_config)
        self.audio_dir = "audio_files"  # 存储音频文件的目录

    def __call__(self, *args, **kwargs):
        return self.processing(*args, **kwargs)

    @logging_time(logger=logger)
    def processing(self, *args, settings: dict) -> io.BytesIO:

        print("\nChat Settings: ", settings)
        if settings is None:
            settings = {}
        user_model_name = settings.get("tts_model_name")
        chroma_collection_id = settings.get("chroma_collection_id")
        user_stream = settings.get('tts_stream')
        print(f"chroma_collection_id: {chroma_collection_id}")
        print(f"tts_model_name: {user_model_name}")

        if user_stream in [None, ""]:
            user_stream = True

        text = args[0]
        current_time = time.time()
        if user_model_name == 'melotts': 
            if chroma_collection_id == 'kiki' or chroma_collection_id is None:
                if self.melo_mode == 'local':
                    audio = self.melotts.tts_to_file(text, self.melo_speaker, speed=self.melo_speed)
                    f = io.BytesIO()
                    sf.write(f, audio, 44100, format='wav')     
                    f.seek(0)

                    # Read the audio data from the buffer
                    data, rate = sf.read(f, dtype='float32')
                    
                    # Peak normalization
                    peak_normalized_audio = pyln.normalize.peak(data, -1.0)
                    
                    # Integrated loudness normalization
                    meter = pyln.Meter(rate)
                    loudness = meter.integrated_loudness(peak_normalized_audio)
                    loudness_normalized_audio = pyln.normalize.loudness(peak_normalized_audio, loudness, -12.0)
                    
                    # Write the loudness normalized audio to an in-memory buffer
                    normalized_audio_buffer = io.BytesIO()
                    sf.write(normalized_audio_buffer, loudness_normalized_audio, rate, format='wav')
                    normalized_audio_buffer.seek(0)
                    
                    print("#### MeloTTS Service consume - local : ", (time.time() - current_time))
                    return normalized_audio_buffer.read()   

                else:
                    message = {
                        "text": text
                    }
                    response = requests.post(self.melo_url, json=message)
                    print("#### MeloTTS Service consume - docker : ", (time.time()-current_time))
                    return response.content 
            elif chroma_collection_id == 'boss':
                if self.cosyvoice_mode == 'local':
                    set_all_random_seed(35616313)
                    audio = self.cosyvoicetts.inference_sft(text, '中文男', speed=1.5, stream=False)
                    for i, j in enumerate(audio):
                        f = io.BytesIO()
                        sf.write(f, j['tts_speech'].cpu().numpy().squeeze(0), 22050, format='wav')
                        f.seek(0)
                    print("#### CosyVoiceTTS Service consume - local : ", (time.time() - current_time))
                    return f.read()
                else:
                    message = {
                        "text": text
                    }
                    response = requests.post(self.cosyvoice_url, json=message)
                    print("#### CosyVoiceTTS Service consume - docker : ", (time.time()-current_time))
                    return response.content             
            
        elif user_model_name == 'cosyvoicetts':
            if chroma_collection_id == 'kiki' or chroma_collection_id is None:
                if self.cosyvoice_mode == 'local':
                    set_all_random_seed(56056558)
                    print("*"*90)
                    audio = self.cosyvoicetts.inference_sft(text, self.cosyvoice_language, stream=True)
                    for i, j in enumerate(audio):
                        f = io.BytesIO()
                        sf.write(f, j['tts_speech'].cpu().numpy().squeeze(0), 22050, format='wav')   
                        f.seek(0)
                    print("#### CosyVoiceTTS Service consume - local : ", (time.time() - current_time))
                    return f.read()
                else:
                    message = {
                        "text": text
                    }
                    response = requests.post(self.cosyvoice_url, json=message)
                    print("#### CosyVoiceTTS Service consume - docker : ", (time.time()-current_time))
                    return response.content            
            elif chroma_collection_id == 'boss':
                if self.cosyvoice_mode == 'local':
                    set_all_random_seed(35616313)
                    audio = self.cosyvoicetts.inference_sft(text, '中文男', speed=1.5, stream=False)
                    for i, j in enumerate(audio):
                        f = io.BytesIO()
                        sf.write(f, j['tts_speech'].cpu().numpy().squeeze(0), 22050, format='wav')
                        f.seek(0)
                    print("#### CosyVoiceTTS Service consume - local : ", (time.time() - current_time))
                    return f.read()
                else:
                    message = {
                        "text": text
                    }
                    response = requests.post(self.cosyvoice_url, json=message)
                    print("#### CosyVoiceTTS Service consume - docker : ", (time.time()-current_time))
                    return response.content 

        elif user_model_name == 'sovitstts':
            if chroma_collection_id == 'kiki' or chroma_collection_id is None:
                if self.sovits_mode == 'local':
                    set_all_random_seed(56056558)                   
                    audio = self.cosyvoicetts.inference_sft(text, self.cosyvoice_language, stream=True)
                    for i, j in enumerate(audio):
                        f = io.BytesIO()
                        sf.write(f, j['tts_speech'].cpu().numpy().squeeze(0), 22050, format='wav')   
                        f.seek(0)
                    print("#### CosyVoiceTTS Service consume - local : ", (time.time() - current_time))
                    return f.read()
                else:
                    message = {
                        "text": text,
                        "text_lang": self.sovits_text_lang,
                        "ref_audio_path": self.sovits_ref_audio_path,
                        "prompt_lang": self.sovits_prompt_lang,
                        "prompt_text": self.sovits_prompt_text,
                        "text_split_method": self.sovits_text_split_method,
                        "batch_size": self.sovits_batch_size,
                        "media_type": self.sovits_media_type,
                        "streaming_mode": self.sovits_streaming_mode
                    }
                    if user_stream:
                        response = requests.get(self.sovits_url, params=message, stream=True)
                        print("#### SoVITS Service consume - docker : ", (time.time()-current_time))
                        return response
                    else:
                        response = requests.get(self.sovits_url, params=message)
                        print("#### SoVITS Service consume - docker : ", (time.time()-current_time))
                        return response.content
                    print("#### SoVITS Service consume - docker : ", (time.time()-current_time))
                    
                        
            # elif chroma_collection_id == 'boss':
            #     if self.cosyvoice_mode == 'local':
            #         set_all_random_seed(35616313)
            #         audio = self.cosyvoicetts.inference_sft(text, '中文男', speed=1.5, stream=False)
            #         for i, j in enumerate(audio):
            #             f = io.BytesIO()
            #             sf.write(f, j['tts_speech'].cpu().numpy().squeeze(0), 22050, format='wav')
            #             f.seek(0)
            #         print("#### CosyVoiceTTS Service consume - local : ", (time.time() - current_time))
            #         return f.read()
            #     else:
            #         message = {
            #             "text": text
            #         }
            #         response = requests.post(self.cosyvoice_url, json=message)
            #         print("#### CosyVoiceTTS Service consume - docker : ", (time.time()-current_time))
            #         return response.content   


        elif user_model_name == 'man':
            if self.cosyvoice_mode == 'local':
                set_all_random_seed(35616313)
                audio = self.cosyvoicetts.inference_sft(text, '中文男', speed=1.5, stream=False)
                for i, j in enumerate(audio):
                    f = io.BytesIO()
                    sf.write(f, j['tts_speech'].cpu().numpy().squeeze(0), 22050, format='wav')
                    f.seek(0)
                print("#### CosyVoiceTTS Service consume - local : ", (time.time() - current_time))
                return f.read()
            else:
                message = {
                    "text": text
                }
                response = requests.post(self.cosyvoice_url, json=message)
                print("#### CosyVoiceTTS Service consume - docker : ", (time.time()-current_time))
                return response.content                   

        else:           
            audio = self.tts_service.read(text)
            print("#### TTS Service consume : ", (time.time()-current_time))
            return audio.read()

    def valid(self, *args, **kwargs) -> bool:
        text = args[0]
        return isinstance(text, str)

    async def fast_api_handler(self, request: Request) -> Response:
        try:
            data = await request.json()
            print(f"data: {data}")
        except:
            return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)
        text = data.get("text")
        setting = data.get("settings")
        tts_model_name = setting.get("tts_model_name")
        user_stream = setting.get("tts_stream")

        if text is None:
            return JSONResponse(content={"error": "text is required"}, status_code=status.HTTP_400_BAD_REQUEST)
        by = self.processing(text, settings=setting)
        # return Response(content=by, media_type="audio/wav", headers={"Content-Disposition": "attachment; filename=audio.wav"})

        if user_stream:
            if by.status_code == 200:
                print("*"*90)
                def audio_stream():
                    # 从上游服务器流式读取数据并逐块发送
                    for chunk in by.iter_content(chunk_size=1024):
                        if chunk:
                            yield chunk

                return StreamingResponse(audio_stream(), media_type="audio/wav")
            else:
                raise HTTPException(status_code=by.status_code, detail="请求失败")

        # if user_stream and tts_model_name == 'sovitstts':
        #     if by.status_code == 200:
        #         # 保存 WAV 文件
        #         wav_filename = os.path.join(self.audio_dir, 'audio.wav')
        #         with open(wav_filename, 'wb') as f:
        #             for chunk in by.iter_content(chunk_size=1024):
        #                 if chunk:
        #                     f.write(chunk)

        #         try:
        #             # 检查文件是否存在
        #             if not os.path.exists(wav_filename):
        #                 return JSONResponse(content={"error": "File not found"}, status_code=404)

        #             # 转换 WAV 为 MP3
        #             mp3_filename = convert_wav_to_mp3(wav_filename)

        #             # 返回 WAV 和 MP3 文件的下载链接
        #             wav_url = f"/audio/{wav_filename}"
        #             mp3_url = f"/audio/{mp3_filename}"
        #             return JSONResponse(content={"wav_url": wav_url, "mp3_url": mp3_url}, status_code=200)

        #         except Exception as e:
        #             return JSONResponse(content={"error": str(e)}, status_code=500)

        #     else:
        #         raise HTTPException(status_code=by.status_code, detail="请求失败")


        else:
            wav_filename = os.path.join(self.audio_dir, 'audio.wav')
            with open(wav_filename, 'wb') as f:
                f.write(by)

            try:
                # 先检查文件是否存在
                if not os.path.exists(wav_filename):
                    return JSONResponse(content={"error": "File not found"}, status_code=404)

                # 转换 WAV 为 MP3
                mp3_filename = convert_wav_to_mp3(wav_filename)

                # 返回 MP3 文件的下载链接
                wav_url = f"/audio/{wav_filename}"
                mp3_url = f"/audio/{mp3_filename}"
                return JSONResponse(content={"wav_url": wav_url, "mp3_url": mp3_url}, status_code=200)
            
            except Exception as e:
                return JSONResponse(content={"error": str(e)}, status_code=500)