"""Khmer TTS demo Space (free ZeroGPU tier).

Model weights live in a PRIVATE HF repo; this Space reads them at startup with
the HF_TOKEN secret (Space settings -> Variables and secrets). Visitors get the
demo, never the weights. ZeroGPU allocates a GPU slice per generation via the
@spaces.GPU decorator; visitors have per-day GPU quotas managed by HF.
"""

from __future__ import annotations

import os
from functools import lru_cache
from importlib.resources import files
from pathlib import Path

import gradio as gr
import numpy as np
import spaces

from kfe import chunk_khmer_text

MODEL_REPO = os.environ.get("MODEL_REPO", "CHANGE_ME/khmer-tts-v2-private")
MAX_CHARS = 500
VOICES = {
    "ស្រី (female)": ("ref_female.wav",
        "បើចង់ដឹងថាឆ្អិនល្អឬនៅយើងអាចសាកដោយយកឈើចាក់ធ្មេញចាក់ចូលទៅកណ្តាលសង់ខ្យា បើគ្មានជាប់មកវិញគឺឆ្អិនហើយ។"),
    "ប្រុស (male)": ("ref_male.wav",
        "នេះធ្វើឲ្យល្ពៅក្លាយជាបន្លែមួយដែលមានភាពបត់បែនខ្ពស់ និងពេញនិយម។"),
}


# load once at startup (ZeroGPU pattern: build on CPU at import, the decorator
# moves compute to the allocated GPU per call)
from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf

import f5_tts.model as m
from f5_tts.infer.utils_infer import (infer_process, load_model, load_vocoder,
                                      preprocess_ref_audio_text)

_token = os.environ.get("HF_TOKEN")
_ckpt = hf_hub_download(MODEL_REPO, "khmer_serving_v2.pt", token=_token)
_vocab = hf_hub_download(MODEL_REPO, "vocab.txt", token=_token)
_cfg = OmegaConf.load(str(files("f5_tts").joinpath("configs/F5TTS_v1_Base.yaml")))
MODEL = load_model(getattr(m, _cfg.model.backbone), _cfg.model.arch, _ckpt,
                   vocab_file=_vocab, use_ema=False)
VOCODER = load_vocoder("vocos")


@spaces.GPU(duration=90)
def tts(text: str, voice: str, nfe: int, progress=gr.Progress()):
    text = (text or "").strip()
    if not text:
        raise gr.Error("សូមបញ្ចូលអត្ថបទខ្មែរ / please enter Khmer text")
    if len(text) > MAX_CHARS:
        raise gr.Error(f"demo is capped at {MAX_CHARS} characters "
                       f"(got {len(text)}) — try a shorter text")

    clip, ref_text = VOICES[voice]
    ref_a, ref_t = preprocess_ref_audio_text(str(Path(__file__).parent / clip), ref_text)

    chunks = chunk_khmer_text(text)
    waves, sr = [], 24000
    for i, c in enumerate(chunks):
        progress((i, len(chunks)), desc=f"synthesizing chunk {i+1}/{len(chunks)}")
        audio, sr, _ = infer_process(ref_a, ref_t, c, MODEL, VOCODER,
                                     nfe_step=int(nfe))
        waves.append(audio)
        waves.append(np.zeros(int(0.18 * sr)))
    return sr, np.concatenate(waves[:-1])


with gr.Blocks(title="Khmer TTS") as demo:
    gr.Markdown(
        "# 🇰🇭 Khmer text-to-speech\n"
        "ការបម្លែងអត្ថបទជាសំឡេងខ្មែរ — fine-tuned on ~1,000 h of Khmer speech "
        "(F5-TTS base).\n\n"
        "Speech data: Digital Divide Data (CC-BY-SA-4.0). "
        "Non-commercial research demo."
    )
    text = gr.Textbox(label="អត្ថបទខ្មែរ (Khmer text)", lines=3,
                      value="សួស្តី! សូមស្វាគមន៍មកកាន់ការបម្លែងអត្ថបទជាសំឡេងខ្មែរ។")
    voice = gr.Radio(list(VOICES), value=list(VOICES)[0], label="សំឡេង (voice)")
    nfe = gr.Slider(16, 48, value=32, step=8,
                    label="quality steps (higher = better + slower)")
    btn = gr.Button("និយាយ (Speak)", variant="primary")
    out = gr.Audio(label="output", autoplay=True)
    btn.click(tts, [text, voice, nfe], out)

demo.queue(max_size=20).launch()