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

"""视频封面抽帧工具 — 从视频中抽取帧作为封面，支持标题文字叠加。

统一封面管道：
- 从已渲染视频抽帧：标题已通过 ASS 字幕烧进视频，帧天然带标题，无需再叠加。
- 从源素材抽帧（API E2 兜底）：源素材无标题，通过 Pillow 在帧上绘制标题文字。
"""

from __future__ import annotations

import logging
import tempfile
from pathlib import Path

logger = logging.getLogger(__name__)

# ── 标题叠加（Pillow）──────────────────────────────────────────────────────
# 实现统一放在 packages/shared/title_overlay.py，API 和 Worker 共用。


def apply_title_overlay(
    image_path: str,
    title_text: str,
    *,
    color: str = "#ffffff",
    position: str = "bottom",
    font_size: int | None = None,
    margin_ratio: float = 0.06,
    stroke_width_ratio: float = 0.04,
) -> str:
    """在图片上绘制标题文字（指定颜色 + 黑色描边/阴影）。

    委托给 packages.shared.title_overlay.apply_title_to_image，
    保持 Worker 内调用方式不变。title_text 为空时直接返回原路径。
    """
    from packages.shared.title_overlay import apply_title_to_image

    if not title_text or not title_text.strip():
        return image_path
    result = apply_title_to_image(
        image_path,
        title_text,
        color=color,
        position=position,
        font_size=font_size,
        margin_ratio=margin_ratio,
        stroke_width_ratio=stroke_width_ratio,
    )
    return result or image_path


def extract_first_frame(
    video_path: str,
    output_path: str | None = None,
    *,
    width: int = -1,
    height: int = -1,
    timeout: int = 30,
    seek_ratio: float = 0.15,
    min_seek_seconds: float = 1.0,
) -> str:
    """抽取视频封面帧（默认取视频时长 15% 处的帧，避开片头纯色画面）。

    因为视频渲染时标题已通过 ASS 字幕烧录，抽取的帧天然带标题。

    Args:
        video_path: 视频文件路径
        output_path: 输出图片路径，不传则用临时文件
        width: 输出宽度（默认 -1，保持原始分辨率）
        height: 输出高度（默认 -1，保持原始分辨率）
        timeout: 超时时间（秒）
        seek_ratio: 抽帧位置占视频时长的比例（默认 0.15，即 15% 处）
        min_seek_seconds: 最小抽帧时间（秒），避免极短视频 seek 到 0

    Returns:
        生成的封面帧文件路径

    Raises:
        RuntimeError: ffmpeg 执行失败或输出文件为空
    """
    from video_processing.ffmpeg_utils import FFMPEG_BIN, probe_duration, run_ffmpeg

    _is_temp_output = False
    if output_path is None:
        tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
        tmp.close()
        output_path = tmp.name
        _is_temp_output = True

    try:
        # 计算抽帧时间点：取视频时长 * seek_ratio，最少 min_seek_seconds 秒
        try:
            duration = probe_duration(video_path)
            seek_time = max(min_seek_seconds, duration * seek_ratio)
        except Exception:
            # probe 失败时 fallback 到第1秒
            seek_time = min_seek_seconds

        # 格式化为 HH:MM:SS.xx
        seek_str = _format_seek_time(seek_time)

        # 构建 scale filter：如果指定了宽高则缩放，否则保持原始分辨率。
        # NOTE: scale_filter 在此处通过 if/else 分支赋值，之后不再被覆盖，
        #       后续 cmd / cmd2 均复用同一变量，逻辑无变化。
        if width > 0 or height > 0:
            w_str = str(width) if width > 0 else "-1"
            h_str = str(height) if height > 0 else "-1"
            scale_filter = f"scale={w_str}:{h_str}:force_original_aspect_ratio=decrease,format=yuvj420p"
        else:
            # 保持原始分辨率，只确保格式兼容
            scale_filter = "format=yuvj420p"

        # -ss 放在 -i 前面（input seeking，更快）
        # -vframes 1 只取一帧
        # -q:v 2 jpeg 高质量
        cmd = [
            FFMPEG_BIN,
            "-y",
            "-ss",
            seek_str,
            "-i",
            video_path,
            "-vframes",
            "1",
            "-vf",
            scale_filter,
            "-q:v",
            "2",
            output_path,
        ]

        try:
            run_ffmpeg(cmd, capture_output=True, timeout=timeout)
        except Exception:
            # 失败时退回到第0帧兜底
            cmd2 = [
                FFMPEG_BIN,
                "-y",
                "-i",
                video_path,
                "-ss",
                "00:00:00",
                "-vframes",
                "1",
                "-vf",
                scale_filter,
                "-q:v",
                "2",
                output_path,
            ]
            run_ffmpeg(cmd2, capture_output=True, timeout=timeout)

        if not Path(output_path).exists() or Path(output_path).stat().st_size == 0:
            raise RuntimeError(f"Cover frame extraction failed: {output_path}")

        return output_path
    except Exception:
        # 失败时清理自己创建的临时文件
        if _is_temp_output and output_path:
            try:
                Path(output_path).unlink(missing_ok=True)
            except Exception:
                pass
        raise


