xiaoxia-saas/apps/worker/video_processing/processor.py

"""
视频处理核心类
"""

import logging
import os
import tempfile
from dataclasses import dataclass
from typing import List

import ffmpeg

logger = logging.getLogger(__name__)


@dataclass
class VideoResult:
    """视频生成结果"""

    output_path: str
    thumbnail_path: str
    duration: float
    width: int
    height: int
    fps: float
    file_size: int


class VideoProcessor:
    """视频处理器"""

    def __init__(self, temp_dir: str = None):
        """
        初始化视频处理器

        Args:
            temp_dir: 临时文件目录，默认使用系统临时目录
        """
        self.temp_dir = temp_dir or tempfile.gettempdir()

    def concatenate_videos(
        self,
        input_paths: List[str],
        output_path: str,
        resolution: tuple[int, int] = (1920, 1080),
        fps: int = 25,
    ) -> VideoResult:
        """
        拼接多个视频

        Args:
            input_paths: 输入视频路径列表
            output_path: 输出视频路径
            resolution: 输出分辨率 (width, height)
            fps: 输出帧率

        Returns:
            VideoResult: 生成结果
        """
        if not input_paths:
            raise ValueError("input_paths cannot be empty")

        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # P2-6 Fix: 使用 try-finally 确保临时文件清理
        concat_file = None
        try:
            # 使用 NamedTemporaryFile 确保临时文件正确清理
            concat_file = tempfile.NamedTemporaryFile(
                mode="w",
                suffix=".txt",
                prefix="ffmpeg_concat_",
                dir=self.temp_dir,
                delete=True,
            )
            for path in input_paths:
                # FFmpeg concat demuxer 格式
                concat_file.write(f"file '{os.path.abspath(path)}'\n")
            concat_file.flush()
            concat_file_path = concat_file.name

            # 使用 FFmpeg 拼接视频
            width, height = resolution
            (
                ffmpeg.input(concat_file_path, format="concat", safe=0)
                .output(
                    output_path,
                    vcodec="libx264",
                    acodec="aac",
                    s=f"{width}x{height}",
                    r=fps,
                    preset="medium",
                    crf=23,
                )
                .overwrite_output()
                .run(capture_stdout=True, capture_stderr=True)
            )

            # 获取视频元数据
            probe = ffmpeg.probe(output_path)
            video_info = next(s for s in probe["streams"] if s["codec_type"] == "video")

            duration = float(probe["format"]["duration"])
            width = int(video_info["width"])
            height = int(video_info["height"])

            # 计算帧率
            fps_str = video_info.get("r_frame_rate", "25/1")
            fps_parts = fps_str.split("/")
            fps_value = float(fps_parts[0]) / float(fps_parts[1]) if len(fps_parts) == 2 else float(fps_parts[0])

            file_size = os.path.getsize(output_path)

            # 生成缩略图
            thumbnail_path = self.generate_thumbnail(output_path)

            return VideoResult(
                output_path=output_path,
                thumbnail_path=thumbnail_path,
                duration=duration,
                width=width,
                height=height,
                fps=fps_value,
                file_size=file_size,
            )

        except ffmpeg.Error as e:
            stderr = e.stderr.decode() if e.stderr else ""
            raise RuntimeError(f"FFmpeg error: {stderr}") from e
        finally:
            # P2-6 Fix: 确保临时文件在所有情况下都被清理
            if concat_file is not None:
                try:
                    concat_file.close()
                except OSError as close_err:
                    logger.warning("临时文件关闭失败: %s", close_err)

    def generate_thumbnail(
        self,
        video_path: str,
        timestamp: float = 1.0,
        output_path: str = None,
    ) -> str:
        """
        生成视频缩略图

        Args:
            video_path: 视频文件路径
            timestamp: 截图时间点（秒）
            output_path: 输出路径，默认为视频路径 + .jpg

        Returns:
            缩略图路径
        """
        if output_path is None:
            output_path = f"{os.path.splitext(video_path)[0]}_thumb.jpg"

        try:
            (
                ffmpeg.input(video_path, ss=timestamp)
                .output(output_path, vframes=1, format="image2", vcodec="mjpeg", pix_fmt="yuvj420p")
                .overwrite_output()
                .run(capture_stdout=True, capture_stderr=True)
            )

            return output_path

        except ffmpeg.Error as e:
            stderr = e.stderr.decode() if e.stderr else ""
            raise RuntimeError(f"FFmpeg thumbnail error: {stderr}") from e

    def get_video_info(self, video_path: str) -> dict:
        """
        获取视频信息

        Args:
            video_path: 视频文件路径

        Returns:
            视频元数据字典
        """
        try:
            probe = ffmpeg.probe(video_path)
            video_info = next(s for s in probe["streams"] if s["codec_type"] == "video")

            duration = float(probe["format"]["duration"])
            width = int(video_info["width"])
            height = int(video_info["height"])

            fps_str = video_info.get("r_frame_rate", "25/1")
            fps_parts = fps_str.split("/")
            fps_value = float(fps_parts[0]) / float(fps_parts[1]) if len(fps_parts) == 2 else float(fps_parts[0])

            return {
                "duration": duration,
                "width": width,
                "height": height,
                "fps": fps_value,
                "codec": video_info.get("codec_name"),
                "bitrate": int(probe["format"].get("bit_rate", 0)),
            }

        except ffmpeg.Error as e:
            stderr = e.stderr.decode() if e.stderr else ""
            raise RuntimeError(f"FFmpeg probe error: {stderr}") from e