import os
import numpy as np
import torch
from torch import LongTensor
import soundfile as sf
# vits
from .vits import utils, commons
from .vits.models import SynthesizerTrn
from .vits.text import text_to_sequence

def tts_model_init(model_path='./vits_model', device='cuda'):
    hps_ms = utils.get_hparams_from_file(os.path.join(model_path, 'config.json'))
    # hps_ms = utils.get_hparams_from_file('vits_model/config.json')
    net_g_ms = SynthesizerTrn(
        len(hps_ms.symbols),
        hps_ms.data.filter_length // 2 + 1,
        hps_ms.train.segment_size // hps_ms.data.hop_length,
        n_speakers=hps_ms.data.n_speakers,
        **hps_ms.model)
    net_g_ms = net_g_ms.eval().to(device)
    speakers = hps_ms.speakers
    utils.load_checkpoint(os.path.join(model_path, 'G_953000.pth'), net_g_ms, None)
    # utils.load_checkpoint('vits_model/G_953000.pth', net_g_ms, None)
    return hps_ms, net_g_ms, speakers

class TextToSpeech:
    def __init__(self, 
                 model_path="./utils/tts/vits_model", 
                 device='cuda',
                 RATE=22050,
                 debug=False,
                 ):
        self.debug = debug
        self.RATE = RATE
        self.device = torch.device(device)
        self.limitation = os.getenv("SYSTEM") == "spaces"  # 在huggingface spaces中限制文本和音频长度
        self.hps_ms, self.net_g_ms, self.speakers = self._tts_model_init(model_path)
        
    def _tts_model_init(self, model_path):
        hps_ms = utils.get_hparams_from_file(os.path.join(model_path, 'config.json'))
        net_g_ms = SynthesizerTrn(
            len(hps_ms.symbols),
            hps_ms.data.filter_length // 2 + 1,
            hps_ms.train.segment_size // hps_ms.data.hop_length,
            n_speakers=hps_ms.data.n_speakers,
            **hps_ms.model)
        net_g_ms = net_g_ms.eval().to(self.device)
        speakers = hps_ms.speakers
        utils.load_checkpoint(os.path.join(model_path, 'G_953000.pth'), net_g_ms, None)
        if self.debug:
            print("Model loaded.")
        return hps_ms, net_g_ms, speakers

    def _get_text(self, text):
        text_norm, clean_text = text_to_sequence(text, self.hps_ms.symbols, self.hps_ms.data.text_cleaners)
        if self.hps_ms.data.add_blank:
            text_norm = commons.intersperse(text_norm, 0)
        text_norm = LongTensor(text_norm)
        return text_norm, clean_text

    def _preprocess_text(self, text, language):
        if language == 0:
            return f"[ZH]{text}[ZH]"
        elif language == 1:
            return f"[JA]{text}[JA]"
        return text

    def _generate_audio(self, text, speaker_id, noise_scale, noise_scale_w, length_scale):
        import time
        start_time = time.time()
        stn_tst, clean_text = self._get_text(text)
        with torch.no_grad():
            x_tst = stn_tst.unsqueeze(0).to(self.device)
            x_tst_lengths = LongTensor([stn_tst.size(0)]).to(self.device)
            speaker_id = LongTensor([speaker_id]).to(self.device)
            audio = self.net_g_ms.infer(x_tst, x_tst_lengths, sid=speaker_id, noise_scale=noise_scale, noise_scale_w=noise_scale_w,
                                        length_scale=length_scale)[0][0, 0].data.cpu().float().numpy()
        if self.debug:
            print(f"Synthesis time: {time.time() - start_time} s")
        return audio
    
    def synthesize(self, text, language, speaker_id, noise_scale, noise_scale_w, length_scale, save_audio=False, return_bytes=False):
        if not len(text):
            return "输入文本不能为空！", None
        text = text.replace('\n', ' ').replace('\r', '').replace(" ", "")
        if len(text) > 100 and self.limitation:
            return f"输入文字过长！{len(text)}>100", None
        text = self._preprocess_text(text, language)
        audio = self._generate_audio(text, speaker_id, noise_scale, noise_scale_w, length_scale)
        if self.debug or save_audio:
            self.save_audio(audio, self.RATE, 'output_file.wav')
        if return_bytes:
            audio = self.convert_numpy_to_bytes(audio)
        return self.RATE, audio

    def convert_numpy_to_bytes(self, audio_data):
        if isinstance(audio_data, np.ndarray):
            if audio_data.dtype == np.dtype('float32'):
                audio_data = np.int16(audio_data * np.iinfo(np.int16).max)
            audio_data = audio_data.tobytes()
            return audio_data
        else:
            raise TypeError("audio_data must be a numpy array")

    def save_audio(self, audio, sample_rate, file_name='output_file.wav'):
        sf.write(file_name, audio, samplerate=sample_rate)
        print(f"VITS Audio saved to {file_name}")