def _format_seek_time(seconds: float) -> str:
    """将秒数格式化为 HH:MM:SS.xx 格式。"""
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = seconds % 60
    return f"{h:02d}:{m:02d}:{s:05.2f}"


def generate_and_upload_thumbnail(
    video_path: str,
    storage_key: str,
    *,
    seek_ratio: float = 0.15,
) -> str:
    """从视频中提取一帧缩略图并上传到 OSS。

    Args:
        video_path: 视频文件路径
        storage_key: OSS 存储 key
        seek_ratio: 抽帧位置比例（默认 0.15）

    Returns:
        上传后的 URL 字符串

    Raises:
        RuntimeError: 抽帧或上传失败
    """
    from video_processing.oss_helpers import upload_to_oss

    tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    tmp.close()
    try:
        frame_path = extract_first_frame(video_path, output_path=tmp.name, seek_ratio=seek_ratio)
        url = upload_to_oss(frame_path, storage_key)
        if not url:
            raise RuntimeError(f"上传缩略图到 OSS 失败: {storage_key}")
        return url
    finally:
        Path(tmp.name).unlink(missing_ok=True)


def _extract_frames_via_mediakit(
    video_path: str,
    plan_id: str,
    num_frames: int,
) -> list[dict] | None:
    """使用 MediaKit 智能抽帧 API 提取封面帧。

    Args:
        video_path: 本地视频文件路径
        plan_id: 编辑计划 ID
        num_frames: 需要的帧数

    Returns:
        帧列表 [{"image_url": str, "timestamp": float}, ...]，失败返回 None
    """
    import uuid

    from video_processing.oss_helpers import upload_to_oss

    from packages.shared.mediakit_client import get_mediakit_client

    client = get_mediakit_client()
    if not client.is_available:
        logger.info("[thumbnail] MediaKit 未配置，跳过智能抽帧")
        return None

    # 1. 上传视频到 OSS 获取 URL
    try:
        video_storage_key = f"temp/{plan_id}/{uuid.uuid4().hex[:8]}_{Path(video_path).name}"
        video_url = upload_to_oss(video_path, video_storage_key)
        if not video_url:
            logger.warning("[thumbnail] 视频上传 OSS 失败，无法使用 MediaKit")
            return None
        logger.info("[thumbnail] 视频已上传 OSS: %s", video_url[:80])
    except Exception as e:
        logger.warning("[thumbnail] 视频上传 OSS 异常: %s，降级到 ffmpeg", e)
        return None

    # 2. 调用 MediaKit 智能抽帧
    try:
        frames = client.extract_frames(
            video_url=video_url,
            strategy="SceneChange",
            max_frames=num_frames * 2,  # 多取一些帧供选择
        )
        if not frames:
            logger.warning("[thumbnail] MediaKit 抽帧返回空，降级到 ffmpeg")
            return None

        # 选取最均匀的 num_frames 个帧
        if len(frames) > num_frames:
            step = len(frames) // num_frames
            frames = [frames[i * step] for i in range(num_frames)]

        logger.info("[thumbnail] MediaKit 抽帧成功: %d 帧", len(frames))
        return frames

    except Exception as e:
        logger.warning("[thumbnail] MediaKit 抽帧异常: %s，降级到 ffmpeg", e)
        return None
    finally:
        # 清理临时视频文件
        try:
            from video_processing.oss_helpers import delete_from_oss

            delete_from_oss(video_storage_key)
        except Exception:
            pass


