import spaces
import gradio as gr
import torch
from huggingface_hub import hf_hub_download
import os
import sys
import tempfile
from scipy.io.wavfile import write
import numpy as np
from tqdm import tqdm
from underthesea import sent_tokenize

try:
    from TTS.tts.configs.xtts_config import XttsConfig
    from TTS.tts.models.xtts import Xtts
except ImportError:
    os.system("git clone https://github.com/hellcatmon/XTTSv2-Finetuning-for-New-Languages.git")
    if os.path.exists("XTTSv2-Finetuning-for-New-Languages/TTS"):
        os.system("mv XTTSv2-Finetuning-for-New-Languages/TTS ./")
    sys.path.append("./TTS")
    from TTS.tts.configs.xtts_config import XttsConfig
    from TTS.tts.models.xtts import Xtts

# Шляхі да файлаў (цяпер як радкі)
repo_id = "archivartaunik/BE_XTTS_V2_60epoch3Dataset"
model_dir = "./model"  # Дырэкторыя для захавання мадэлі
os.makedirs(model_dir, exist_ok=True) # Ствараем дырэкторыю, калі яе няма
checkpoint_file = os.path.join(model_dir, "model.pth")
config_file = os.path.join(model_dir, "config.json")
vocab_file = os.path.join(model_dir, "vocab.json")
default_voice_file = os.path.join(model_dir, "voice.wav")

if not os.path.exists(checkpoint_file):
    hf_hub_download(repo_id, filename="model.pth", local_dir=model_dir)
if not os.path.exists(config_file):
    hf_hub_download(repo_id, filename="config.json", local_dir=model_dir)
if not os.path.exists(vocab_file):
    hf_hub_download(repo_id, filename="vocab.json", local_dir=model_dir)
if not os.path.exists(default_voice_file):
    hf_hub_download(repo_id, filename="voice.wav", local_dir=model_dir)

# Загрузка канфігурацыі і мадэлі адзін раз
config = XttsConfig()
config.load_json(config_file)
XTTS_MODEL = Xtts.init_from_config(config)
XTTS_MODEL.load_checkpoint(config, checkpoint_path=checkpoint_file, vocab_path=vocab_file, use_deepspeed=False) # Тут выпраўленне
device = "cuda:0" if torch.cuda.is_available() else "cpu"
XTTS_MODEL.to(device)
sampling_rate = XTTS_MODEL.config.audio["sample_rate"]

@spaces.GPU(duration=60)
def text_to_speech(belarusian_story, speaker_audio_file=None):
    if not speaker_audio_file or (not isinstance(speaker_audio_file, str) and speaker_audio_file.name == ""):
        speaker_audio_file = default_voice_file

    try:
        gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(
            audio_path=speaker_audio_file,
            gpt_cond_len=XTTS_MODEL.config.gpt_cond_len,
            max_ref_length=XTTS_MODEL.config.max_ref_len,
            sound_norm_refs=XTTS_MODEL.config.sound_norm_refs,
        )
    except Exception as e:
        return f"Error getting conditioning latents: {e}"

    try:
        tts_texts = sent_tokenize(belarusian_story)
    except Exception as e:
        return f"Error tokenizing text: {e}"

    all_wavs = []
    for text in tqdm(tts_texts):
        try:
            with torch.no_grad():
                wav_chunk = XTTS_MODEL.inference(
                    text=text,
                    language="be",
                    gpt_cond_latent=gpt_cond_latent,
                    speaker_embedding=speaker_embedding,
                    temperature=0.1,
                    length_penalty=1.0,
                    repetition_penalty=10.0,
                    top_k=10,
                    top_p=0.3,
                )
            all_wavs.append(wav_chunk["wav"])
        except Exception as e:
            return f"Error generating audio: {e}"

    try:
        out_wav = np.concatenate(all_wavs)
    except ValueError:
        return "Немагчыма згенерыраваць аўдыё. Праверце ўваходны тэкст і аўдыёфайл."
    except Exception as e:
        return f"Error concatenating audio: {e}"

    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    write(temp_file.name, sampling_rate, out_wav)

    return temp_file.name

examples = [
    ["Такім чынам, клуб стаў уладальнікам усіх існых на сёння міжнародных трафеяў паўднёваамерыканскага футболу.", "Nestarka.wav", "krai.wav"],
    ["Яму не ўдалося палепшыць фінансавае становішча каралеўства, а, наадварот, прыйшлося распрадаваць каштоўнасці чэшскай кароны.", "muzh.wav", "examples/цуды.wav"],
    ["Кампілятарамі называюць праграмы, якія пераўтвараюць код вышэйшага ўзроўню ў код ніжэйшага ўзроўню.", "chunk_100.wav", "examples/надВозерам.wav"],
    ["Акрамя таго, ліхачы аддаюць перавагу рэгі, хіп-хопу і класічнай музыцы.", "d1015.mp3", "examples/Беларусь.wav"],
    ["Позірк можа быць уважлівым, зацікаўленым, захопленым, але бывае і нахабным, задзірлівым, пагардлівым, напышлівым.", "donarka_ench.wav", "examples/цуды.wav"],
    ["Такі нават шчыры, ці што: родная мова народу – трасянка, а беларуская яму чужая!", "muzhcynski.wav", "examples/цуды.wav"],
    ]
# Генерацыя адсутных файлаў
for example in examples:
    input_text, input_audio, output_audio = example
    if not os.path.exists(output_audio):
        print(f"Creating missing file: {output_audio}")
        generated_audio = text_to_speech(input_text, input_audio)
        os.rename(generated_audio, output_audio)
        
demo = gr.Interface(
    fn=text_to_speech,
    inputs=[
        gr.Textbox(lines=5, label="Тэкст на беларускай мове"),
        gr.Audio(type="filepath", label="Прыклад голасу (без іншых гукаў) не карацей 7 секунд", interactive=True),
    ],
    outputs=gr.Audio(label="Згенераванае аўдыя"),
    title="Belarusian TTS Demo",
    description="Увядзіце тэкст, і мадэль пераўтворыць яго ў аўдыя. Вы можаце выкарыстоўваць голас па змаўчанні, абраць голас з прыкладаў унізе ці загрузіць уласны файл або запісаць аўдыё. Выкарыстоўваце прыклады з добрай якасцю, без іншых гукаў і разнастайнай інтанацыяй, ад яе моцна залежыць вынік. Таксама яна ўплывае на націскі. Прыклады галасоў могуць быць на любой мове. /n Каб палепшыць якасць мадэлі, як па націскам, так і па дакладнасці кланавання голасаў, трэба больш датасэтаў. Ахвяруйце свой голас праз Donar.by, запішыце свой голас з высокай якасцю станьце ўзорным донарам (запісацца тут https://forms.gle/TTirxTJNd3Ngw3dD6) /n Далучайцеся да беларускаймоўнай суполкі у ТГ каб дапамагчы ці даведацца пра навіны ШІ https://t.me/belarusai/",
    examples=examples,
    cache_examples=False
)

if __name__ == "__main__":
    demo.launch()