#!/usr/bin/env python3
"""Generate the PrimeTTS training corpus with VoxCPM2, cloning ONE reference (young-girl 'stacy')
so en + zh + code-mix are all in a single consistent voice. Output 48 kHz (VoxCPM2 audiovae SR).
Resumable (skips ids already in the manifest). Usage:
  CUDA_VISIBLE_DEVICES=N python gen_voxcpm_corpus.py --texts <jsonl> --manifest <out.jsonl>
"""
import argparse, json, os, time
import numpy as np, soundfile as sf
import text_norm                                  # entity normalizer (phone/email/price/serial/date...)

SR = 48000

def trim(w, thr=0.02):
    x = np.abs(w)
    if x.max() < 1e-5: return w
    idx = np.where(x > thr * x.max())[0]
    if len(idx) == 0: return w
    return w[max(0, idx[0] - int(0.05*SR)):min(len(w), idx[-1] + int(0.12*SR))]

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--texts", required=True)
    ap.add_argument("--ref", default="clone_ref_girl.wav")
    ap.add_argument("--ref-text", default="clone_ref_girl.txt")
    ap.add_argument("--out-dir", default="voxcpm_corpus")
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--no-normalize", dest="normalize", action="store_false",
                    help="disable text_norm entity pre-normalization (default: on)")
    a = ap.parse_args()
    os.makedirs(a.out_dir, exist_ok=True)
    rt = open(a.ref_text).read().strip()
    from voxcpm import VoxCPM
    m = VoxCPM.from_pretrained("openbmb/VoxCPM2")
    done = set()
    if os.path.exists(a.manifest):
        for l in open(a.manifest):
            try: done.add(json.loads(l)["id"])
            except Exception: pass
    mf = open(a.manifest, "a", encoding="utf-8")
    rows = [json.loads(l) for l in open(a.texts) if l.strip()]
    t0 = time.time(); n = 0
    for r in rows:
        if r["id"] in done: continue
        out = os.path.join(a.out_dir, r["id"] + ".wav")
        # PRE-NORMALIZE entities (phone/email/price/serial/date...) so the teacher reads the spoken form
        # and the manifest text matches the audio (digit-by-digit, not cardinalized). Idempotent.
        txt = text_norm.normalize(r["text"]) if a.normalize else r["text"]
        try:
            w = np.asarray(m.generate(text=txt, prompt_wav_path=a.ref, prompt_text=rt), dtype="float32").reshape(-1)
            w = trim(w)
            if len(w) < int(0.3*SR): print("SHORT skip", r["id"], flush=True); continue
            sf.write(out, w, SR)
            mf.write(json.dumps({"id": r["id"], "text": txt, "lang": r["lang"],
                                 "target_audio": os.path.abspath(out), "dur": round(len(w)/SR, 2)}, ensure_ascii=False) + "\n")
            mf.flush(); n += 1
            if n % 20 == 0: print(f"{n} done | {n/(time.time()-t0)*3600:.0f}/h", flush=True)
        except Exception as e:
            print("FAIL", r["id"], str(e)[:90], flush=True)
    print(f"DONE {n} new clips", flush=True)

if __name__ == "__main__":
    main()