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

"""统一渲染引擎 — 输入 EditPlan + EditPlanClips，按时间线+图层渲染视频.

核心原则（灵应）：渲染引擎是统一的，不判断模式，只按 clip_type/config.role
分组为图层再合成。

图层分组：
  main (无 config.role)     → main (z=0)
  main + config.role=b_roll → broll (z=0，与 main 同层替换)
  overlay                   → overlay (z=1，画中画叠加)
  background                → background (z=0，全屏底图)
  corner_voice              → corner_voice (z=1，右上角小窗)
  b_roll                    → broll (z=0)
  intro / outro             → main (z=0，按 order 排在首/尾)

合成流程：
  1. 每个 clip 先 trim + scale + setpts 预处理
  2. 同层 clips 按 order 用 xfade 串联
  3. overlay/corner_voice 层 overlay 到主层
  4. 如有独立音频轨，amix 混入
"""

from __future__ import annotations

import logging
import os
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from video_processing.ffmpeg_utils import (
    DEFAULT_FPS,
    DEFAULT_OUTPUT_HEIGHT,
    DEFAULT_OUTPUT_WIDTH,
    DEFAULT_TRANSITION_DURATION,
    FFMPEG_BIN,
    build_xfade_filter_chain,
    probe_duration,
    probe_video_info,
    run_ffmpeg,
)

logger = logging.getLogger(__name__)


# ── 数据结构 ──────────────────────────────────────────────────────────────────


@dataclass
class ResolvedClip:
    """已解析到本地路径的片段。"""

    clip_id: str
    asset_id: str
    local_path: Path
    clip_type: str
    order: int
    start_time: float = 0.0
    duration: float = 0.0  # 0 表示使用素材完整时长
    transition_effect: str = "cut"
    config: dict[str, Any] = field(default_factory=dict)

    # 运行时填充
    actual_duration: float = 0.0  # 素材实际时长（probe 后填充）


@dataclass
class RenderLayer:
    """渲染图层。"""

    role: str  # "main" | "overlay" | "pip" | "background" | "corner_voice" | "broll" | "audio"
    clips: list[ResolvedClip] = field(default_factory=list)
    z_index: int = 0
    opacity: float = 1.0
    position: tuple[int, int] | None = None  # (x, y) 偏移，None 表示全屏


@dataclass
class RenderResult:
    """渲染结果。"""

    output_path: Path
    duration: float
    file_size: int
    width: int
    height: int


# ── clip_type → layer role 映射 ──────────────────────────────────────────────


def _resolve_layer_role(clip_type: str, config: dict[str, Any]) -> str:
    """根据 clip_type 和 config.role 确定图层角色。

    映射规则：
      intro / outro         → "main"（按 order 排在首/尾）
      overlay               → "overlay"（画中画叠加，z=1）
      corner_voice          → "corner_voice"（右上角小窗，z=1）
      background            → "background"（全屏底图，z=0）
      b_roll                → "broll"（z=0）
      main + config.role=b_roll → "broll"
      main (default)        → "main"
    """
    role = config.get("role", "")

    if clip_type in ("intro", "outro"):
        return "main"
    if clip_type == "overlay":
        return "overlay"
    if clip_type == "corner_voice":
        return "corner_voice"
    if clip_type == "background":
        return "background"
    if clip_type == "b_roll":
        return "broll"
    # main type
    if role == "b_roll":
        return "broll"
    return "main"


# ── 图层默认 z_index ─────────────────────────────────────────────────────────

_LAYER_Z_INDEX: dict[str, int] = {
    "background": -1,
    "broll": 0,
    "main": 0,
    "overlay": 1,
    "corner_voice": 1,
    "audio": 2,
}

# 图层默认 PiP 位置（相对输出画布的偏移）
_PIP_SCALE = 0.25  # PiP 占主画面的比例


# ── 统一渲染引擎 ─────────────────────────────────────────────────────────────


