audio_to_text

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""本地课堂录音批量转文字（混合引擎：FunASR SeACo-Paraformer + faster-whisper）。

中文主导的录音路由到 FunASR SeACo-Paraformer（原生热词机制），
英语/法语等其他语言路由到 faster-whisper。按文件整体路由，不做句级混切。
所有推理均在本地完成，不使用云端 API。

用法：
    ./run.sh                      # 全自动：探测语言 → 路由引擎 → 转写
    ./run.sh --language zh        # 强制语言（zh→Paraformer；en/fr/…→Whisper）
    ./run.sh --engine whisper     # 强制引擎：paraformer | whisper
    ./run.sh --model large-v3     # 覆盖 Whisper 模型（默认 large-v3-turbo）
    ./run.sh --dry-run            # 只解码+探测+打印路由结果，不转写
    ./run.sh --download-only      # 只预下载全部模型到 models/ 后退出
"""

from __future__ import annotations

import argparse
import gc
import inspect
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import wave
from pathlib import Path

import numpy as np

# ---------------------------------------------------------------- 路径与常量

PROJECT_DIR = Path(__file__).resolve().parent
INPUT_DIR = PROJECT_DIR / "input_audio"
OUTPUT_DIR = PROJECT_DIR / "output_text"
MODELS_DIR = PROJECT_DIR / "models"
FUNASR_CACHE_DIR = MODELS_DIR / "funasr"
WHISPER_CACHE_DIR = MODELS_DIR / "faster-whisper"

# funasr(≥1.2) 不再读取 AutoModel 的 cache_dir 参数，统一走 modelscope 缓存，
# 只认 MODELSCOPE_CACHE 环境变量 —— 必须在 import funasr 之前设置，保证模型缓存在项目内
os.environ["MODELSCOPE_CACHE"] = str(FUNASR_CACHE_DIR)

SUPPORTED_EXTS = {".m4a", ".mp3", ".wav", ".flac", ".aac", ".ogg", ".mp4", ".wma", ".amr"}

PARAFORMER_MODEL = "iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch"
VAD_MODEL = "iic/speech_fsmn_vad_zh-cn-16k-common-pytorch"
PUNC_MODEL = "iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch"

WHISPER_MODEL_DEFAULT = "large-v3-turbo"
WHISPER_CPU_THREADS = 16          # 20 线程机器留一点给系统
WHISPER_BATCH_SIZE = 8            # 15GB 内存的保守批大小
WHISPER_PROMPT_HOTWORDS = 50      # initial_prompt 携带的热词个数（超长会被截断）

SAMPLE_RATE = 16000
PROBE_SLICE_SECONDS = 45.0        # 语言探测切片长度
ZH_PROB_THRESHOLD = 0.6           # zh 平均概率超过该值且为最高才走 Paraformer
SEGMENTS_PER_PARAGRAPH = 6        # Whisper 每 ~6 个 segment 合并成一个自然段
SENTENCES_PER_PARAGRAPH = 6       # Paraformer 每 ~6 句插一个空行

# 将来有 NVIDIA GPU：export A2T_DEVICE=cuda（代码保留 device 参数位）
DEVICE = os.environ.get("A2T_DEVICE", "cpu")
COMPUTE_TYPE = "int8" if DEVICE == "cpu" else "float16"

WS_RE = re.compile(r"[ \t\r\n\f\v]+")
SENT_RE = re.compile(r"[^。！？；]*[。！？；]+|[^。！？；]+$")


def log(msg: str) -> None:
    print(msg, flush=True)


def short_err(exc: BaseException, limit: int = 150) -> str:
    return re.sub(r"\s+", " ", str(exc))[:limit] or type(exc).__name__


# ---------------------------------------------------------------- CLI

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="run.sh",
        description="本地课堂录音批量转文字（中文主导→Paraformer，英/法等→Whisper）",
    )
    p.add_argument("--language", metavar="CODE", default=None,
                   help="强制语言：zh→Paraformer；en/fr/…→Whisper（语言码传给 Whisper）")
    p.add_argument("--engine", choices=("paraformer", "whisper"), default=None,
                   help="强制指定转写引擎")
    p.add_argument("--model", default=WHISPER_MODEL_DEFAULT,
                   help=f"Whisper 模型（默认 {WHISPER_MODEL_DEFAULT}）")
    p.add_argument("--dry-run", action="store_true",
                   help="只解码+探测+打印路由结果，不加载 FunASR、不转写")
    p.add_argument("--download-only", action="store_true",
                   help="只预下载/校验全部模型到 models/ 后退出")
    args = p.parse_args()
    if args.language:
        args.language = args.language.strip().lower()
        args.lang_primary = args.language.split("-")[0]
    else:
        args.lang_primary = None
    return args


# ---------------------------------------------------------------- 基础工具

def collect_audio_files() -> list[Path]:
    if not INPUT_DIR.is_dir():
        return []
    return [p for p in sorted(INPUT_DIR.iterdir())
            if p.is_file() and p.suffix.lower() in SUPPORTED_EXTS]


def output_is_fresh(audio: Path) -> bool:
    """断点续传：输出 md 存在且 mtime 晚于音频 mtime → 视为最新，跳过。"""
    out = OUTPUT_DIR / (audio.stem + ".md")
    return out.exists() and out.stat().st_mtime > audio.stat().st_mtime


def load_hotwords() -> list[str]:
    p = PROJECT_DIR / "hotwords.txt"
    if not p.exists():
        return []
    # 「#」开头的行为注释/分组标题，不参与热词
    return [ln.strip() for ln in p.read_text(encoding="utf-8").splitlines()
            if ln.strip() and not ln.strip().startswith("#")]


def decode_audio(path: Path, offset: float = 0.0, duration: float | None = None) -> np.ndarray:
    """解码为 16kHz float32 单声道。主路径 PyAV（faster-whisper 自带），异常回退 ffmpeg。"""
    from faster_whisper.audio import decode_audio as fw_decode
    try:
        # 该实现签名只有 (input_file, sampling_rate, split_stereo)，默认即输出单声道 float32
        audio = fw_decode(str(path), sampling_rate=SAMPLE_RATE)
    except Exception as exc:
        log(f"    [解码] PyAV 失败（{type(exc).__name__}: {short_err(exc)}），回退 ffmpeg 转码 …")
        audio = decode_audio_ffmpeg(path)
    if offset > 0 or duration is not None:
        i0 = max(0, int(offset * SAMPLE_RATE))
        i1 = len(audio) if duration is None else max(i0, min(len(audio), int((offset + duration) * SAMPLE_RATE)))
        audio = audio[i0:i1]
    return audio


def decode_audio_ffmpeg(path: Path) -> np.ndarray:
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("PyAV 解码失败，且系统中没有 ffmpeg 可作兜底")
    fd, name = tempfile.mkstemp(suffix=".wav", prefix="a2t_")
    os.close(fd)
    tmp = Path(name)
    try:
        # 文件名可能含空格/中文：一律参数列表形式，shell=False
        proc = subprocess.run(
            [ffmpeg, "-y", "-i", str(path), "-ac", "1", "-ar", str(SAMPLE_RATE), str(tmp)],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )
        if proc.returncode != 0:
            raise RuntimeError(f"ffmpeg 兜底转码也失败（返回码 {proc.returncode}），文件可能损坏或为空")
        with wave.open(str(tmp), "rb") as w:
            if w.getframerate() != SAMPLE_RATE or w.getnchannels() != 1:
                raise RuntimeError("ffmpeg 输出 wav 参数异常")
            data = w.readframes(w.getnframes())
        return np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
    finally:
        tmp.unlink(missing_ok=True)


def write_temp_wav(audio: np.ndarray) -> Path:
    """float32 [-1,1] 单声道 16k → 临时 16bit PCM wav（FunASR ndarray 兜底用）。"""
    fd, name = tempfile.mkstemp(suffix=".wav", prefix="a2t_")
    os.close(fd)
    pcm = np.clip(audio, -1.0, 1.0)
    pcm = (pcm * 32767.0).astype(np.int16)
    with wave.open(name, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(SAMPLE_RATE)
        w.writeframes(pcm.tobytes())
    return Path(name)


# ---------------------------------------------------------------- 模型加载

def dir_has_files(d: Path) -> bool:
    return d.is_dir() and any(d.iterdir())


def load_whisper(model_name: str):
    from faster_whisper import WhisperModel
    if not dir_has_files(WHISPER_CACHE_DIR):
        log(f"    首次使用：本地无缓存，正在从 HuggingFace 下载 {model_name}"
            f"（体积较大，仅此一次，之后离线可用）…")
        log("    下载慢/失败可设置国内镜像后重试：export HF_ENDPOINT=https://hf-mirror.com")
    return WhisperModel(
        model_name,
        device=DEVICE,
        compute_type=COMPUTE_TYPE,
        cpu_threads=WHISPER_CPU_THREADS,
        download_root=str(WHISPER_CACHE_DIR),
    )


def make_batched_pipeline(model):
    from faster_whisper import BatchedInferencePipeline
    return BatchedInferencePipeline(model=model)


def load_paraformer():
    import torch
    torch.set_num_threads(WHISPER_CPU_THREADS)
    if not dir_has_files(FUNASR_CACHE_DIR):
        log("    首次使用：本地无缓存，正在从 ModelScope 下载 SeACo-Paraformer/VAD/标点"
            "（约 1.3GB，仅此一次，之后离线可用）…")
    from funasr import AutoModel
    return AutoModel(
        model=PARAFORMER_MODEL,
        vad_model=VAD_MODEL,
        vad_kwargs={"max_single_segment_time": 30000},
        punc_model=PUNC_MODEL,
        device=DEVICE,
        cache_dir=str(FUNASR_CACHE_DIR),
        disable_update=True,
    )


def release(*objs) -> None:
    for o in objs:
        del o
    gc.collect()


# ---------------------------------------------------------------- 语言探测与路由

def probe_slices(path: Path) -> list[np.ndarray]:
    """探测切片：开头 45s + 时长中点 45s；文件短于 90s 则只用全片（单次完整解码）。"""
    audio = decode_audio(path)
    total_s = len(audio) / SAMPLE_RATE
    if total_s <= PROBE_SLICE_SECONDS * 2:
        return [audio]
    first_len = int(PROBE_SLICE_SECONDS * SAMPLE_RATE)
    mid_start = max(0, len(audio) // 2 - first_len // 2)
    return [audio[:first_len], audio[mid_start:mid_start + first_len]]


def probe_language(whisper_model, audio_slices: list[np.ndarray]) -> tuple[str, float]:
    """对每个切片 transcribe(language=None) 读 info.language/概率，各语言概率取平均。"""
    per_lang: dict[str, list[float]] = {}
    for idx, sl in enumerate(audio_slices, 1):
        _segments, info = whisper_model.transcribe(sl, language=None, beam_size=5)
        lang = info.language or "unknown"
        prob = float(info.language_probability or 0.0)
        per_lang.setdefault(lang, []).append(prob)
        log(f"    [探测] 切片 {idx}/{len(audio_slices)}：{lang}（p={prob:.2f}）")
    avg = {lang: sum(v) / len(v) for lang, v in per_lang.items()}
    best = max(avg, key=lambda k: avg[k])
    return best, avg[best]


def forced_route(args: argparse.Namespace) -> tuple[str, str | None, str]:
    """强制模式下的路由：--engine 优先于 --language。返回 (引擎, whisper语言码, 依据)。"""
    if args.engine and args.language:
        eng = args.engine
        basis = f"强制 --engine={eng} + --language={args.language}"
    elif args.engine:
        eng = args.engine
        basis = f"强制 --engine={eng}"
    else:
        eng = "paraformer" if args.lang_primary == "zh" else "whisper"
        basis = f"强制 --language={args.language}"
    lang = args.lang_primary if (eng == "whisper" and args.language) else None
    return eng, lang, basis


def auto_route(detected: str, prob: float) -> tuple[str, str]:
    if detected.startswith("zh") and prob > ZH_PROB_THRESHOLD:
        return "paraformer", f"自动：探测到 {detected}（平均p={prob:.2f} > {ZH_PROB_THRESHOLD}）"
    return "whisper", f"自动：探测到 {detected}（平均p={prob:.2f}）"


# ---------------------------------------------------------------- 转写

def build_whisper_prompt(hotwords: list[str], language: str | None) -> str | None:
    if not hotwords:
        return None
    terms = "、".join(hotwords[:WHISPER_PROMPT_HOTWORDS])
    prompt = f"这是一堂土木工程专业课，可能涉及以下术语：{terms}"
    if language and language.startswith("zh"):
        prompt = "以下是普通话的句子。" + prompt  # 防止中文输出成繁体
    return prompt


def _accepted_kwargs(func) -> set[str]:
    try:
        params = inspect.signature(func).parameters
    except (TypeError, ValueError):
        return set()
    return {n for n, v in params.items()
            if v.kind in (v.POSITIONAL_OR_KEYWORD, v.KEYWORD_ONLY)}


def whisper_transcribe(model, pipe, audio: np.ndarray, language: str | None, initial_prompt: str | None):
    """优先 BatchedInferencePipeline（快）；不支持 initial_prompt 时按兜底方案退回普通 transcribe。"""
    want = {
        "language": language,
        "beam_size": 5,
        "vad_filter": True,
        "vad_parameters": {"min_silence_duration_ms": 500},
        "condition_on_previous_text": True,
        "initial_prompt": initial_prompt,
        "batch_size": WHISPER_BATCH_SIZE,
    }
    candidates = []
    if pipe is not None and (not initial_prompt or "initial_prompt" in _accepted_kwargs(pipe.transcribe)):
        candidates.append(pipe.transcribe)
    candidates.append(model.transcribe)
    last_err: TypeError | None = None
    for fn in candidates:
        accepted = _accepted_kwargs(fn)
        kwargs = {k: v for k, v in want.items() if v is not None and k in accepted}
        try:
            return fn(audio, **kwargs)
        except TypeError as exc:  # 该实现不认某个参数 → 换下一个候选
            last_err = exc
    raise RuntimeError(f"transcribe 调用失败：{last_err}")


def merge_whisper_segments(segments) -> str:
    """连续约 6 个 segment 合并为一个自然段，段间空行。"""
    paragraphs: list[str] = []
    buf: list[str] = []
    for i, seg in enumerate(segments, 1):
        buf.append(seg.text)
        if i % 50 == 0:
            log(f"    … 已转写 {i} 段")
        if i % SEGMENTS_PER_PARAGRAPH == 0:
            paragraphs.append(WS_RE.sub(" ", "".join(buf)).strip())
            buf = []
    if buf:
        paragraphs.append(WS_RE.sub(" ", "".join(buf)).strip())
    return "\n\n".join(p for p in paragraphs if p)


def paraformer_transcribe(am, audio: np.ndarray, hotword_str: str) -> str:
    gen_kwargs = {"fs": SAMPLE_RATE, "batch_size_s": 300}
    if hotword_str:
        gen_kwargs["hotword"] = hotword_str
    try:
        res = am.generate(input=audio, **gen_kwargs)
    except Exception as exc:
        log(f"    [Paraformer] ndarray 输入不被接受（{type(exc).__name__}: {short_err(exc)}），"
            f"按兜底方案改传临时 wav 路径重试")
        wav = write_temp_wav(audio)
        try:
            res = am.generate(input=str(wav), **gen_kwargs)
        finally:
            wav.unlink(missing_ok=True)
    if not res or not isinstance(res[0], dict) or "text" not in res[0]:
        raise RuntimeError(f"FunASR 返回结果异常：{short_err(Exception(repr(res)))}")
    return str(res[0]["text"])


def paraformer_paragraphs(text: str) -> str:
    """整段带标点文本按句号等断句，每 ~6 句插一个空行。"""
    sents = [s for s in SENT_RE.findall(text) if s.strip()]
    if len(sents) <= SENTENCES_PER_PARAGRAPH:
        return text.strip()
    return "\n\n".join(
        "".join(sents[i:i + SENTENCES_PER_PARAGRAPH]).strip()
        for i in range(0, len(sents), SENTENCES_PER_PARAGRAPH)
    )


def write_output(audio: Path, body: str) -> Path:
    out = OUTPUT_DIR / (audio.stem + ".md")
    out.write_text(f"# {audio.stem}\n\n{body}\n", encoding="utf-8")
    return out


# ---------------------------------------------------------------- 子命令

def cmd_download_only(args: argparse.Namespace) -> int:
    log(f"==> 预下载/校验 Whisper 模型：{args.model}（来自 HuggingFace）…")
    m = load_whisper(args.model)
    log(f"    缓存目录：{WHISPER_CACHE_DIR}")
    del m
    gc.collect()
    log("==> 预下载/校验 FunASR 模型：SeACo-Paraformer + VAD + 标点（来自 ModelScope）…")
    am = load_paraformer()
    log(f"    缓存目录：{FUNASR_CACHE_DIR}")
    del am
    gc.collect()
    log("全部模型就绪，之后可离线转写。")
    return 0


# ---------------------------------------------------------------- 主流程

def main() -> int:
    args = parse_args()
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    INPUT_DIR.mkdir(parents=True, exist_ok=True)

    if args.download_only:
        return cmd_download_only(args)

    files = collect_audio_files()
    if not files:
        log("input_audio/ 中没有音频文件"
            f"（支持：{' '.join(sorted(SUPPORTED_EXTS))}）。把录音放进去后再运行。")
        return 0
    log(f"在 input_audio/ 找到 {len(files)} 个音频文件。")

    todo: list[Path] = []
    skipped = 0
    for f in files:
        if output_is_fresh(f):
            skipped += 1
            log(f"  跳过（输出已存在且比音频新）：{f.name}")
        else:
            todo.append(f)
    if not todo:
        log(f"==== 汇总 ====")
        log(f"成功 0 / 跳过 {skipped} / 失败 0（全部文件均为最新）")
        return 0

    hotwords = load_hotwords()
    log(f"热词：{len(hotwords)} 个（hotwords.txt）")

    forced = bool(args.engine or args.language)
    whisper_model = None
    whisper_items: list[tuple[Path, str | None]] = []   # (音频, whisper 语言码或 None=自动)
    para_items: list[Path] = []
    ok_n = fail_n = 0
    failures: list[tuple[str, str]] = []

    # ---------- 阶段一：解码 + 语言探测 + 路由 ----------
    if forced:
        for f in todo:
            eng, lang, basis = forced_route(args)
            log(f"  路由：{f.name} → {eng}（{basis}）")
            if eng == "whisper":
                whisper_items.append((f, lang))
            else:
                para_items.append(f)
    else:
        log(f"==> 加载 Whisper（{args.model}）：先做语言探测，之后复用同一实例转写 …")
        whisper_model = load_whisper(args.model)
        for f in todo:
            log(f"  探测：{f.name}")
            try:
                slices = probe_slices(f)
                detected, prob = probe_language(whisper_model, slices)
                del slices
            except Exception as exc:
                fail_n += 1
                failures.append((f.name, f"解码/探测失败：{short_err(exc, 200)}"))
                log(f"  路由：{f.name} → 失败（解码/探测失败：{short_err(exc)}）")
                continue
            eng, basis = auto_route(detected, prob)
            log(f"  路由：{f.name} → {eng}（{basis}）")
            if eng == "whisper":
                whisper_items.append((f, detected))
            else:
                para_items.append(f)

    if args.dry_run:
        log("==== dry-run 结束（未转写）====")
        log(f"Whisper 组 {len(whisper_items)} 个 / Paraformer 组 {len(para_items)} 个 / 跳过 {skipped} 个")
        return 0

    # ---------- 阶段二：Whisper 组（探测时模型已加载，直接复用） ----------
    pipe = None
    if whisper_items:
        if whisper_model is None:
            log(f"==> 加载 Whisper（{args.model}）…")
            whisper_model = load_whisper(args.model)
        try:
            pipe = make_batched_pipeline(whisper_model)
        except Exception:
            pipe = None  # 旧版无 BatchedInferencePipeline → 直接用普通 transcribe
        for f, lang in whisper_items:
            t0 = time.monotonic()
            log(f"  [Whisper] 开始：{f.name}（language={lang or '自动'}）")
            try:
                audio = decode_audio(f)
                prompt = build_whisper_prompt(hotwords, lang)
                segments, _info = whisper_transcribe(whisper_model, pipe, audio, lang, prompt)
                body = merge_whisper_segments(segments)
                if not body.strip():
                    raise RuntimeError("转写结果为空（可能是纯静音或音频异常）")
                out = write_output(f, body)
                log(f"  [Whisper] 完成：{f.name} → {out.name}（耗时 {time.monotonic() - t0:.1f}s）")
                ok_n += 1
            except Exception as exc:
                fail_n += 1
                failures.append((f.name, f"{type(exc).__name__}: {short_err(exc, 300)}"))
                log(f"  [Whisper] 失败：{f.name}（{short_err(exc, 200)}）")

    # 两引擎不同时常驻：先彻底释放 Whisper，再加载 Paraformer
    if whisper_model is not None:
        del whisper_model, pipe
        whisper_model = None
        gc.collect()
        log("==> 已释放 Whisper 模型")

    # ---------- 阶段三：Paraformer 组 ----------
    if para_items:
        log("==> 加载 FunASR（SeACo-Paraformer + VAD + 标点）…")
        am = load_paraformer()
        hotword_str = " ".join(hotwords)
        for f in para_items:
            t0 = time.monotonic()
            log(f"  [Paraformer] 开始：{f.name}")
            try:
                audio = decode_audio(f)
                text = paraformer_transcribe(am, audio, hotword_str)
                body = paraformer_paragraphs(text)
                if not body.strip():
                    raise RuntimeError("转写结果为空（可能是纯静音或音频异常）")
                out = write_output(f, body)
                log(f"  [Paraformer] 完成：{f.name} → {out.name}（耗时 {time.monotonic() - t0:.1f}s）")
                ok_n += 1
            except Exception as exc:
                fail_n += 1
                failures.append((f.name, f"{type(exc).__name__}: {short_err(exc, 300)}"))
                log(f"  [Paraformer] 失败：{f.name}（{short_err(exc, 200)}）")
        del am
        gc.collect()

    # ---------- 汇总 ----------
    log("==== 汇总 ====")
    log(f"成功 {ok_n} / 跳过 {skipped} / 失败 {fail_n}")
    for name, reason in failures:
        log(f"  ✗ {name}：{reason}")
    return 1 if fail_n else 0


if __name__ == "__main__":
    sys.exit(main())