#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
影语学习台 · 本地字幕助手

用途：
  在用户电脑本地从视频生成 .srt 字幕。电影文件不会上传到服务器。

准备：
  1. 安装 ffmpeg
     macOS: brew install ffmpeg
     Windows: 安装 ffmpeg 并加入 PATH

  2. 安装 Python 依赖
     python3 -m pip install faster-whisper

用法：
  python3 local-subtitle-helper.py "/path/to/movie.mp4"
  python3 local-subtitle-helper.py "/path/to/movie.mp4" --language en --model small
"""

import argparse
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path


def fail(message):
    print(f"\n错误：{message}", file=sys.stderr)
    sys.exit(1)


def check_command(name):
    if shutil.which(name) is None:
        fail(f"未找到 {name}。请先安装 {name} 并确认命令行可用。")


def format_time(seconds):
    if seconds < 0:
        seconds = 0
    millis = int(round((seconds - int(seconds)) * 1000))
    total = int(seconds)
    h = total // 3600
    m = (total % 3600) // 60
    s = total % 60
    return f"{h:02d}:{m:02d}:{s:02d},{millis:03d}"


def extract_audio(video_path, audio_path):
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(video_path),
        "-vn",
        "-ac",
        "1",
        "-ar",
        "16000",
        "-f",
        "wav",
        str(audio_path),
    ]
    subprocess.run(cmd, check=True)


def write_srt(segments, output_path):
    with open(output_path, "w", encoding="utf-8") as fp:
        for index, seg in enumerate(segments, 1):
            text = " ".join((seg.text or "").strip().split())
            if not text:
                continue
            fp.write(f"{index}\n")
            fp.write(f"{format_time(seg.start)} --> {format_time(seg.end)}\n")
            fp.write(f"{text}\n\n")


def main():
    parser = argparse.ArgumentParser(description="本地生成电影字幕，不上传电影文件。")
    parser.add_argument("video", help="本地视频文件路径")
    parser.add_argument("--output", help="输出 .srt 路径，默认与视频同名")
    parser.add_argument("--language", default="auto", help="语言，如 en/zh/auto，默认 auto")
    parser.add_argument("--model", default="small", help="模型：tiny/base/small/medium/large-v3，默认 small")
    parser.add_argument("--device", default="auto", help="设备：auto/cpu/cuda，默认 auto")
    args = parser.parse_args()

    video_path = Path(args.video).expanduser().resolve()
    if not video_path.exists():
        fail(f"视频不存在：{video_path}")
    check_command("ffmpeg")

    try:
        from faster_whisper import WhisperModel
    except ImportError:
        fail("未安装 faster-whisper。请执行：python3 -m pip install faster-whisper")

    output_path = Path(args.output).expanduser().resolve() if args.output else video_path.with_suffix(".srt")

    print("影语学习台 · 本地字幕助手")
    print(f"视频：{video_path}")
    print(f"输出：{output_path}")
    print(f"模型：{args.model}")
    print("\n正在提取音频...")

    with tempfile.TemporaryDirectory() as tmp:
        audio_path = Path(tmp) / "audio.wav"
        extract_audio(video_path, audio_path)

        print("正在识别字幕，第一次运行会下载模型，可能需要较久...")
        model = WhisperModel(args.model, device=args.device, compute_type="auto")
        language = None if args.language == "auto" else args.language
        segments, info = model.transcribe(str(audio_path), language=language, vad_filter=True, beam_size=5)

        print(f"识别语言：{getattr(info, 'language', 'unknown')}")
        write_srt(list(segments), output_path)

    print("\n完成。")
    print(f"字幕文件：{output_path}")
    print("回到 https://english.autoflow.asia，在“高级工具”里手动选择这份字幕即可开始学习。")


if __name__ == "__main__":
    main()