#!/usr/bin/env python3
"""
VibeVoice 分角色转录(手动按需后端,非自动管线)。
把任意音频转成带说话人标签 + 时间戳的草稿。默认管线仍是 Qwen3(asr-transcribe.sh);
这个工具用于"需要谁说了什么"的场景。

用法:
  vibevoice-transcribe.py <audio> [--out DIR] [--chunk 200] [--unify]

输出:
  <out>/<basename>-diarized.md   # 带 `mm:ss` 时间戳 + S# 说话人(每块独立编号)
  <out>/<basename>-diarized.txt  # 纯文本版
  --unify 时额外 -global.md      # ECAPA+KMeans 全局统一说话人(需 venv,单麦只能锁主讲)

注意(踩坑见 vault: ASR服务横评-2026-06):
  - 服务端 segments 是 JSON 字符串,需 json.loads;非语音段无 Speaker 键。
  - 200s 是最佳分段(被服务端 max_time 限,单请求 ~280s 音频上限)。
  - 仅出草稿,正式纪要仍需 Kai 确认(memo-inbox 纪律)。
"""
import json, subprocess, os, math, re, sys, argparse

URL = os.environ.get("VIBE_URL", "http://192.168.50.100:10112/v1/audio/transcriptions")
MODEL = "vibevoice-asr-hf-nf4"
GLOSSARY = os.path.expanduser("~/obsidian-vault/minutes/_glossary.md")
SEG_RE = re.compile(r'\{"Start":\s*([\d.]+)\s*,\s*"End":\s*([\d.]+)\s*,\s*"Speaker":\s*(\d+|null)\s*,\s*"Content":\s*"((?:[^"\\]|\\.)*)"\s*\}')


def build_prompt():
    """Reuse glossary wikilinks + terms as a hotword prompt (same spirit as asr-transcribe.sh)."""
    terms = []
    p = GLOSSARY
    if os.path.exists(p):
        for raw in open(p, errors="ignore").read().splitlines():
            s = raw.strip()
            if not s or s.startswith("#") or s.startswith(">"):
                continue
            terms += [m.strip() for m in re.findall(r'\[\[([^\]|#]+)', s)]
            s = re.sub(r'^[-*+\d\.\)\s]+', '', s)
            s = re.split(r'\s[—–-]\s|[:：]', s, maxsplit=1)[0].strip()
            s = re.sub(r'\[\[|\]\]', '', s).strip()
            if s and len(s) <= 80:
                terms.append(s)
    seen, out = set(), []
    for t in terms:
        if t and t not in seen:
            seen.add(t); out.append(t)
    if not out:
        return ""
    return "以下是会议中可能出现的人名、公司、项目、产品和术语，请优先按这些词纠正识别错误：" + " ".join(out[:200])


def parse_segs(d):
    raw = d.get("segments")
    arr = raw if isinstance(raw, list) else None
    if arr is None and isinstance(raw, str):
        try: arr = json.loads(raw)
        except Exception: arr = None
    if arr is None:
        body = re.sub(r'^\s*assistant\s*', '', d.get("text", ""))
        arr = [{"Start": float(a), "End": float(b),
                "Speaker": (None if c == "null" else int(c)), "Content": w}
               for a, b, c, w in SEG_RE.findall(body)]
    out = []
    for s in arr:
        if not isinstance(s, dict):
            continue
        cc = (s.get("Content") or "").strip()
        if not cc:
            continue
        out.append({"Start": float(s.get("Start", 0)), "End": float(s.get("End", 0)),
                    "Speaker": s.get("Speaker"), "Content": cc})
    return out


def is_loop(cc):
    return len(cc) > 200 and cc[:18] and cc.count(cc[:18]) > 8


