#!/usr/bin/env python3
"""
transcribe_audio.py - MP3 音频批量转录文字脚本
功能:
  - 扫描单个 MP3 文件或遍历指定目录中的所有 MP3 文件
  - 维护 manifest.json 进度文件，支持断点续传（key 含模型维度，多模型互不覆盖）
  - 使用系统已安装的 whisper 命令行工具转录音频
  - 支持 MP3 切分（处理可能生成超大 txt 的音频文件）
  - 支持含空格和中文的文件名
  - 支持预下载模型 (--preload-models)
  - 支持多模型对比 (--all-models)
  - 支持自定义输出后缀 (--output-suffix)

支持的模型: tiny / base / small / medium
推荐语言:  zh (中文) / en (英文)，也接受其他 whisper 支持的语言代码

用法:
  python transcribe_audio.py <MP3文件或目录> [选项]

示例:
  # 预下载所有支持的模型
  python transcribe_audio.py --preload-models

  # 使用 medium 模型转录中文
  python transcribe_audio.py /path/to/audio/ --model medium --language zh

  # 一次跑 4 个模型对比精度（生成 file.tiny.txt / file.base.txt / ...）
  python transcribe_audio.py /path/to/audio/ --all-models --language zh

  # 自定义输出后缀
  python transcribe_audio.py /path/to/audio.mp3 --model base --output-suffix _v1

  # 超1800秒切分
  python transcribe_audio.py /path/to/audio/ --split-size 1800

  # 重新处理失败项
  python transcribe_audio.py /path/to/audio/ --retry-failed

  # 查看进度状态
  python transcribe_audio.py /path/to/audio/ --status
"""

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

# ─── 常量 ──────────────────────────────────────────────────────────────────────

MANIFEST_NAME = "manifest.json"
AUDIO_EXTS = {".mp3", ".m4a", ".wav", ".flac", ".ogg", ".aac"}

SUPPORTED_MODELS = ["tiny", "base", "small", "medium"]
DEFAULT_WHISPER_MODEL = "base"
DEFAULT_SPLIT_SEGMENT_SECS = 1800   # 每片30分钟


# ─── 日志 ──────────────────────────────────────────────────────────────────────

def log(msg: str, level: str = "INFO"):
    ts = time.strftime("%H:%M:%S")
    prefix = {"INFO": "  ", "OK": "✅", "ERR": "❌", "WARN": "⚠️ ", "STEP": "▶ "}.get(level, "  ")
    print(f"[{ts}] {prefix} {msg}", flush=True)


# ─── Manifest 操作 (key 包含 model，支持同一音频多模型共存) ────────────────────

def load_manifest(manifest_path: Path) -> dict:
    if manifest_path.exists():
        with open(manifest_path, encoding="utf-8") as f:
            return json.load(f)
    return {"version": 1, "files": {}}


