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

"""音频混音模块 — 从 unified_render_service.py 拆分.

职责：
- 主图层音频 concat 拼接
- 独立音频轨 amix 混音
- 音视频合并（mux）

所有函数接收 RenderContext 获取共享依赖（work_dir、plan_id 等），
避免直接依赖 UnifiedRenderService 类。
"""

from __future__ import annotations

import logging
import subprocess
from dataclasses import dataclass, field
from pathlib import Path

from video_processing.ffmpeg_utils import FFMPEG_BIN, probe_has_audio, run_ffmpeg

# 延迟导入避免循环依赖：unified_render_service 定义 ResolvedClip / RenderLayer，
# 本模块提供音频函数供 unified_render_service 调用。
# 使用 from __future__ import annotations + TYPE_CHECKING 解决类型引用。
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from video_processing.unified_render_service import ResolvedClip, RenderLayer

logger = logging.getLogger(__name__)


@dataclass
class RenderContext:
    """渲染上下文 — 提供音频混音所需的共享依赖."""

    work_dir: Path
    plan_id: str
    # 音频探测缓存（避免同一 clip 被多次 ffprobe）
    _audio_cache: dict[str, bool] = field(default_factory=dict)


# ── 工具函数 ──────────────────────────────────────────────────────────────────


def clip_effective_duration(clip: ResolvedClip) -> float:
    """计算 clip 的有效时长.

    与 UnifiedRenderService._clip_effective_duration 逻辑一致。
    """
    if clip.duration > 0:
        return min(clip.duration, clip.actual_duration) if clip.actual_duration > 0 else clip.duration
    return clip.actual_duration if clip.actual_duration > 0 else 0.0


def clip_has_audio(ctx: RenderContext, clip: ResolvedClip) -> bool:
    """探测 clip 是否有音频流（带缓存）.

    避免同一个 clip 被多次 ffprobe 探测。
    """
    key = str(clip.local_path)
    if key not in ctx._audio_cache:
        ctx._audio_cache[key] = probe_has_audio(clip.local_path)
    return ctx._audio_cache[key]


# ── 音频混音 ──────────────────────────────────────────────────────────────────


def mix_audio(
    ctx: RenderContext,
    layers: list[RenderLayer],
    video_duration: float,
) -> Path | None:
    """音频后处理混音.

    处理逻辑：
    1. 主音频源按优先级查找：main > broll（background 不参与主音频，通常是图片无音轨）
    2. 主图层音频按顺序 concat 拼接
    3. 独立音频轨（audio role）用 amix 混入
    4. 输出时长截断到 video_duration
    5. 无音频流的 clip 会被自动跳过，避免 FFmpeg 引用 [i:a] 失败

    Args:
        ctx: 渲染上下文
        layers: 图层列表
        video_duration: 视频总时长（用于截断音频）

    Returns:
        混音后的音频文件路径，无音频时返回 None
    """
    # 按优先级精确查找主音频图层：main > broll
    # background 不参与主音频（通常是静态图片，无音轨）
    layer_map = {layer.role: layer for layer in layers}
    main_layer = None
    for role in ("main", "broll"):
        if role in layer_map and layer_map[role].clips:
            main_layer = layer_map[role]
            break

    main_clips: list[ResolvedClip] = main_layer.clips if main_layer else []

    # 没有主视频图层时兜底：检查 overlay/corner_voice 层是否有带音频的素材
    if not main_clips:
        for role in ("overlay", "corner_voice"):
            if role in layer_map and layer_map[role].clips:
                main_clips = layer_map[role].clips
                break

    # 收集独立音频轨
    audio_clips: list[ResolvedClip] = []
    if "audio" in layer_map:
        audio_clips = layer_map["audio"].clips

    # ── 防御：过滤掉无音频流的 clip ──
    main_clips = [c for c in main_clips if clip_has_audio(ctx, c)]
    audio_clips = [c for c in audio_clips if clip_has_audio(ctx, c)]

    if not main_clips and not audio_clips:
        return None

    # 构建音频处理命令
    output_path = ctx.work_dir / f"audio_{ctx.plan_id}.aac"

    # 简单场景：只有主图层 + 无独立音频 → 直接从视频提取音频并拼接
    if main_clips and not audio_clips:
        concat_main_audio(ctx, main_clips, output_path, video_duration)
        return output_path

    # 有独立音频轨 → amix 混音
    mix_with_independent_audio(ctx, main_clips, audio_clips, output_path, video_duration)
    return output_path


def concat_main_audio(
    ctx: RenderContext,
    clips: list[ResolvedClip],
    output_path: Path,
    video_duration: float,
) -> None:
    """主图层音频 concat 拼接（对齐链路A行为）.

    每个 clip 提取音频 → trim → 按顺序 concat。
    """
    if len(clips) == 1:
        # 单 clip，直接提取音频，截断到 min(clip有效时长, 视频总时长)
        clip = clips[0]
        effective_duration = clip_effective_duration(clip)
        # 最终时长：取 clip 有效时长和视频总时长的较小值
        # （视频总时长由主图层决定，但单 clip 场景下两者应该一致，仍做保护）
        final_duration = effective_duration
        if video_duration > 0 and (final_duration <= 0 or final_duration > video_duration):
            final_duration = video_duration

        command = [
            FFMPEG_BIN,
            "-y",
            "-i",
            str(clip.local_path),
            "-vn",
            "-acodec",
            "aac",
            "-b:a",
            "128k",
        ]
        if final_duration > 0:
            command.extend(["-t", f"{final_duration:.3f}"])
        command.append(str(output_path))
        run_ffmpeg(command)
        return

    # 多 clip，用 filter_complex concat
    input_args: list[str] = []
    filter_parts: list[str] = []

    for i, clip in enumerate(clips):
        input_args.extend(["-i", str(clip.local_path)])
        effective_duration = clip_effective_duration(clip)
        if effective_duration > 0:
            filter_parts.append(f"[{i}:a]atrim=0:{effective_duration:.3f},asetpts=PTS-STARTPTS[a{i}]")
        else:
            filter_parts.append(f"[{i}:a]asetpts=PTS-STARTPTS[a{i}]")

    audio_labels = "".join(f"[a{i}]" for i in range(len(clips)))
    filter_parts.append(f"{audio_labels}concat=n={len(clips)}:v=0:a=1[outa]")

    # 截断到视频总时长
    if video_duration > 0:
        filter_parts.append(f"[outa]atrim=0:{video_duration:.3f}[final_audio]")
        final_label = "final_audio"
    else:
        final_label = "outa"

    filter_complex = ";".join(filter_parts)

    command = [
        FFMPEG_BIN,
        "-y",
        *input_args,
        "-filter_complex",
        filter_complex,
        "-map",
        f"[{final_label}]",
        "-acodec",
        "aac",
        "-b:a",
        "128k",
        str(output_path),
    ]
    run_ffmpeg(command)