def hms(x):
    x = int(x); return f"{x // 60:02d}:{x % 60:02d}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("audio")
    ap.add_argument("--out", default=None, help="output dir (default: alongside audio)")
    ap.add_argument("--chunk", type=int, default=200)
    ap.add_argument("--unify", action="store_true", help="also run ECAPA global speaker unification")
    args = ap.parse_args()

    audio = os.path.abspath(os.path.expanduser(args.audio))
    if not os.path.isfile(audio):
        sys.exit(f"audio not found: {audio}")
    base = re.sub(r'\.[^.]+$', '', os.path.basename(audio))
    outdir = os.path.abspath(args.out) if args.out else os.path.dirname(audio)
    os.makedirs(outdir, exist_ok=True)
    work = os.path.join(outdir, f".{base}.vibework"); os.makedirs(work, exist_ok=True)

    prompt = build_prompt()
    dur = float(subprocess.check_output(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=nk=1:nw=1", audio]).decode().strip())
    n = math.ceil(dur / args.chunk)
    print(f"[vibevoice] {os.path.basename(audio)} dur={dur:.0f}s -> {n} chunks of {args.chunk}s", flush=True)

    all_segs = []
    for i in range(n):
        start = i * args.chunk
        length = min(args.chunk, dur - start)
        mp3 = os.path.join(work, f"c{i:02d}.mp3")
        subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-ss", str(start),
                        "-t", str(length), "-i", audio, "-ar", "16000", "-ac", "1", "-b:a", "64k", mp3], check=True)
        out_json = os.path.join(work, f"c{i:02d}.json")
        mnt = min(4096, 1024 + args.chunk * 12)
        rc = subprocess.run(["curl", "-sS", "--max-time", "260", "-F", f"model={MODEL}",
                             "-F", f"file=@{mp3}", "-F", f"prompt={prompt}",
                             "-F", f"max_new_tokens={mnt}", URL, "-o", out_json]).returncode
        try:
            d = json.load(open(out_json)); segs = parse_segs(d)
        except Exception as e:
            print(f"  chunk {i+1}/{n} FAIL ({e}) rc={rc}", flush=True); continue
        loops = sum(1 for s in segs if is_loop(s["Content"]))
        last = max((s["End"] for s in segs), default=0)
        for s in segs:
            cc = s["Content"]
            if len(cc) > 400:
                cc = cc[:120] + " …[复读截断]"
            all_segs.append({"Start": start + s["Start"], "End": start + s["End"],
                             "Speaker": s["Speaker"], "Content": cc})
        tag = f" ⚠loop×{loops}" if loops else ""
        print(f"  chunk {i+1}/{n} @{start:.0f}s: {len(segs)} segs, cov->{last:.0f}/{length:.0f}s{tag}", flush=True)

    all_segs.sort(key=lambda x: x["Start"])
    md = [f"# {base} · VibeVoice 分角色草稿",
          f"> 音频 {dur/60:.0f}min · {args.chunk}s 切块 · 说话人 S# **每块独立编号**(块间需对齐) · 仅草稿待确认\n"]
    txt = []
    for s in all_segs:
        spk = f"S{s['Speaker']}" if s["Speaker"] is not None else "·"
        md.append(f"- `{hms(s['Start'])}` **{spk}**: {s['Content']}")
        txt.append(f"[{hms(s['Start'])}] {spk}: {s['Content']}")
    md_path = os.path.join(outdir, f"{base}-diarized.md")
    txt_path = os.path.join(outdir, f"{base}-diarized.txt")
    open(md_path, "w").write("\n".join(md))
    open(txt_path, "w").write("\n".join(txt))
    json.dump(all_segs, open(os.path.join(work, "segs.json"), "w"), ensure_ascii=False)
    print(f"[vibevoice] {len(all_segs)} segs -> {md_path}", flush=True)

    if args.unify:
        unify = os.path.join(os.path.dirname(os.path.abspath(__file__)), "vibevoice-speaker-unify.py")
        venv_py = os.environ.get("ASR_VENV_PY",
                                 os.path.expanduser("~/claude-workspace/OBSIDIAN/recordings/asr-venv/bin/python"))
        if os.path.exists(venv_py) and os.path.exists(unify):
            print("[vibevoice] running ECAPA global speaker unification ...", flush=True)
            subprocess.run([venv_py, unify, audio, os.path.join(work, "segs.json"),
                            os.path.join(outdir, f"{base}-global.md")])
        else:
            print(f"[vibevoice] --unify skipped: venv ({venv_py}) or unify script missing. "
                  f"See vault ASR服务横评 to rebuild venv.", flush=True)


if __name__ == "__main__":
    main()