def save_manifest(manifest_path: Path, manifest: dict):
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    with open(manifest_path, "w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)


def transcribe_key(model: str, audio_path: str, output_suffix: str = "") -> str:
    """
    manifest key 由 model + output_suffix + 音频路径共同决定，
    因此改变 --output-suffix 会形成新的 key（不会误跳过）。
    """
    if output_suffix:
        return f"transcribe:{model}:{output_suffix}:{audio_path}"
    return f"transcribe:{model}:{audio_path}"


def get_transcribe_entry(manifest: dict, audio_path: str, model: str,
                         output_suffix: str = "") -> dict:
    key = transcribe_key(model, audio_path, output_suffix)
    if key not in manifest["files"]:
        manifest["files"][key] = {
            "transcribed": False,
            "txt_path": None,
            "split_segments": [],
            "error": None,
            "model": model,
            "language": None,
            "output_suffix": output_suffix or None,
            "updated_at": None,
        }
    return manifest["files"][key]


def mark_transcribed(manifest: dict, manifest_path: Path, audio_path: str,
                     txt_path: str, model: str, language: str = None,
                     output_suffix: str = "", segments: list = None):
    entry = get_transcribe_entry(manifest, audio_path, model, output_suffix)
    entry["transcribed"] = True
    entry["txt_path"] = txt_path
    entry["split_segments"] = segments or []
    entry["model"] = model
    entry["language"] = language
    entry["output_suffix"] = output_suffix or None
    entry["error"] = None
    entry["updated_at"] = time.strftime("%Y-%m-%dT%H:%M:%S")
    save_manifest(manifest_path, manifest)


def mark_transcribe_failed(manifest: dict, manifest_path: Path,
                           audio_path: str, model: str, error: str,
                           output_suffix: str = ""):
    entry = get_transcribe_entry(manifest, audio_path, model, output_suffix)
    entry["transcribed"] = False
    entry["error"] = error
    entry["updated_at"] = time.strftime("%Y-%m-%dT%H:%M:%S")
    save_manifest(manifest_path, manifest)


# ─── 模型预下载 ────────────────────────────────────────────────────────────────

def preload_models(models: list) -> bool:
    """
    调用 whisper Python API 预下载模型到 ~/.cache/whisper/。
    比调用 whisper CLI 更快（无需额外音频输入）。
    """
    try:
        import whisper as whisper_lib
    except ImportError:
        log("openai-whisper Python 库未安装，无法预下载", "ERR")
        log("请安装：pip install openai-whisper", "ERR")
        return False

    log(f"开始预下载 {len(models)} 个模型：{', '.join(models)}", "STEP")
    for m in models:
        log(f"  下载 {m}...")
        t0 = time.time()
        try:
            whisper_lib.load_model(m)
            log(f"  {m} 就绪 ({time.time() - t0:.1f}s)", "OK")
        except Exception as e:
            log(f"  {m} 下载失败: {e}", "ERR")
            return False
    log(f"所有模型下载完成 ✅", "OK")
    return True


# ─── Whisper / FFmpeg 工具函数 ─────────────────────────────────────────────────

def check_whisper() -> bool:
    try:
        subprocess.run(["whisper", "--help"], capture_output=True, check=False)
        return True
    except FileNotFoundError:
        log("未找到 whisper 命令，请确认已安装：pip install openai-whisper", "ERR")
        return False


def check_ffmpeg() -> bool:
    try:
        subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False


def get_audio_duration(audio_path: str) -> float:
    try:
        r = subprocess.run(
            ["ffprobe", "-v", "error", "-show_entries", "format=duration",
             "-of", "default=noprint_wrappers=1:nokey=1", audio_path],
            capture_output=True, text=True
        )
        return float(r.stdout.strip())
    except Exception:
        return 0.0


def split_audio(audio_path: str, output_dir: str, segment_secs: int) -> list:
    """
    将 MP3 切分为多个片段，返回片段路径列表。
    使用 ffmpeg segment 模式，避免重编码（codec copy）。
    """
    stem = Path(audio_path).stem
    ext = Path(audio_path).suffix
    out_pattern = os.path.join(output_dir, f"{stem}_seg%03d{ext}")

    cmd = [
        "ffmpeg", "-y", "-i", audio_path,
        "-f", "segment",
        "-segment_time", str(segment_secs),
        "-c", "copy",
        "-reset_timestamps", "1",
        out_pattern
    ]
    r = subprocess.run(cmd, capture_output=True)
    if r.returncode != 0:
        return []

    return sorted(str(p) for p in Path(output_dir).glob(f"{stem}_seg*{ext}"))


def run_whisper(audio_path: str, output_dir: str, model: str,
                language: str = None) -> tuple:
    """
    调用 whisper 命令转录音频，返回 (success: bool, txt_path_or_err: str)。
    whisper 会自动在 output_dir 生成 <stem>.txt 等文件。
    注意: 多次调用同 output_dir 且相同 stem 时会覆盖，需由调用方即时重命名。
    """
    cmd = [
        "whisper", audio_path,
        "--model", model,
        "--output_dir", output_dir,
        "--output_format", "txt",
        "--verbose", "False",
    ]
    if language:
        cmd += ["--language", language]

    r = subprocess.run(cmd, capture_output=True, text=True)

    stem = Path(audio_path).stem
    txt_path = os.path.join(output_dir, f"{stem}.txt")

    if r.returncode == 0 and os.path.exists(txt_path):
        return True, txt_path

    err = r.stderr[-500:] if r.stderr else "unknown error"
    return False, err


def merge_txt_files(txt_files: list, merged_path: str) -> bool:
    """将多个转录片段 txt 按顺序合并为一个完整文件"""
    try:
        with open(merged_path, "w", encoding="utf-8") as fout:
            for i, tf in enumerate(txt_files):
                if i > 0:
                    fout.write("\n\n")
                with open(tf, encoding="utf-8") as fin:
                    fout.write(fin.read().strip())
        return True
    except Exception:
        return False


# ─── 文件扫描 ──────────────────────────────────────────────────────────────────

def scan_audios(source: str) -> list:
    p = Path(source)
    if p.is_file():
        if p.suffix.lower() in AUDIO_EXTS:
            return [str(p.resolve())]
        else:
            log(f"不支持的音频格式: {p.suffix}", "ERR")
            return []
    elif p.is_dir():
        return sorted(str(f.resolve()) for f in p.rglob("*")
                      if f.suffix.lower() in AUDIO_EXTS)
    else:
        log(f"路径不存在: {source}", "ERR")
        return []


def txt_exists(txt_path: str) -> bool:
    return bool(txt_path) and os.path.exists(txt_path) and os.path.getsize(txt_path) > 0


# ─── 输出文件命名 ──────────────────────────────────────────────────────────────

def compute_target_txt(out_dir: Path, stem: str, model: str,
                       output_suffix: str, tag_model: bool) -> str:
    """
    生成最终 txt 路径。命名规则：
    - tag_model=True (多模型模式):     <stem>.<model>.txt
    - output_suffix 非空:              <stem><output_suffix>.txt
    - 都无:                            <stem>.txt (默认)
    """
    if tag_model:
        return str(out_dir / f"{stem}.{model}.txt")
    if output_suffix:
        return str(out_dir / f"{stem}{output_suffix}.txt")
    return str(out_dir / f"{stem}.txt")


# ─── 状态展示 ──────────────────────────────────────────────────────────────────

def show_status(manifest: dict, all_audios: list):
    print("\n" + "=" * 70)
    print("📊 音频转录进度状态（按 模型 x 音频 展示）")
    print("=" * 70)

    entries = {k: v for k, v in manifest["files"].items() if k.startswith("transcribe:")}
    done = sum(1 for v in entries.values() if v.get("transcribed"))
    failed = sum(1 for v in entries.values() if v.get("error"))
    print(f"  Manifest 记录: {len(entries)}  已完成: {done}  失败: {failed}")
    print("=" * 70)

    for a in all_audios:
        print(f"\n  📄 {Path(a).name}")
        matching = [(k, v) for k, v in entries.items() if k.endswith(f":{a}")]
        if not matching:
            print(f"     ⏳ 无记录（未转录）")
            continue
        for key, entry in matching:
            m = entry.get("model", "?")
            suffix = entry.get("output_suffix") or ""
            suffix_tag = f" suffix={suffix}" if suffix else ""
            if entry.get("transcribed"):
                size = ""
                if entry.get("txt_path") and os.path.exists(entry["txt_path"]):
                    size = f"  {os.path.getsize(entry['txt_path']) // 1024}KB"
                lang = entry.get("language") or "auto"
                txt_name = Path(entry['txt_path']).name if entry.get('txt_path') else '?'
                print(f"     ✅ [{m:6s}] lang={lang}{suffix_tag}{size}  → {txt_name}")
            elif entry.get("error"):
                print(f"     ❌ [{m:6s}]{suffix_tag} {entry['error'][:60]}")
    print()


# ─── 核心：单个音频 + 单个模型 转录 ─────────────────────────────────────────────

def transcribe_one(audio: str, model: str, out_dir: Path, args,
                   manifest: dict, manifest_path: Path,
                   has_ffmpeg: bool, tag_model: bool) -> str:
    """
    返回状态字符串：'ok' / 'skip' / 'skip-failed' / 'error:<msg>'
    """
    stem = Path(audio).stem
    final_txt = compute_target_txt(out_dir, stem, model, args.output_suffix, tag_model)

    # 多模型模式用 .<model>.txt 命名，output_suffix 在这种情况下强制为空
    effective_suffix = "" if tag_model else args.output_suffix
    entry = get_transcribe_entry(manifest, audio, model, effective_suffix)
    already_done = entry.get("transcribed") and txt_exists(entry.get("txt_path"))
    is_failed = bool(entry.get("error"))

    if already_done:
        return "skip"
    if is_failed and not args.retry_failed:
        return "skip-failed"

    try:
        duration = get_audio_duration(audio) if has_ffmpeg else 0
        need_split = args.split_size and duration > 0 and duration > args.split_size

        if need_split:
            log(f"    音频时长 {duration:.0f}s > {args.split_size}s，启动切分转录")
            seg_dir = out_dir / f"{stem}_segments"
            seg_dir.mkdir(exist_ok=True)

            segments = split_audio(audio, str(seg_dir), args.split_segment)
            if not segments:
                raise RuntimeError("音频切分失败")

            # 每个模型独立子目录，避免片段 txt 互相覆盖
            seg_txt_dir = seg_dir / model
            seg_txt_dir.mkdir(exist_ok=True)
            log(f"    切分为 {len(segments)} 个片段，逐段转录中...")

            seg_txts = []
            for j, seg in enumerate(segments, 1):
                ok, result = run_whisper(seg, str(seg_txt_dir), model, args.language)
                if ok:
                    seg_txts.append(result)
                    log(f"      [{j}/{len(segments)}] {Path(seg).name} → done", "OK")
                else:
                    raise RuntimeError(f"片段 {Path(seg).name} 转录失败: {result}")

            if not merge_txt_files(seg_txts, final_txt):
                raise RuntimeError("合并转录片段失败")

            size = os.path.getsize(final_txt)
            log(f"    合并 → {Path(final_txt).name} ({size // 1024} KB)", "OK")
            mark_transcribed(manifest, manifest_path, audio, final_txt,
                             model, language=args.language,
                             output_suffix=effective_suffix, segments=seg_txts)
        else:
            log(f"    🎙 转录中（模型: {model}, 语言: {args.language or 'auto'}）...")
            ok, whisper_out = run_whisper(audio, str(out_dir), model, args.language)
            if not ok:
                raise RuntimeError(f"转录失败: {whisper_out}")

            if whisper_out != final_txt:
                os.replace(whisper_out, final_txt)

            size = os.path.getsize(final_txt)
            log(f"    输出: {Path(final_txt).name} ({size // 1024} KB)", "OK")
            mark_transcribed(manifest, manifest_path, audio, final_txt,
                             model, language=args.language,
                             output_suffix=effective_suffix)

        return "ok"

    except Exception as e:
        err_msg = str(e)
        mark_transcribe_failed(manifest, manifest_path, audio, model, err_msg,
                               output_suffix=effective_suffix)
        # 清理不完整的输出
        if os.path.exists(final_txt) and os.path.getsize(final_txt) == 0:
            os.remove(final_txt)
        return f"error: {err_msg}"


# ─── 主流程 ────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="MP3 音频批量转录工具（基于 openai-whisper）",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__
    )
    parser.add_argument("source", nargs="?", default=None,
                        help="MP3 文件路径或包含音频的目录（--preload-models 时可省略）")
    parser.add_argument("--output", "-o", default=None,
                        help="TXT 输出目录（默认与音频文件同目录）")
    parser.add_argument("--manifest", "-m", default=None,
                        help="manifest.json 路径（默认在 source 目录下）")
    parser.add_argument("--model", default=DEFAULT_WHISPER_MODEL,
                        choices=SUPPORTED_MODELS,
                        help=f"Whisper 模型（默认 {DEFAULT_WHISPER_MODEL}，可选：{'/'.join(SUPPORTED_MODELS)}）")
    parser.add_argument("--language", default=None,
                        help="音频语言代码，推荐 zh(中文) / en(英文)（默认自动检测）")
    parser.add_argument("--all-models", action="store_true",
                        help=f"对每个音频依次使用所有支持的模型转录 ({'/'.join(SUPPORTED_MODELS)})，"
                             f"输出为 <stem>.<model>.txt 便于对比精度")
    parser.add_argument("--output-suffix", default="",
                        help="输出文件名后缀，如 _v1 → <stem>_v1.txt（与 --all-models 互斥时被忽略）")
    parser.add_argument("--split-size", type=int, default=None,
                        help="切分阈值（秒）：超过此时长的音频会被切分后分别转录")
    parser.add_argument("--split-segment", type=int, default=DEFAULT_SPLIT_SEGMENT_SECS,
                        help=f"每个切分片段时长（秒，默认 {DEFAULT_SPLIT_SEGMENT_SECS}）")
    parser.add_argument("--retry-failed", action="store_true",
                        help="重新处理上次失败的文件")
    parser.add_argument("--status", action="store_true",
                        help="仅查看进度状态，不执行转录")
    parser.add_argument("--preload-models", action="store_true",
                        help=f"预下载所有支持的模型（{'/'.join(SUPPORTED_MODELS)}）到本地缓存后退出")
    args = parser.parse_args()

    # ── 模式1: 仅预下载模型 ──
    if args.preload_models:
        if not preload_models(SUPPORTED_MODELS):
            sys.exit(1)
        return

    # 从这里开始 source 是必需的
    if not args.source:
        parser.error("需要指定 source 参数（音频文件或目录），除非使用 --preload-models")

    if not check_whisper():
        sys.exit(1)

    has_ffmpeg = check_ffmpeg()
    if args.split_size and not has_ffmpeg:
        log("--split-size 需要 ffmpeg，但未找到 ffmpeg 命令", "ERR")
        sys.exit(1)

    source_path = Path(args.source).resolve()
    base_dir = source_path if source_path.is_dir() else source_path.parent
    manifest_path = Path(args.manifest) if args.manifest else base_dir / MANIFEST_NAME

    manifest = load_manifest(manifest_path)

    all_audios = scan_audios(args.source)
    if not all_audios:
        log("未找到任何音频文件", "WARN")
        sys.exit(0)

    # 决定使用哪些模型
    if args.all_models:
        models_to_run = SUPPORTED_MODELS
        if args.output_suffix:
            log(f"--all-models 已启用，--output-suffix 将被忽略（自动使用 .<model>.txt 命名）", "WARN")
    else:
        models_to_run = [args.model]

    tag_model = args.all_models
    log(f"扫描到 {len(all_audios)} 个音频文件，将使用 {len(models_to_run)} 个模型：{', '.join(models_to_run)}")

    # ── 模式2: 仅查看状态 ──
    if args.status:
        show_status(manifest, all_audios)
        return

    # ── 模式3: 执行转录 ──
    total_jobs = len(all_audios) * len(models_to_run)
    stats = {"ok": 0, "skip": 0, "skip-failed": 0, "error": 0}
    failed_details = []

    log(f"共 {total_jobs} 个 (音频 × 模型) 任务待处理")

    job_idx = 0
    for audio in all_audios:
        audio_name = Path(audio).name
        audio_dir = Path(audio).parent
        out_dir = Path(args.output).resolve() if args.output else audio_dir
        out_dir.mkdir(parents=True, exist_ok=True)

        log(f"\n📄 {audio_name}", "STEP")

        for model in models_to_run:
            job_idx += 1
            log(f"  [{job_idx}/{total_jobs}] 模型: {model}")

            t0 = time.time()
            result = transcribe_one(audio, model, out_dir, args, manifest,
                                    manifest_path, has_ffmpeg, tag_model)
            elapsed = time.time() - t0

            if result == "ok":
                stats["ok"] += 1
                log(f"    ⏱  耗时 {elapsed:.1f}s", "OK")
            elif result == "skip":
                stats["skip"] += 1
                log(f"    ⏭  已完成，跳过", "INFO")
            elif result == "skip-failed":
                stats["skip-failed"] += 1
                log(f"    ⚠  上次失败，加 --retry-failed 可重试", "WARN")
            else:
                stats["error"] += 1
                failed_details.append((audio_name, model, result))
                log(f"    ❌ {result}", "ERR")

    # ── 汇总 ──
    print(f"\n{'='*70}")
    log(f"🏁 完成: 成功 {stats['ok']}  已跳过 {stats['skip']}  "
        f"失败跳过 {stats['skip-failed']}  失败 {stats['error']}", "INFO")

    if failed_details:
        log("失败详情:", "ERR")
        for audio_name, model, err in failed_details:
            print(f"    - {audio_name} [模型:{model}] : {err[:80]}")
        print("  运行时加 --retry-failed 可重新处理失败项")


if __name__ == "__main__":
    main()