#!/usr/bin/env python3
"""tts.py - text to speech with Kokoro-82M (open weights, Apache 2.0).

    python tts.py "Hello, this is a test." --out hello.wav              # female, American English
    python tts.py "Hello." --voice am_michael --out hello_m.wav          # male
    python tts.py --file chapter.txt --voice bf_emma --lang b --out chapter.wav
    python tts.py "Mixed." --voice af_heart,am_michael --out mix.wav     # blend of two voices
    python tts.py --list-voices

Setup (once):  pip install kokoro soundfile      (downloads ~330 MB of weights on first run)
               Languages other than English need:  apt/brew install espeak-ng
Output: 24 kHz mono 16-bit WAV. Long texts are split at blank lines / sentences,
synthesized chunk by chunk and concatenated, so there is no length limit.
"""
import argparse, sys, time

# Kokoro-82M v1.0 voices (prefix: a/b = American/British English, f/m = female/male)
VOICES = {
    "female American": ["af_heart", "af_bella", "af_nicole", "af_sarah", "af_sky", "af_aoede", "af_kore", "af_nova", "af_river", "af_jessica", "af_alloy"],
    "male American":   ["am_michael", "am_adam", "am_fenrir", "am_puck", "am_echo", "am_eric", "am_liam", "am_onyx", "am_santa"],
    "female British":  ["bf_emma", "bf_isabella", "bf_alice", "bf_lily"],
    "male British":    ["bm_george", "bm_lewis", "bm_daniel", "bm_fable"],
    "other languages": ["ef_dora em_alex (es)", "ff_siwis (fr)", "hf_alpha hm_omega (hi)", "if_sara im_nicola (it)",
                        "jf_alpha jm_kumo (ja)", "pf_dora pm_alex (pt-br)", "zf_xiaobei zm_yunjian (zh)"],
}
LANGS = {"a": "American English", "b": "British English", "e": "Spanish", "f": "French", "h": "Hindi",
         "i": "Italian", "j": "Japanese", "p": "Portuguese (BR)", "z": "Mandarin"}   # no German in Kokoro-82M v1.0


def synthesize(text: str, voice: str = "af_heart", lang: str = "a", speed: float = 1.0):
    """Returns (samples: float32 numpy array, sample_rate=24000). Loads the model on first call."""
    import numpy as np
    from kokoro import KPipeline  # heavy import (torch); keep it inside the function
    global _PIPELINES
    try:
        _PIPELINES
    except NameError:
        _PIPELINES = {}
    if lang not in _PIPELINES:
        _PIPELINES[lang] = KPipeline(lang_code=lang)
    chunks = []
    # the pipeline yields (graphemes, phonemes, audio) per text chunk; split_pattern decides the chunks
    for _graphemes, _phonemes, audio in _PIPELINES[lang](text, voice=voice, speed=speed, split_pattern=r"\n+|(?<=[.!?])\s+"):
        if audio is not None:
            chunks.append(audio.numpy() if hasattr(audio, "numpy") else np.asarray(audio))
            chunks.append(np.zeros(int(24000 * 0.25), dtype=np.float32))  # 250 ms pause between chunks
    if not chunks:
        raise ValueError("no audio produced - empty text?")
    return np.concatenate(chunks).astype(np.float32), 24000


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("text", nargs="?", help="text to speak (or use --file)")
    ap.add_argument("--file", help="read the text from this file")
    ap.add_argument("--voice", default="af_heart", help="voice name or 'name1,name2' blend (default af_heart = female)")
    ap.add_argument("--lang", default="a", choices=LANGS, help="language/G2P: a=American, b=British English, ...")
    ap.add_argument("--speed", type=float, default=1.0, help="0.5 .. 2.0")
    ap.add_argument("--out", default="out.wav", help="WAV file to write")
    ap.add_argument("--play", action="store_true", help="also play it (needs 'pip install sounddevice')")
    ap.add_argument("--list-voices", action="store_true")
    a = ap.parse_args()
    if a.list_voices:
        for group, names in VOICES.items():
            print(f"{group:17s} {' '.join(names)}")
        print("languages:", ", ".join(f"{k}={v}" for k, v in LANGS.items()))
        return 0
    text = open(a.file, encoding="utf-8").read() if a.file else a.text
    if not text or not text.strip():
        ap.error("give a text or --file")
    t0 = time.time()
    samples, sr = synthesize(text, a.voice, a.lang, a.speed)
    import soundfile as sf
    sf.write(a.out, samples, sr)
    print(f"wrote {a.out}: {len(samples) / sr:.1f} s of audio, voice {a.voice}, {time.time() - t0:.1f} s")
    if a.play:
        import sounddevice as sd
        sd.play(samples, sr); sd.wait()
    return 0


if __name__ == "__main__":
    sys.exit(main())