class UnifiedRenderService:
    """统一渲染引擎。

    输入 EditPlan + EditPlanClips + 素材路径映射，按时间线+图层执行渲染。
    """

    def __init__(
        self,
        plan: Any,  # EditPlan
        clips: list[Any],  # list[EditPlanClip]
        asset_path_map: dict[str, Path],  # asset_id → local_path
        work_dir: Path,
        *,
        output_width: int = DEFAULT_OUTPUT_WIDTH,
        output_height: int = DEFAULT_OUTPUT_HEIGHT,
        output_fps: int = DEFAULT_FPS,
        transition_duration: float = DEFAULT_TRANSITION_DURATION,
    ):
        self.plan = plan
        self.clips = clips
        self.asset_path_map = asset_path_map
        self.work_dir = work_dir
        self.output_width = output_width
        self.output_height = output_height
        self.output_fps = output_fps
        self.transition_duration = transition_duration

    def render(self) -> RenderResult:
        """执行渲染，返回 RenderResult。

        Raises:
            ValueError: 没有可渲染的片段时抛出
        """
        # 1. 解析 clips → ResolvedClips（跳过无素材的 clip）
        resolved = self._resolve_clips()
        if not resolved:
            raise ValueError("没有可渲染的片段（所有片段素材缺失或下载失败）")

        # 2. 分组为 RenderLayers
        layers = self._group_clips_into_layers(resolved)

        # 3. 构建 filter_complex
        output_path = self.work_dir / f"rendered_{self.plan.id}.mp4"
        filter_complex, input_args = self._build_filter_complex(layers)

        # 4. 执行 FFmpeg
        self._execute_ffmpeg(filter_complex, input_args, output_path)

        # 5. 探测输出
        duration, file_size, width, height = self._probe_output(output_path)

        return RenderResult(
            output_path=output_path,
            duration=duration,
            file_size=file_size,
            width=width,
            height=height,
        )

    # ── 内部方法 ──────────────────────────────────────────────────────────────

    def _resolve_clips(self) -> list[ResolvedClip]:
        """将 EditPlanClip 列表解析为 ResolvedClip 列表。

        跳过 asset_id 为空或在 asset_path_map 中找不到的片段。
        """
        resolved: list[ResolvedClip] = []
        for clip in self.clips:
            asset_id = clip.asset_id
            if not asset_id:
                logger.warning("片段无素材: clip_id=%s", clip.id)
                continue

            local_path = self.asset_path_map.get(asset_id)
            if local_path is None or not local_path.exists():
                logger.warning("素材不存在: clip_id=%s asset_id=%s", clip.id, asset_id)
                continue

            # 探测实际时长
            try:
                actual_duration = probe_duration(local_path)
            except Exception:
                actual_duration = clip.duration or 5.0

            rc = ResolvedClip(
                clip_id=clip.id,
                asset_id=asset_id,
                local_path=local_path,
                clip_type=clip.clip_type,
                order=clip.order,
                start_time=clip.start_time,
                duration=clip.duration,
                transition_effect=clip.transition_effect or "cut",
                config=clip.config or {},
                actual_duration=actual_duration,
            )
            resolved.append(rc)

        # 按 order 排序
        resolved.sort(key=lambda c: c.order)
        return resolved

    def _group_clips_into_layers(self, resolved_clips: list[ResolvedClip]) -> list[RenderLayer]:
        """将 ResolvedClips 分组为 RenderLayers。

        分组规则见 _resolve_layer_role 函数文档。
        """
        layer_map: dict[str, RenderLayer] = {}

        for clip in resolved_clips:
            role = _resolve_layer_role(clip.clip_type, clip.config)
            if role not in layer_map:
                z = _LAYER_Z_INDEX.get(role, 0)
                layer_map[role] = RenderLayer(role=role, z_index=z)
            layer_map[role].clips.append(clip)

        # 每个 layer 内的 clips 按 order 排序
        for layer in layer_map.values():
            layer.clips.sort(key=lambda c: c.order)

        # 计算 PiP 位置
        pip_width = int(self.output_width * _PIP_SCALE)
        pip_height = int(self.output_height * _PIP_SCALE)
        margin = 20  # 边距

        if "overlay" in layer_map:
            layer_map["overlay"].position = (
                self.output_width - pip_width - margin,
                margin,
            )
        if "corner_voice" in layer_map:
            layer_map["corner_voice"].position = (
                self.output_width - pip_width - margin,
                margin,
            )

        # 按 z_index 排序返回
        layers = sorted(layer_map.values(), key=lambda lyr: lyr.z_index)
        return layers

    def _build_filter_complex(self, layers: list[RenderLayer]) -> tuple[str, list[str]]:
        """构建 FFmpeg filter_complex 字符串和输入参数列表。

        Returns:
            (filter_complex_str, input_args_list)
            input_args_list 是 ["-i", path1, "-i", path2, ...] 格式
        """
        if not layers:
            raise ValueError("没有可渲染的图层")

        # 收集所有 clips（按图层顺序，同层按 order）
        all_clips: list[ResolvedClip] = []
        for layer in layers:
            all_clips.extend(layer.clips)

        # 构建输入参数
        input_args: list[str] = []
        clip_to_input_idx: dict[str, int] = {}
        for i, clip in enumerate(all_clips):
            input_args.extend(["-i", str(clip.local_path)])
            clip_to_input_idx[clip.clip_id] = i

        filter_parts: list[str] = []

        # Step 1: 预处理每个 clip — scale + setpts
        # 为每个 clip 生成预处理后的标签 [v0], [v1], ...
        preprocessed_labels: list[str] = []
        for i, clip in enumerate(all_clips):
            label = f"v{i}"
            role = _resolve_layer_role(clip.clip_type, clip.config)

            filters: list[str] = []

            # trim — 始终将输出截断到有效时长，防止 xfade offset 与实际时长不匹配
            # 有效时长 = min(指定时长, 实际时长)；若均未设置则跳过
            effective_duration = 0.0
            if clip.duration > 0:
                effective_duration = (
                    min(clip.duration, clip.actual_duration) if clip.actual_duration > 0 else clip.duration
                )
            elif clip.actual_duration > 0:
                effective_duration = clip.actual_duration

            if effective_duration > 0:
                filters.append(f"trim=duration={effective_duration}")
                filters.append("setpts=PTS-STARTPTS")

            # scale
            if role in ("overlay", "corner_voice"):
                pip_w = int(self.output_width * _PIP_SCALE)
                pip_h = int(self.output_height * _PIP_SCALE)
                filters.append(f"scale={pip_w}:{pip_h}")
            elif role == "background":
                filters.append(
                    f"scale={self.output_width}:{self.output_height}" ":force_original_aspect_ratio=increase"
                )
                filters.append(f"crop={self.output_width}:{self.output_height}")
            else:
                # main / broll: scale + pad 保持宽高比
                filters.append(
                    f"scale={self.output_width}:{self.output_height}" ":force_original_aspect_ratio=decrease"
                )
                filters.append(f"pad={self.output_width}:{self.output_height}" ":(ow-iw)/2:(oh-ih)/2:black")

            filters.append("setpts=PTS-STARTPTS")
            filters.append(f"fps={self.output_fps}")

            filter_str = f"[{i}:v]{','.join(filters)}[{label}]"
            filter_parts.append(filter_str)
            preprocessed_labels.append(label)

        # Step 2: 同层 clips 用 xfade 串联
        layer_output_labels: dict[str, str] = {}
        for layer in layers:
            layer_clip_indices = [all_clips.index(c) for c in layer.clips]
            layer_labels = [preprocessed_labels[i] for i in layer_clip_indices]
            # 使用 trim 后的有效时长，与 Step 1 的 trim=duration 保持一致
            layer_durations = []
            for i in layer_clip_indices:
                c = all_clips[i]
                if c.duration > 0:
                    eff = min(c.duration, c.actual_duration) if c.actual_duration > 0 else c.duration
                else:
                    eff = c.actual_duration if c.actual_duration > 0 else 0.0
                layer_durations.append(eff)
            layer_transitions = [all_clips[i].transition_effect for i in layer_clip_indices]

            if len(layer_labels) == 1:
                # 单 clip 层，直接使用预处理标签
                layer_output_labels[layer.role] = layer_labels[0]
            else:
                # 多 clip 层，用 xfade 串联
                out_label = f"{layer.role}_merged"
                xfade_filter, _ = build_xfade_filter_chain(
                    clip_durations=layer_durations,
                    clip_video_labels=layer_labels,
                    transitions=layer_transitions,
                    transition_duration=self.transition_duration,
                    output_label=out_label,
                )
                if xfade_filter:
                    filter_parts.append(xfade_filter)
                layer_output_labels[layer.role] = out_label

        # Step 3: 合成各层
        # 找到主层 — background 优先作为底图，其次 broll / main
        final_video_label = None

        if "background" in layer_output_labels:
            final_video_label = layer_output_labels["background"]
            # b_roll / main 叠加到 background 上
            for role in ("broll", "main"):
                if role in layer_output_labels:
                    base_label = layer_output_labels[role]
                    combined_label = f"combined_{role}"
                    filter_parts.append(
                        f"[{final_video_label}][{base_label}]" f"overlay=(W-w)/2:(H-h)/2[{combined_label}]"
                    )
                    final_video_label = combined_label
        else:
            # 无 background 时，取 broll 或 main 作为基础
            for role in ("broll", "main"):
                if role in layer_output_labels:
                    final_video_label = layer_output_labels[role]
                    break

        if final_video_label is None:
            # 没有任何主层，使用第一个层
            final_video_label = layer_output_labels[layers[0].role]

        # 叠加 overlay 层
        for layer in layers:
            if layer.role in ("overlay", "corner_voice"):
                if layer.role not in layer_output_labels:
                    continue
                overlay_label = layer_output_labels[layer.role]
                x, y = layer.position or (
                    self.output_width - int(self.output_width * _PIP_SCALE) - 20,
                    20,
                )
                combined_label = f"combined_{layer.role}"
                filter_parts.append(f"[{final_video_label}][{overlay_label}]" f"overlay={x}:{y}[{combined_label}]")
                final_video_label = combined_label

        filter_parts.append(f"[{final_video_label}]format=yuv420p[final_video]")

        filter_complex = ";".join(filter_parts)
        return filter_complex, input_args

    def _execute_ffmpeg(
        self,
        filter_complex: str,
        input_args: list[str],
        output_path: Path,
    ) -> None:
        """执行 FFmpeg 渲染命令。

        失败时记录完整 filter_complex 以便排查（如 exit code 183）。
        """
        command = [
            FFMPEG_BIN,
            "-y",
            *input_args,
            "-filter_complex",
            filter_complex,
            "-map",
            "[final_video]",
            "-c:v",
            "libx264",
            "-crf",
            "23",
            "-preset",
            "medium",
            "-pix_fmt",
            "yuv420p",
            "-movflags",
            "+faststart",
            str(output_path),
        ]

        logger.info(
            "执行渲染: plan_id=%s inputs=%d output=%s",
            self.plan.id,
            input_args.count("-i"),
            output_path,
        )
        try:
            run_ffmpeg(command)
        except subprocess.CalledProcessError as e:
            # 额外记录 filter_complex，方便排查滤镜链构建问题
            logger.error(
                "渲染失败: plan_id=%s exit_code=%d\nfilter_complex:\n%s",
                self.plan.id,
                e.returncode,
                filter_complex[:5000],
            )
            raise

    def _probe_output(self, output_path: Path) -> tuple[float, int, int, int]:
        """探测输出文件的时长、大小、宽高。

        Returns:
            (duration, file_size, width, height)
        """
        info = probe_video_info(str(output_path))
        file_size = output_path.stat().st_size if output_path.exists() else 0
        return (
            info["duration"],
            file_size,
            info["width"],
            info["height"],
        )