def mix_with_independent_audio(
    ctx: RenderContext,
    main_clips: list[ResolvedClip],
    audio_clips: list[ResolvedClip],
    output_path: Path,
    video_duration: float,
) -> None:
    """主音频 + 独立音频轨 amix 混音.

    Args:
        ctx: 渲染上下文
        main_clips: 主视频 clips（提取音频后 concat）
        audio_clips: 独立音频轨 clips
        output_path: 输出路径
        video_duration: 视频总时长
    """
    input_args: list[str] = []
    filter_parts: list[str] = []
    mix_labels: list[str] = []

    input_idx = 0

    # 1. 主图层音频 concat
    if main_clips:
        for clip in main_clips:
            input_args.extend(["-i", str(clip.local_path)])
            effective_duration = clip_effective_duration(clip)
            if effective_duration > 0:
                filter_parts.append(
                    f"[{input_idx}:a]atrim=0:{effective_duration:.3f},asetpts=PTS-STARTPTS[ma{input_idx}]"
                )
            else:
                filter_parts.append(f"[{input_idx}:a]asetpts=PTS-STARTPTS[ma{input_idx}]")
            input_idx += 1

        if len(main_clips) == 1:
            mix_labels.append("ma0")
        else:
            main_labels = "".join(f"[ma{i}]" for i in range(len(main_clips)))
            filter_parts.append(f"{main_labels}concat=n={len(main_clips)}:v=0:a=1[main_audio]")
            mix_labels.append("main_audio")

    # 2. 独立音频轨
    for j, clip in enumerate(audio_clips):
        input_args.extend(["-i", str(clip.local_path)])
        effective_duration = clip_effective_duration(clip)
        volume = clip.config.get("volume", 1.0) if clip.config else 1.0
        label = f"ia{j}"
        filters = []
        if effective_duration > 0:
            filters.append(f"atrim=0:{effective_duration:.3f}")
        filters.append("asetpts=PTS-STARTPTS")
        if volume != 1.0:
            filters.append(f"volume={volume}")
        filter_parts.append(f"[{input_idx}:a]{','.join(filters)}[{label}]")
        mix_labels.append(label)
        input_idx += 1

    # 3. amix 混音
    mix_inputs = "".join(f"[{label}]" for label in mix_labels)
    n_inputs = len(mix_labels)
    # normalized=0 保持音量，duration=shortest 取最短
    filter_parts.append(f"{mix_inputs}amix=inputs={n_inputs}:duration=longest:normalize=0[mixed_audio]")

    # 4. 截断到视频时长
    if video_duration > 0:
        filter_parts.append(f"[mixed_audio]atrim=0:{video_duration:.3f}[final_audio]")
        final_label = "final_audio"
    else:
        final_label = "mixed_audio"

    filter_complex = ";".join(filter_parts)

    command = [
        FFMPEG_BIN,
        "-y",
        *input_args,
        "-filter_complex",
        filter_complex,
        "-map",
        f"[{final_label}]",
        "-acodec",
        "aac",
        "-b:a",
        "128k",
        str(output_path),
    ]

    logger.info(
        "音频混音: plan_id=%s main_clips=%d audio_clips=%d",
        ctx.plan_id,
        len(main_clips),
        len(audio_clips),
    )
    try:
        run_ffmpeg(command)
    except subprocess.CalledProcessError as e:
        logger.error(
            "音频混音失败: plan_id=%s exit_code=%d\nfilter_complex:\n%s",
            ctx.plan_id,
            e.returncode,
            filter_complex[:3000],
        )
        raise


def merge_audio_video(
    ctx: RenderContext,
    video_path: Path,
    audio_path: Path,
    output_path: Path,
) -> None:
    """将音频合并到视频中（视频流拷贝，音频直接复用）.

    Args:
        ctx: 渲染上下文
        video_path: 无声视频路径
        audio_path: 音频文件路径
        output_path: 输出文件路径
    """
    command = [
        FFMPEG_BIN,
        "-y",
        "-i",
        str(video_path),
        "-i",
        str(audio_path),
        "-c:v",
        "copy",
        "-c:a",
        "aac",
        "-b:a",
        "128k",
        "-map",
        "0:v:0",
        "-map",
        "1:a:0",
        "-shortest",
        "-movflags",
        "+faststart",
        str(output_path),
    ]

    logger.info("合并音视频: plan_id=%s", ctx.plan_id)
    try:
        run_ffmpeg(command)
    except subprocess.CalledProcessError as e:
        logger.error(
            "合并音视频失败: plan_id=%s exit_code=%d",
            ctx.plan_id,
            e.returncode,
        )
        raise