import spaces
import tempfile
import soundfile as sf
import torch
import torchaudio
import gradio as gr

# ---- PATCH: torchcodec missing on HF Spaces; use soundfile instead ----
import myanmar_tts.api as _api
from myanmar_tts.utils import audio as _audio_utils

def _sf_load(path, target_sr):
    wav, sr = sf.read(path, dtype="float32")
    wav = torch.from_numpy(wav).float()
    if wav.dim() == 1:
        wav = wav.unsqueeze(0)
    if sr != target_sr:
        wav = torchaudio.functional.resample(wav, sr, target_sr)
    return wav

_api.load_and_resample_audio = _sf_load
_audio_utils.load_and_resample_audio = _sf_load
# ---- END PATCH ----

from myanmar_tts import MyanmarTTS

_model = None


def _ensure_model():
    global _model
    if _model is None:
        _model = MyanmarTTS(device="cuda")
        _model._load()


@spaces.GPU(duration=30)
def generate_speech(text, step, cfg):
    _ensure_model()
    if not text or not text.strip():
        raise gr.Error("Please enter Burmese text.")
    audio = _model.tts(text.strip(), solver="euler", step=int(step), cfg=float(cfg))
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    sf.write(tmp.name, audio, 44100)
    return tmp.name


# Natural Burmese sentences (female voice appropriate)
EXAMPLES = [
    ["ဟယ်လို... ဒါလင်... မတွေ့ရတာ... ကြာပီ။", 32, 3.0],
    ["ဒီနေ့ ရာသီဥတု အလွန်ကောင်းမွန်ပါတယ်ရှင်။", 12, 3.0],
    ["လူသားတွေ အားလုံးကို အရမ်း ချစ်ပါတယ်ရှင့်။", 12, 3.0],
    ["မင်္ဂလာပါရှင်၊ ကျွန်မ မြန်မာလူမျိုး ဖြစ်ပါတယ်ရှင်။", 8, 3.0],
]

with gr.Blocks(title="MyanmarTTS") as demo:
    gr.Markdown("# MyanmarTTS\nBurmese Text-to-Speech, 31M params, CC0 1.0")
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(value="ဟယ်လို... ဒါလင်... မတွေ့ရတာ... ကြာပီ။", label="Burmese Text", lines=4)
            with gr.Row():
                step_slider = gr.Slider(4, 32, value=32, step=4, label="Steps")
                cfg_slider = gr.Slider(1.0, 5.0, value=3.0, step=0.5, label="CFG")
            btn = gr.Button("Generate", variant="primary")
        with gr.Column():
            out = gr.Audio(label="Output", type="filepath")
    gr.Examples(EXAMPLES, inputs=[text_input, step_slider, cfg_slider],
                outputs=out, fn=generate_speech, cache_examples=False)
    btn.click(fn=generate_speech,
              inputs=[text_input, step_slider, cfg_slider], outputs=out)

if __name__ == "__main__":
    demo.queue(max_size=20).launch()