# -*- coding: utf-8 -*-
"""
@author:XuMing(xuming624@qq.com)
@description: Re-train by TWMAN
"""
import hashlib
import os
import ssl

import gradio as gr
import torch
from loguru import logger

ssl._create_default_https_context = ssl._create_unverified_context
import nltk

# 檢查是否已下載資源，若未下載則進行下載
nltk_data_path = os.path.expanduser('~/nltk_data')
if not os.path.exists(os.path.join(nltk_data_path, 'corpora/cmudict.zip')):
    nltk.download('cmudict', download_dir=nltk_data_path)

if not os.path.exists(os.path.join(nltk_data_path, 'taggers/averaged_perceptron_tagger.zip')):
    nltk.download('averaged_perceptron_tagger', download_dir=nltk_data_path)

from parrots import TextToSpeech

# 設定裝置與模式
device = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"device: {device}")
half = True if device == "cuda" else False

# 初始化語音合成模型
m = TextToSpeech(speaker_model_path="DeepLearning101/GPT-SoVITS_TWMAN", speaker_name="TWMAN", device=device, half=half)

# 用於檢查和生成語音的音訊檔案
def get_text_hash(text: str):
    return hashlib.md5(text.encode('utf-8')).hexdigest()

def do_tts_wav_predict(text: str, output_path: str = None):
    if output_path is None:
        output_path = f"output_audio_{get_text_hash(text)}.wav"
    if not os.path.exists(output_path):
        m.predict(text, text_language="auto", output_path=output_path)
    return output_path

# 建立 Gradio WebUI
with gr.Blocks(title="TTS WebUI") as app:
    gr.Markdown("""
    # 線上語音合成 (TWMAN)
    #### 請嚴格遵守法規，發布二創作品請標註本專案作者及連結，並標註生成工具 GPT-SoVITS AI！
    ⚠️ 注意：在線生成可能較慢，建議在本地進行推理。 
    
    更多相關內容：
    - [語音處理技術](https://www.twman.org/AI/ASR)
    - [語音處理常見問題](https://blog.twman.org/2021/04/ASR.html)
    - [Parrots專案](https://github.com/shibing624/parrots)
    - [模型使用說明](https://github.com/RVC-Boss/GPT-SoVITS)
    """)

    # 設定語音合成輸入與按鈕
    with gr.Group():
        gr.Markdown("*請在下方輸入要進行語音合成的文字*")
        with gr.Row():
            text = gr.Textbox(label="想語音合成的文字 (100字以内)", value="床前明月光，疑是地上霜。舉頭望明月，低頭思故鄉。", placeholder="請輸入您想要的文字", lines=3)
            inference_button = gr.Button("語音合成", variant="primary")
            output = gr.Audio(label="合成的語音")

        # 設定按鈕點擊事件
        inference_button.click(
            do_tts_wav_predict,
            [text],
            [output],
        )

# 啟動 Gradio 應用
app.queue(max_size=10)
app.launch(share=True, inbrowser=True)