#!/usr/bin/env python3
"""CuteTTS (AX650) command line demo.

    python3 demo.py --model-dir models --text "hello" --reference ref.wav --output out.wav
"""
import argparse
import os
import sys

import numpy as np

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

from cutetts_ax import CuteTTSAX  # noqa: E402
from cutetts_ax import audio as A  # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model-dir", default=os.environ.get("CUTETTS_MODEL_DIR", "models"))
    ap.add_argument("--assets", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets"))
    ap.add_argument("--engine", default="ax", choices=["ax", "onnx"])
    ap.add_argument("--text", default="The voice generated by this model sounds amazing!")
    ap.add_argument("--reference", default="samples/reference.wav")
    ap.add_argument("--output", default="output.wav")
    ap.add_argument("--cfg", type=float, default=None)
    ap.add_argument("--max-new", type=int, default=None)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--ref-seconds", type=float, default=5.0)
    ap.add_argument("--stream", action="store_true", help="true streaming: emit audio while generating")
    ap.add_argument("--first-chunk", type=int, default=16, help="latents in the first streamed chunk")
    ap.add_argument("--chunk", type=int, default=32, help="latents per streamed chunk afterwards")
    ap.add_argument("--windowed", action="store_true",
                    help="exact windowed streaming (bit-identical to offline, higher latency)")
    a = ap.parse_args()

    tts = CuteTTSAX(a.model_dir, a.assets, engine=a.engine)
    if a.stream:
        import wave
        with wave.open(a.output, "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(tts.sr)
            for chunk in tts.synthesize_stream(a.text, a.reference, max_new=a.max_new,
                                               cfg_strength=a.cfg, seed=a.seed,
                                               ref_max_seconds=a.ref_seconds,
                                               first_chunk=a.first_chunk, chunk=a.chunk,
                                               windowed=a.windowed):
                pcm = (np.clip(chunk, -1.0, 1.0) * 32767.0).astype("<i2")
                wf.writeframes(pcm.tobytes())
        tm = tts.last_timing
        print(f"streaming: {tm['chunks']} chunks, TTFT {tm['ttft']:.2f}s, "
              f"audio {tm['audio_seconds']:.2f}s in {tm['total']:.2f}s "
              f"(RTF {tm['total']/max(tm['audio_seconds'],1e-6):.3f})")
        print("saved", a.output)
        return
    wav = tts.synthesize(a.text, a.reference, max_new=a.max_new, cfg_strength=a.cfg,
                         seed=a.seed, ref_max_seconds=a.ref_seconds)
    A.write_wav(a.output, wav, tts.sr)
    tm = tts.last_timing
    print(f"audio {tm['audio_seconds']:.2f}s in {tm['total']:.2f}s "
          f"(RTF {tm['total']/max(tm['audio_seconds'],1e-6):.2f}), {tm['steps']} patches")
    print("  vae_enc %.3f  speaker %.3f  prompt %.3f  prefill %.3f  ar %.3f "
          "(dit %.3f, lm_dec %.3f)  vae_dec %.3f"
          % (tm["vae_enc"], tm["speaker_enc"], tm["prompt"], tm["prefill"], tm["ar_loop"],
             tm["dit"], tm["lm_decode"], tm["vae_dec"]))
    print("saved", a.output)


if __name__ == "__main__":
    main()