import os
import gradio as gr
import soundfile as sf
from huggingface_hub import snapshot_download

MODEL_REPO = "LyngualLabs/YorubaEnglish-CodeSwitching-TTS"
DEFAULT_REF_WAV = "reference.wav"
DEFAULT_REF_TEXT = "Ìròyìn ti sọ pé the government will ensure electricity tariff goes down ní January."

print("Downloading model...", flush=True)
model_dir = snapshot_download(MODEL_REPO)

from voxcpm import VoxCPM  # noqa: E402  (import after snapshot_download so the repo is cached)

print("Loading model...", flush=True)
model = VoxCPM.from_pretrained(model_dir, load_denoiser=False)
SR = model.tts_model.sample_rate
print("Model loaded.", flush=True)


def synthesize(text, ref_audio, ref_text, cfg_value, inference_timesteps):
    if not text or not text.strip():
        raise gr.Error("Please enter some text to synthesize.")

    prompt_wav_path = ref_audio if ref_audio else DEFAULT_REF_WAV
    prompt_text = ref_text.strip() if (ref_audio and ref_text and ref_text.strip()) else (
        DEFAULT_REF_TEXT if not ref_audio else ""
    )
    if ref_audio and not prompt_text:
        raise gr.Error("Please provide the exact transcript of your uploaded reference audio.")

    wav = model.generate(
        text=text.strip(),
        prompt_wav_path=prompt_wav_path,
        prompt_text=prompt_text,
        cfg_value=cfg_value,
        inference_timesteps=int(inference_timesteps),
    )
    out_path = "output.wav"
    sf.write(out_path, wav, SR)
    return out_path


EXAMPLES = [
    ["Mo ní meeting pẹ̀lú marketing team ní aago mẹ́wàá lónìí, nítorí náà ẹ jẹ́ ká yára."],
    ["Ìròyìn ti sọ pé the government will ensure electricity tariff goes down ní January."],
    ["Wọ́n ní ìpàdé ní hotel tó wà ní Victoria Island, kì í ṣe ní office wa."],
]

with gr.Blocks(title="Yoruba-English Code-Switching TTS") as demo:
    gr.Markdown(
        """
        # 🎙️ Yoruba–English Code-Switching TTS
        A voice-cloning text-to-speech model for Yoruba, including natural
        Yoruba–English code-switching — full fine-tuned from
        [VoxCPM2](https://huggingface.co/openbmb/VoxCPM2) on ~1,039 hours pooled from
        DSN African Voices, NaijaVoices, YECS, and WAXAL.

        Model card: [LyngualLabs/YorubaEnglish-CodeSwitching-TTS](https://huggingface.co/LyngualLabs/YorubaEnglish-CodeSwitching-TTS)

        ⚠️ **Experimental** — trained for only ~1 epoch. Quality varies by sentence;
        rarer words/proper nouns are more likely to be mispronounced.
        """
    )
    with gr.Row():
        with gr.Column():
            text_in = gr.Textbox(
                label="Text to speak",
                placeholder="Type Yoruba, English, or mixed Yoruba-English text...",
                lines=4,
            )
            with gr.Accordion("Voice cloning (optional) — upload your own reference", open=False):
                ref_audio = gr.Audio(label="Reference audio (3-10s, clean)", type="filepath")
                ref_text = gr.Textbox(label="Exact transcript of the reference audio")
                gr.Markdown("_Leave blank to use the default built-in voice._")
            with gr.Accordion("Advanced", open=False):
                cfg_value = gr.Slider(1.0, 4.0, value=2.0, step=0.1, label="cfg_value (higher = sticks closer to reference)")
                inference_timesteps = gr.Slider(10, 35, value=22, step=1, label="inference_timesteps (higher = smoother, slower)")
            btn = gr.Button("Generate", variant="primary")
        with gr.Column():
            audio_out = gr.Audio(label="Generated speech", type="filepath")

    gr.Examples(examples=EXAMPLES, inputs=[text_in])
    btn.click(
        synthesize,
        inputs=[text_in, ref_audio, ref_text, cfg_value, inference_timesteps],
        outputs=[audio_out],
    )

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