ForcePilot/backend/package/yuxi/knowledge/parser/ffmpeg.py

"""ffmpeg/ffprobe 执行基础设施，提供音视频处理的底层能力。"""

import json
import subprocess
import threading
from pathlib import Path

from yuxi.utils import logger

# ── 限制常量 ──────────────────────────────────────────────
FFMPEG_MAX_BUFFER_BYTES = 10 * 1024 * 1024  # 10MB
FFPROBE_TIMEOUT_MS = 10_000  # 10s
FFMPEG_TIMEOUT_MS = 45_000  # 45s
FFMPEG_MAX_AUDIO_DURATION_SECS = 1200  # 20 minutes
MAX_AUDIO_BYTES = 16 * 1024 * 1024  # 16MB
MAX_VIDEO_BYTES = 16 * 1024 * 1024  # 16MB
FFMPEG_MAX_CONCURRENT = 4

# ── 并发信号量 ────────────────────────────────────────────
_semaphore = threading.Semaphore(FFMPEG_MAX_CONCURRENT)


# ── 异常 ──────────────────────────────────────────────────
class FFmpegError(Exception):
    """ffmpeg/ffprobe 执行错误。"""

    def __init__(self, message: str, returncode: int = -1):
        super().__init__(message)
        self.message = message
        self.returncode = returncode


# ── 核心执行函数 ──────────────────────────────────────────
def run_ffmpeg(args: list[str], timeout_ms: int = FFMPEG_TIMEOUT_MS, *, return_stderr: bool = False) -> str:
    """同步执行 ffmpeg 命令，默认返回 stdout 内容；return_stderr=True 时返回 stderr。"""
    cmd = ["ffmpeg", *args]
    timeout_sec = timeout_ms / 1000

    with _semaphore:
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                timeout=timeout_sec,
            )
        except subprocess.TimeoutExpired:
            raise FFmpegError(
                f"ffmpeg 执行超时 ({timeout_ms}ms): {' '.join(cmd)}",
                returncode=-1,
            )
        except OSError as e:
            raise FFmpegError(
                f"ffmpeg 执行失败（请确认 ffmpeg 已安装）: {e}",
                returncode=-1,
            )

    output = result.stderr if return_stderr else result.stdout

    if len(output) > FFMPEG_MAX_BUFFER_BYTES:
        raise FFmpegError(
            f"ffmpeg 输出超过缓冲区限制 ({FFMPEG_MAX_BUFFER_BYTES} bytes)",
            returncode=result.returncode,
        )

    if result.returncode != 0:
        stderr = result.stderr.decode(errors="replace")[:512]
        raise FFmpegError(
            f"ffmpeg 返回非零退出码 {result.returncode}: {stderr}",
            returncode=result.returncode,
        )

    return output.decode(errors="replace")


def run_ffprobe(args: list[str], timeout_ms: int = FFPROBE_TIMEOUT_MS) -> str:
    """同步执行 ffprobe 命令，返回 stdout 内容。"""
    cmd = ["ffprobe", *args]
    timeout_sec = timeout_ms / 1000

    with _semaphore:
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                timeout=timeout_sec,
            )
        except subprocess.TimeoutExpired:
            raise FFmpegError(
                f"ffprobe 执行超时 ({timeout_ms}ms): {' '.join(cmd)}",
                returncode=-1,
            )
        except OSError as e:
            raise FFmpegError(
                f"ffprobe 执行失败（请确认 ffprobe 已安装）: {e}",
                returncode=-1,
            )

    if len(result.stdout) > FFMPEG_MAX_BUFFER_BYTES:
        raise FFmpegError(
            f"ffprobe 输出超过缓冲区限制 ({FFMPEG_MAX_BUFFER_BYTES} bytes)",
            returncode=result.returncode,
        )

    if result.returncode != 0:
        stderr = result.stderr.decode(errors="replace")[:512]
        raise FFmpegError(
            f"ffprobe 返回非零退出码 {result.returncode}: {stderr}",
            returncode=result.returncode,
        )

    return result.stdout.decode(errors="replace")


# ── 业务函数 ──────────────────────────────────────────────
def probe_audio_duration(file_path: str) -> float | None:
    """使用 ffprobe 获取音频时长（秒），失败返回 None。"""
    try:
        output = run_ffprobe(
            ["-v", "quiet", "-print_format", "json", "-show_format", file_path],
        )
        info = json.loads(output)
        return float(info["format"]["duration"])
    except (FFmpegError, KeyError, ValueError, json.JSONDecodeError) as e:
        logger.warning("获取音频时长失败: %s — %s", file_path, e)
        return None


def convert_to_wav(input_path: str, output_path: str) -> None:
    """将音频转换为 WAV 格式（16kHz, 单声道, PCM s16le）。"""
    run_ffmpeg([
        "-i", input_path,
        "-vn",
        "-acodec", "pcm_s16le",
        "-ar", "16000",
        "-ac", "1",
        output_path,
    ])


def extract_audio_from_video(video_path: str, output_path: str) -> None:
    """从视频中提取音轨，输出为 WAV 格式（16kHz, 单声道, PCM s16le）。"""
    run_ffmpeg([
        "-i", video_path,
        "-vn",
        "-acodec", "pcm_s16le",
        "-ar", "16000",
        "-ac", "1",
        output_path,
    ])


def extract_keyframes(
    video_path: str,
    output_dir: str,
    scene_threshold: float = 0.3,
) -> list[tuple[str, float]]:
    """通过场景变化检测提取关键帧，返回 [(路径, 时间戳秒), ...] 按时间排序。

    分两步执行：
    1. 用 ffmpeg select+showinfo 检测场景切换帧的 PTS
    2. 用 -ss 精确跳转到每个时间点提取帧
    """
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    # 第一步：检测场景切换帧的 PTS 时间戳
    timestamps = _detect_scene_change_pts(video_path, scene_threshold)
    if not timestamps:
        return []

    # 第二步：在每个时间点提取帧
    frames: list[tuple[str, float]] = []
    for i, ts in enumerate(timestamps, start=1):
        output_path = str(Path(output_dir) / f"frame_{i:04d}.png")
        try:
            run_ffmpeg([
                "-ss", f"{ts:.6f}",
                "-i", video_path,
                "-frames:v", "1",
                output_path,
            ])
            frames.append((output_path, ts))
        except FFmpegError:
            logger.debug("关键帧提取跳过: ts=%.2f", ts)

    return frames


def _detect_scene_change_pts(video_path: str, scene_threshold: float) -> list[float]:
    """使用 ffmpeg select+showinfo 检测场景切换帧的 PTS 时间戳。"""
    import re

    stderr_text = run_ffmpeg([
        "-i", video_path,
        "-vf", f"select=gt(scene\\,{scene_threshold}),showinfo",
        "-vsync", "vfr",
        "-f", "null", "-",
    ], return_stderr=True)

    timestamps: list[float] = []
    for match in re.finditer(r"pts_time:(\d+\.?\d*)", stderr_text):
        try:
            timestamps.append(float(match.group(1)))
        except ValueError:
            pass

    return timestamps