def extract_and_upload_cover_frames(
    video_path: str,
    plan_id: str,
    *,
    task_id: str = "",
    num_frames: int = 3,
    title_text: str = "",
    title_color: str = "#ffffff",
    title_position: str = "bottom",
    title_font_size: int | None = None,
) -> list[dict]:
    """从视频中抽取多帧作为封面候选，上传到 OSS。

    优先使用 MediaKit 智能抽帧，失败时降级到 ffmpeg 直接抽帧。

    Args:
        video_path: 视频文件路径
        plan_id: 编辑计划 ID（用于生成 storage key）
        task_id: 任务 ID（用于生成独立的 storage key，避免标题变更时封面冲突）
        num_frames: 抽取帧数（默认 3）
        title_text: 标题文字；非空时用 Pillow 叠加到每帧。
            从已渲染视频抽帧时通常传空（标题已烧录）；从源素材抽帧时传标题。
        title_color: 标题字体颜色（#RRGGBB）
        title_position: 标题位置 top/center/bottom
        title_font_size: 标题字号，None 时自动计算

    Returns:
        封面候选列表，每项包含 {"url": str, "position": float}
    """
    import httpx
    from video_processing.ffmpeg_utils import probe_duration
    from video_processing.oss_helpers import upload_to_oss

    try:
        duration = probe_duration(video_path)
    except Exception:
        duration = 0.0

    candidates: list[dict] = []

    # 优先尝试 MediaKit 智能抽帧
    mediakit_frames = _extract_frames_via_mediakit(video_path, plan_id, num_frames)
    if mediakit_frames:
        for i, frame in enumerate(mediakit_frames):
            frame_url = frame.get("image_url")
            if not frame_url:
                continue
            tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
            tmp.close()
            try:
                # 下载 MediaKit 返回的帧图
                resp = httpx.get(frame_url, timeout=30, follow_redirects=True)
                resp.raise_for_status()
                with open(tmp.name, "wb") as f:
                    f.write(resp.content)

                # 叠加标题文字（如需要）
                if title_text and title_text.strip():
                    apply_title_overlay(
                        tmp.name,
                        title_text,
                        color=title_color,
                        position=title_position,
                        font_size=title_font_size,
                    )

                storage_key = f"covers/{plan_id}/{task_id}/mediakit_frame_{i}.jpg"
                url = upload_to_oss(tmp.name, storage_key)
                if url:
                    seek_time = frame.get("timestamp", 0.0)
                    candidates.append({"url": url, "position": round(seek_time, 2)})
            except Exception as e:
                logger.warning("[thumbnail] MediaKit 帧 %d 处理失败: %s", i, e)
            finally:
                Path(tmp.name).unlink(missing_ok=True)

        if len(candidates) >= num_frames:
            logger.info("[thumbnail] MediaKit 智能抽帧完成: %d 帧", len(candidates))
            return candidates[:num_frames]

        logger.warning("[thumbnail] MediaKit 抽帧不足 %d 帧，降级到 ffmpeg", num_frames)

    # Fallback: ffmpeg 直接抽帧
    logger.info("[thumbnail] 使用 ffmpeg 抽帧")
    # 均匀分布抽帧点：从 10% 到 90%
    for i in range(num_frames):
        ratio = 0.1 + 0.8 * i / max(num_frames - 1, 1)
        tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
        tmp.close()
        try:
            frame_path = extract_first_frame(
                video_path,
                output_path=tmp.name,
                seek_ratio=ratio,
                min_seek_seconds=0.5,
            )
            # 从源素材抽帧时叠加标题文字；已渲染视频标题已烧录时传空字符串跳过
            if title_text and title_text.strip():
                apply_title_overlay(
                    frame_path,
                    title_text,
                    color=title_color,
                    position=title_position,
                    font_size=title_font_size,
                )
            storage_key = f"covers/{plan_id}/{task_id}/frame_{i}.jpg"
            url = upload_to_oss(frame_path, storage_key)
            if url:
                seek_time = max(0.5, duration * ratio) if duration > 0 else 0.0
                candidates.append({"url": url, "position": round(seek_time, 2)})
        except Exception as e:
            logger.warning("[thumbnail] 封面候选帧 %d 提取失败: %s", i, e)
        finally:
            Path(tmp.name).unlink(missing_ok=True)

    return candidates