"""Text -> 24 kHz WAV with the Maya 2 MLX port.

    .venv/bin/python -I speak.py "Hello there." -o hello.wav [--lang en] [--voice Malini] [--seed 7]
    .venv/bin/python -I speak.py --voices          # the named voices

Same defaults as the official Maya3TTS.generate: plain voice, English, ~20 s pieces with 0.15 s gaps, no text
normalisation unless --normalize, a fresh engine-wide seed per piece unless --seed.
"""
import argparse
import os
import sys
import time
import wave

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import numpy as np  # noqa: E402

from maya2_mlx import Maya2  # noqa: E402


def write_wav(path, wav, sr=24000):
    pcm = np.rint(np.asarray(wav, np.float32) * np.float32(32767)).astype("<i2")   # the engine's int16 samples
    with wave.open(path, "wb") as f:
        f.setnchannels(1)
        f.setsampwidth(2)
        f.setframerate(sr)
        f.writeframes(pcm.tobytes())


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("text", nargs="?")
    ap.add_argument("-o", "--out", default="maya2.wav")
    ap.add_argument("--lang", default="en")
    ap.add_argument("--voice", default=None)
    ap.add_argument("--seed", type=int, default=None)
    ap.add_argument("--temperature", type=float, default=None)
    ap.add_argument("--top-k", type=int, default=None)
    ap.add_argument("--top-p", type=float, default=None)
    ap.add_argument("--repetition-penalty", type=float, default=None)
    ap.add_argument("--max-frames", type=int, default=None)
    ap.add_argument("--normalize", action="store_true", help="run the bundled text normaliser first")
    ap.add_argument("--voices", action="store_true", help="list the voices and exit")
    ap.add_argument("--models", default=None, help="weights folder (default: models/maya2-8bit, else models/maya2)")
    a = ap.parse_args()
    tts = Maya2(a.models) if a.models else Maya2()
    if a.voices:
        for vid in tts.voices:
            m = tts.registry.get(vid).meta
            print(f"{vid:10s} {m.get('gender', ''):7s} {m.get('age_band', ''):7s} {m.get('description', '')}")
        return
    if not a.text:
        ap.error("text is required")
    t0 = time.perf_counter()
    wav, info = tts.generate(a.text, a.lang, a.voice, max_frames=a.max_frames, return_info=True, normalize=a.normalize,
                             seed=a.seed, temperature=a.temperature, top_k=a.top_k, top_p=a.top_p,
                             repetition_penalty=a.repetition_penalty)
    el = time.perf_counter() - t0
    write_wav(a.out, wav)
    print(f"{a.out}: {len(wav) / 24000:.2f} s audio in {el:.1f} s ({os.path.basename(tts.models_dir)}), pieces {info['pieces']}, "
          f"stops {[s['reason'] for s in info['stops']]}, seeds {[s['seed'] for s in info['stops']]}")


if __name__ == "__main__":
    main()