import argparse
import json
import os
import sys
import time
from pathlib import Path

os.environ.setdefault("TQDM_DISABLE", "1")  # keep sampling bars out of job logs

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

text = (p.get("text") or "").strip()
if not text:
    print("ERROR: text is required")
    sys.exit(1)
if not a.input:
    print("ERROR: needs a reference voice WAV as input")
    sys.exit(1)

import torch  # noqa: E402
import torchaudio  # noqa: E402
from chatterbox.tts import ChatterboxTTS  # noqa: E402

device = "mps" if torch.backends.mps.is_available() else "cpu"
t0 = time.time()
model = ChatterboxTTS.from_pretrained(device=device)
print(f"model loaded on {device} in {time.time() - t0:.1f}s", flush=True)

t1 = time.time()
wav = model.generate(
    text,
    audio_prompt_path=a.input[0],
    exaggeration=float(p.get("exaggeration", 0.5)),
    cfg_weight=float(p.get("cfg_weight", 0.5)),
)
out = Path(a.outdir) / f"clone_{Path(a.input[0]).stem}.wav"
torchaudio.save(str(out), wav, model.sr)
dur = wav.shape[-1] / model.sr
print(f"done in {time.time() - t1:.1f}s: {out.name} ({dur:.1f}s)", flush=True)