# Калі запускаеце ў чыстым асяроддзі (раскаментуйце):
# !pip install -q gradio spaces huggingface_hub torch scipy tqdm gitpython

import os
import sys
import time
import tempfile
import subprocess

import spaces
import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from scipy.io.wavfile import write
import numpy as np
from tqdm import tqdm

# ---------------------------------------------------------
# 1) Клануем і падключаем coqui-ai-TTS (fork з падтрымкай BE)
# ---------------------------------------------------------
REPO_URL = "https://github.com/tuteishygpt/coqui-ai-TTS.git"
REPO_DIR = "coqui-ai-TTS"

if not os.path.exists(REPO_DIR):
    subprocess.run(["git", "clone", REPO_URL, REPO_DIR], check=True)

repo_root = os.path.abspath(REPO_DIR)
if repo_root not in sys.path:
    sys.path.insert(0, repo_root)

from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
from TTS.tts.layers.xtts.tokenizer import split_sentence, VoiceBpeTokenizer

# ---------------------------------------------------------
# 2) Файлы мадэлі
# ---------------------------------------------------------
repo_id = "archivartaunik/BE_XTTS_V2_10ep250k"
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)

# ---------------------------------------------------------
# 3) Загрузка мадэлі і токенайзера
# ---------------------------------------------------------
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).eval()
sampling_rate = int(XTTS_MODEL.config.audio["sample_rate"])

tokenizer = VoiceBpeTokenizer(vocab_file=vocab_file)
XTTS_MODEL.tokenizer = tokenizer

# ---------------------------------------------------------
# 4) Патокавая TTS-функцыя: кожны чанк да канца, потым наступны
# ---------------------------------------------------------
@spaces.GPU(duration=60)
def text_to_speech(belarusian_story, speaker_audio_file=None):
    """
    Streaming для gr.Audio:
      - На кожным кроку: yield (sr, chunk) — прайграецца толькі гэты чанк.
      - Перад наступным yield робім sleep на даўжыню chunk, каб першы праграўся ДА КАНЦА.
      - У фінале: yield шлях да поўнага WAV (для загрузкі).
    """
    if not belarusian_story or str(belarusian_story).strip() == "":
        raise gr.Error("Увядзі хоць нейкі тэкст 🙂")

    # Голас па змаўчанні
    if not speaker_audio_file or (
        not isinstance(speaker_audio_file, str)
        and getattr(speaker_audio_file, "name", "") == ""
    ):
        speaker_audio_file = default_voice_file

    # Conditioning latents
    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:
        raise gr.Error(f"Памылка пры атрыманні латэнтаў голасу: {e}")

    # Разбіўка на сказы/чанкі
    try:
        lang = "be"
        chunk_limit = tokenizer.char_limits.get(lang, 250)
        tts_texts = split_sentence(
            str(belarusian_story).strip(),
            lang=lang,
            text_split_length=chunk_limit,
        )
        tts_texts = [s.strip() for s in tts_texts if s and s.strip()]
        if not tts_texts:
            raise gr.Error("Не атрымалася падзяліць тэкст на сказы/чанкі.")
    except Exception as e:
        raise gr.Error(f"Памылка пры падзеле тэксту на сказы: {e}")

    all_chunks = []

    for text in tqdm(tts_texts):
        try:
            # 1) генеруем чанк
            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,
                )
            cur = wav_chunk["wav"].astype(np.float32)
            all_chunks.append(cur)

            # 2) адразу аддаём яго ў плэер
            yield (sampling_rate, cur)

            # 3) чакаем, пакуль ён ПРАЙГРАЕЦЦА цалкам
            #    (грубая, але надзейная сінхранізацыя: па даўжыні chunk)
            duration = len(cur) / float(sampling_rate)
            # невялікі запас, каб не перакрываць канец
            time.sleep(duration + 0.05)

        except Exception as e:
            raise gr.Error(f"Памылка пры генерырацыі аўдыя: {e}")

    if not all_chunks:
        raise gr.Error("Нічога не згенеравана. Праверце ўваходныя даныя.")

    # Фінальны WAV у temp-файл
    try:
        final_audio = np.concatenate(all_chunks, axis=0)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
        write(temp_file.name, sampling_rate, final_audio)
        yield temp_file.name
    except Exception as e:
        raise gr.Error(f"Памылка пры запісе фінальнага WAV: {e}")

# ---------------------------------------------------------
# 5) Прыклады (тэкст + файл голасу)
# ---------------------------------------------------------
examples = [
    [
        "Такім чынам, клуб стаў уладальнікам усіх існых на сёння міжнародных трафеяў паўднёваамерыканскага футболу.",
        "Nestarka.wav",
    ],
    [
        "Яму не ўдалося палепшыць фінансавае становішча каралеўства, а, наадварот, прыйшлося распрадаваць каштоўнасці чэшскай кароны.",
        "muzh.wav",
    ],
    [
        "Кампілятарамі называюць праграмы, якія пераўтвараюць код вышэйшага ўзроўню ў код ніжэйшага ўзроўню.",
        "chunk_100.wav",
    ],
    [
        "Акрамя таго, ліхачы аддаюць перавагу рэгі, хіп-хопу і класічнай музыцы.",
        "d1015.mp3",
    ],
    [
        "Позірк можа быць уважлівым, зацікаўленым, захопленым, але бывае і нахабным, задзірлівым, пагардлівым, напышлівым.",
        "donarka_ench.wav",
    ],
    [
        "Такі нават шчыры, ці што: родная мова народу – трасянка, а беларуская яму чужая!",
        "muzhcynski.wav",
    ],
]

analytics_script = """
<script async src="https://www.googletagmanager.com/gtag/js?id=G-TKDCRCQ7FK"></script>
<script>
  window.dataLayer = window.dataLayer || [];
  function gtag(){dataLayer.push(arguments);}
  gtag('js', new Date());
  gtag('config', 'G-TKDCRCQ7FK');
</script>
"""

# ---------------------------------------------------------
# 6) Gradio UI (autoplay)
# ---------------------------------------------------------
with gr.Blocks() as demo:
    gr.HTML(analytics_script)
    gr.Interface(
        fn=text_to_speech,  # генератар
        inputs=[
            gr.Textbox(lines=5, label="Тэкст на беларускай мове"),
            gr.Audio(
                type="filepath",
                label="Прыклад голасу (без іншых гукаў) не карацей 7 секунд",
                interactive=True,
            ),
        ],
        outputs=gr.Audio(
            type="filepath",   # прымае (sr, ndarray) і фінальны шлях
            label="Згенераванае аўдыя (па чарзе, без абрываў)",
            autoplay=True,
        ),
        title="Belarusian TTS Demo — Streaming (па чарзе, без абрываў)",
        description="""
        <p>Кожны чанк прайграецца <b>да канца</b>, потым запускаецца наступны. У фінале будзе даступная загрузка аб’яднанага WAV.</p>
        """,
        examples=examples,
        cache_examples=False,
        allow_flagging="never",
    )

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