#!/usr/bin/env python3
"""Generate the BreezyVoice teacher corpus for distillation into Inflect-Nano.
Clean short reference (no prompt-leak) + INLINE ASR GATE: each clip is transcribed
(faster-whisper) and kept only if Han-CER vs intended text is below threshold; else
retried, then skipped. Writes 22.05kHz wav + manifest {id,text,wav,dur,cer}. Resumable.
Run in .venv-breezy with PYTHONPATH=BreezyVoice repo.
"""
from __future__ import annotations
import argparse, json, os, sys, time
import soundfile as sf, numpy as np

BV = "/home/luigi/jetson-tts/third_party/BreezyVoice"
ZT = "/home/luigi/jetson-tts/mossnano/zhtw8k"
sys.path.insert(0, BV)
from single_inference import CustomCosyVoice, get_bopomofo_rare
from g2pw import G2PWConverter
from cosyvoice.utils.file_utils import load_wav
from faster_whisper import WhisperModel

REF_AUDIO = f"{ZT}/ref/ref_clean.wav"
REF_TEXT = open(f"{ZT}/ref/ref_clean.txt").read().strip()


import re as _re
import opencc as _opencc
_T2S = _opencc.OpenCC('t2s')  # normalize trad<->simp so the gate scores REAL pronunciation, not script
def han(s): return "".join(c for c in s if "一" <= c <= "鿿")
def is_zh(text): return bool(_re.search(r"[一-鿿]", text))
def _lev(r, h):
    if not r: return 0.0
    d = list(range(len(h)+1))
    for i in range(1, len(r)+1):
        prev = d[0]; d[0] = i
        for j in range(1, len(h)+1):
            cur = d[j]; d[j] = min(d[j]+1, d[j-1]+1, prev+(r[i-1] != h[j-1])); prev = cur
    return d[len(h)]/len(r)
def _enwords(s): return _re.findall(r"[a-z']+", s.lower())
def score(ref, hyp):
    """Han-CER for zh/mix; word-error-rate for pure-English. Lower = better."""
    if is_zh(ref):
        return _lev(han(_T2S.convert(ref)), han(_T2S.convert(hyp)))
    return _lev(_enwords(ref), _enwords(hyp))   # word-level Levenshtein ratio


def read_tsv(path, limit, skip):
    rows = []
    for line in open(path, encoding="utf-8"):
        line = line.rstrip("\n")
        if not line.strip(): continue
        parts = line.split("\t")
        rows.append((parts[0] if len(parts) > 1 else f"utt{len(rows):06d}", parts[-1].strip()))
    rows = rows[skip:]
    return rows[:limit] if limit else rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="/home/luigi/jetson-tts/models/BreezyVoice")
    ap.add_argument("--corpus", default="/home/luigi/jetson-tts/data/text/train.tsv")
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--skip", type=int, default=0)
    ap.add_argument("--min-sec", type=float, default=0.8)
    ap.add_argument("--max-sec", type=float, default=20.0)
    ap.add_argument("--cer-thresh", type=float, default=0.30)
    ap.add_argument("--retries", type=int, default=3)
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    man_path = os.path.join(args.out_dir, "manifest.jsonl")
    done = set()
    if os.path.exists(man_path):
        for l in open(man_path):
            try: done.add(json.loads(l)["id"])
            except Exception: pass
    print(f"resuming: {len(done)} done | ref: {REF_TEXT}")

    cv = CustomCosyVoice(args.model); conv = G2PWConverter()
    asr = WhisperModel("SoybeanMilk/faster-whisper-Breeze-ASR-25", device="cuda", compute_type="float16")  # zh-TW gate (traditional output)
    ref_bopo = get_bopomofo_rare(cv.frontend.text_normalize_new(REF_TEXT, split=False), conv)
    ref_wav = load_wav(REF_AUDIO, 16000)

    rows = read_tsv(args.corpus, args.limit, args.skip)
    todo = [(u, t) for u, t in rows if u not in done]
    print(f"to synth: {len(todo)} / {len(rows)}")
    mf = open(man_path, "a", encoding="utf-8")
    t0 = time.time(); n_ok = 0; n_skip = 0; tot = 0.0
    for i, (utt, text) in enumerate(todo):
        bopo = get_bopomofo_rare(cv.frontend.text_normalize_new(text, split=False), conv)
        best = None; best_cer = 9.9
        for attempt in range(args.retries):
            try:
                out = cv.inference_zero_shot_no_normalize(bopo, ref_bopo, ref_wav)
                w = out["tts_speech"].squeeze().cpu().numpy().astype(np.float32)
            except Exception as e:
                print(f"  [err {utt}] {e}"); continue
            dur = len(w) / 22050
            if dur < args.min_sec or dur > args.max_sec: continue
            tmp = os.path.join(args.out_dir, f".{utt}.tmp.wav"); sf.write(tmp, w, 22050)
            segs, _ = asr.transcribe(tmp, language=("zh" if is_zh(text) else "en"), beam_size=5)
            c = score(text, "".join(s.text for s in segs))
            if c < best_cer: best_cer = c; best = (w, dur, tmp)
            if c <= args.cer_thresh: break
        if best is None or best_cer > args.cer_thresh:
            n_skip += 1
            if best and os.path.exists(best[2]): os.remove(best[2])
            if n_skip <= 20 or n_skip % 50 == 0:
                print(f"  [SKIP {utt}] best_cer={best_cer:.2f} | {text[:30]}")
            continue
        w, dur, tmp = best
        wp = os.path.join(args.out_dir, f"{utt}.wav"); os.replace(tmp, wp)
        mf.write(json.dumps({"id": utt, "text": text, "wav": wp, "dur": round(dur, 3),
                             "cer": round(best_cer, 3)}, ensure_ascii=False) + "\n"); mf.flush()
        n_ok += 1; tot += dur
        if (i + 1) % 25 == 0:
            el = time.time() - t0
            print(f"  {i+1}/{len(todo)} ok={n_ok} skip={n_skip} audio={tot/60:.1f}min "
                  f"{el/(i+1):.2f}s/clip eta={(len(todo)-i-1)*el/(i+1)/60:.0f}min")
    mf.close()
    print(f"DONE ok={n_ok} skip={n_skip} audio={tot/60:.1f}min in {(time.time()-t0)/60:.0f}min")


if __name__ == "__main__":
    main()