import os
os.environ.setdefault("MPLBACKEND", "Agg")

import gradio as gr
import sherpa_onnx
import tempfile
from huggingface_hub import hf_hub_download

MODEL_REPO = "willwade/mms-tts-multilingual-models-onnx"
MODEL_FILE = "bak/model.onnx"
TOKENS_FILE = "bak/tokens.txt"

model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
tokens_path = hf_hub_download(repo_id=MODEL_REPO, filename=TOKENS_FILE)

config = sherpa_onnx.OfflineTtsConfig(
    model=sherpa_onnx.OfflineTtsModelConfig(
        vits=sherpa_onnx.OfflineTtsVitsModelConfig(
            model=model_path,
            tokens=tokens_path,
            noise_scale=0.667,
            length_scale=1.0,
            noise_scale_w=0.8,
        ),
        num_threads=2,
    ),
)

tts = sherpa_onnx.OfflineTts(config)


def tts_fn(text):
    if not text.strip():
        return None

    audio = tts.generate(text.strip(), sid=0, speed=1.0)
    samples = audio.samples

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    sherpa_onnx.write_wave(tmp.name, samples, audio.sample_rate)
    return tmp.name


app = gr.Blocks(title="Башҡортса TTS")

with app:
    gr.Markdown("# 🎙️ Башҡортса TTS")
    gr.Markdown(
        "Офлайн озвучка текста на башкирском. "
        "Модель: MMS-TTS (Meta) via sherpa-onnx"
    )

    with gr.Row():
        with gr.Column(scale=3):
            text = gr.Textbox(
                label="Текст (башкирский кириллица)",
                placeholder="Бөгөн көн бик матур...",
                lines=6,
                value="Бөгөн көн бик матур. Һаумыһығыҙ, хәлдәрегеҙ нисек?",
            )
            with gr.Row():
                btn = gr.Button("▶ Озвучить", variant="primary")
                clear = gr.Button("✕ Очистить")

        with gr.Column(scale=2):
            audio = gr.Audio(label="Аудио", type="filepath")
            status = gr.Textbox(
                label="Статус", value="Готов", interactive=False
            )

    btn.click(fn=tts_fn, inputs=text, outputs=audio, api_name="predict")
    text.submit(fn=tts_fn, inputs=text, outputs=audio)
    clear.click(lambda: ("", None, "Готов"), None, [text, audio, status])

    gr.Markdown("---")
    gr.Markdown(
        "Сәләм, быны мин эшләнем — "
        "[instagram.com/ttaymas](https://instagram.com/ttaymas)"
    )
    gr.Markdown(
        "**Модель:** facebook/mms-tts-bak (VITS) · "
        "**Лицензия:** CC-BY-NC 4.0"
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)