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

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

from video_processing.ffmpeg_utils import FFMPEG_BIN, probe_has_audio, run_ffmpeg
from video_processing.reverse_engine import ReverseConfig, ReverseEngine
from video_processing.speed_engine import SpeedEngine

if TYPE_CHECKING:
    from video_processing.unified_render_service import RenderLayer, ResolvedClip

logger = logging.getLogger(__name__)


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

    work_dir: Path
    plan_id: str
    # 音频降噪配置（全局，对最终混音结果应用）
    noise_reduction_config: dict | None = None
    # 音频探测缓存（避免同一 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,
    *,
    bgm_path: str | None = None,
    bgm_config: dict | None = None,
    audio_tracks_config: dict | None = None,
) -> Path | None:
    """音频后处理混音.

    处理逻辑：
    1. 主音频源按优先级查找：main > broll（background 不参与主音频，通常是图片无音轨）
    2. 主图层音频按顺序 concat 拼接
    3. 独立音频轨（audio role）用 amix 混入
    4. 输出时长截断到 video_duration
    5. 无音频流的 clip 会被自动跳过，避免 FFmpeg 引用 [i:a] 失败
    6. 如果提供了 bgm_path，则额外混入 BGM（支持淡入淡出、循环、人声闪避）
    7. 如果配置了 audio_tracks，则混入多轨道音频（配音、音效等）
    8. 如果配置了降噪，最后应用降噪

    Args:
        ctx: 渲染上下文
        layers: 图层列表
        video_duration: 视频总时长（用于截断音频）
        bgm_path: BGM 音频本地路径，为 None 时不混入 BGM
        bgm_config: BGM 配置字典（volume/fade_in/fade_out/sidechain 等）
        audio_tracks_config: 多轨道音频配置（tracks/master_volume 等）

    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:
        # 没有主音频也没有独立音频 → 检查是否有 BGM
        if bgm_path and bgm_config and bgm_config.get("enabled", False):
            from video_processing.bgm_mixer import BGMConfig, build_bgm_only

            bgm_cfg = BGMConfig.from_config_dict(bgm_path, bgm_config)
            try:
                return build_bgm_only(ctx, bgm_cfg, video_duration)
            except Exception:
                logger.exception("[bgm] 纯BGM生成失败: plan_id=%s", ctx.plan_id)
        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)
    else:
        # 有独立音频轨 → amix 混音
        mix_with_independent_audio(ctx, main_clips, audio_clips, output_path, video_duration)

    # ── BGM 混音 ──
    if bgm_path and bgm_config and bgm_config.get("enabled", False):
        from video_processing.bgm_mixer import BGMConfig, mix_bgm_with_main

        bgm_cfg = BGMConfig.from_config_dict(bgm_path, bgm_config)

        try:
            # 这里 main_audio 就是 output_path，先有主音频再混 BGM
            final_path = mix_bgm_with_main(ctx, output_path, bgm_cfg, video_duration)
            output_path = final_path
        except Exception:
            logger.exception("[bgm] BGM 混音失败，回退到无 BGM 音频: plan_id=%s", ctx.plan_id)

    # ── 多轨道混音（配音/音效等） ──
    if audio_tracks_config and audio_tracks_config.get("enabled", False):
        from video_processing.multi_track_mixer import mix_audio_tracks_from_config

        try:
            tracks_config = audio_tracks_config.get("tracks_config") or audio_tracks_config
            multi_output = mix_audio_tracks_from_config(ctx, output_path, tracks_config, video_duration)
            if multi_output and multi_output != output_path:
                output_path = multi_output
        except Exception:
            logger.exception("[multi-track] 多轨道混音失败，回退: plan_id=%s", ctx.plan_id)

    return _apply_noise_reduction_if_needed(ctx, output_path)


def _apply_noise_reduction_if_needed(ctx: RenderContext, audio_path: Path) -> Path:
    """如果配置了音频降噪，对已生成的音频文件应用降噪。

    作为后处理步骤，对最终混音结果统一降噪。
    失败时返回原始文件路径，不阻断主流程。
    """
    if not ctx.noise_reduction_config:
        return audio_path

    try:
        from video_processing.noise_reduction_engine import NoiseReductionConfig, NoiseReductionEngine

        config = NoiseReductionConfig.from_dict(ctx.noise_reduction_config)
        if not config.has_effect():
            return audio_path

        engine = NoiseReductionEngine(config)
        filter_str = engine.build_filter("[0:a]", "[out]")
        # 提取滤镜部分（不带标签）
        filter_part = filter_str[len("[0:a]") : -len("[out]")]

        nr_output_path = audio_path.with_name(f"{audio_path.stem}_nr.aac")
        command = [
            FFMPEG_BIN,
            "-y",
            "-i",
            str(audio_path),
            "-af",
            filter_part,
            "-acodec",
            "aac",
            "-b:a",
            "128k",
            str(nr_output_path),
        ]
        run_ffmpeg(command)

        if nr_output_path.exists():
            return nr_output_path
        logger.warning("[noise-reduction] 降噪输出文件不存在，使用原始音频")
        return audio_path
    except Exception as e:
        logger.warning("[noise-reduction] 音频降噪失败，使用原始音频: %s", e)
        return audio_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)
        trim_start = getattr(clip, "start_time", 0) or 0
        speed = getattr(clip, "playback_speed", 1.0) or 1.0
        if not isinstance(speed, (int, float)) or speed <= 0:
            speed = 1.0

        # 调速后时长
        adjusted_duration = effective_duration / speed if abs(speed - 1.0) >= 1e-6 else effective_duration
        # 最终时长：取调速后时长和视频总时长的较小值
        final_duration = adjusted_duration
        if video_duration > 0 and (final_duration <= 0 or final_duration > video_duration):
            final_duration = video_duration

        # 音频倒放
        reverse_config = ReverseConfig.from_dict(clip.config.get("reverse"))
        has_reverse = reverse_config.enabled and reverse_config.reverse_audio
        has_speed = abs(speed - 1.0) >= 1e-6

        if not has_speed and not has_reverse:
            # 无调速无倒放：简单命令行，-ss 裁剪更高效
            command = [
                FFMPEG_BIN,
                "-y",
                "-i",
                str(clip.local_path),
                "-vn",
                "-acodec",
                "aac",
                "-b:a",
                "128k",
            ]
            if trim_start > 0:
                command.extend(["-ss", f"{trim_start:.3f}"])
            if final_duration > 0:
                command.extend(["-t", f"{final_duration:.3f}"])
            command.append(str(output_path))
            run_ffmpeg(command)
        else:
            # 有调速或倒放：用 filter_complex
            speed_engine = SpeedEngine()
            audio_filters = []
            if effective_duration > 0:
                audio_filters.append(f"atrim=start={trim_start:.3f}:duration={effective_duration:.3f}")
                audio_filters.append("asetpts=PTS-STARTPTS")

            # 音频调速
            if has_speed:
                from video_processing.speed_engine import SpeedConfig

                config = SpeedConfig(speed=float(speed))
                config.clamp()
                atempo_filter = speed_engine.build_audio_filter(config)
                if atempo_filter:
                    audio_filters.append(atempo_filter)

            # 音频倒放
            if has_reverse:
                reverse_filter = ReverseEngine.build_audio_filter(reverse_config, duration=effective_duration)
                if reverse_filter:
                    audio_filters.append(reverse_filter)

            filter_parts: list[str] = [f"[0:a]{','.join(audio_filters)}[outa]"]
            if video_duration > 0 and final_duration < adjusted_duration:
                filter_parts.append(f"[outa]atrim=0:{final_duration:.3f}[final_audio]")
                final_label = "final_audio"
            else:
                final_label = "outa"

            filter_complex = ";".join(filter_parts)
            command = [
                FFMPEG_BIN,
                "-y",
                "-i",
                str(clip.local_path),
                "-filter_complex",
                filter_complex,
                "-map",
                f"[{final_label}]",
                "-acodec",
                "aac",
                "-b:a",
                "128k",
                str(output_path),
            ]
            run_ffmpeg(command)
        return

    # 多 clip，用 filter_complex concat
    input_args: list[str] = []
    filter_parts: list[str] = []
    speed_engine = SpeedEngine()

    for i, clip in enumerate(clips):
        input_args.extend(["-i", str(clip.local_path)])
        effective_duration = clip_effective_duration(clip)
        trim_start = getattr(clip, "start_time", 0) or 0
        speed = getattr(clip, "playback_speed", 1.0) or 1.0
        if not isinstance(speed, (int, float)) or speed <= 0:
            speed = 1.0

        audio_filters: list[str] = []
        if effective_duration > 0:
            audio_filters.append(f"atrim=start={trim_start:.3f}:duration={effective_duration:.3f}")
            audio_filters.append("asetpts=PTS-STARTPTS")

        # 音频调速 — atempo 多级串联
        if abs(speed - 1.0) >= 1e-6:
            from video_processing.speed_engine import SpeedConfig

            config = SpeedConfig(speed=float(speed))
            config.clamp()
            atempo_filter = speed_engine.build_audio_filter(config)
            if atempo_filter:
                audio_filters.append(atempo_filter)
        else:
            audio_filters.append("asetpts=PTS-STARTPTS")

        # 音频倒放
        reverse_config = ReverseConfig.from_dict(clip.config.get("reverse"))
        if reverse_config.enabled and reverse_config.reverse_audio:
            reverse_filter = ReverseEngine.build_audio_filter(reverse_config, duration=effective_duration)
            if reverse_filter:
                audio_filters.append(reverse_filter)

        filter_parts.append(f"[{i}:a]{','.join(audio_filters)}[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)
            trim_start = getattr(clip, "start_time", 0) or 0
            if effective_duration > 0:
                filter_parts.append(
                    f"[{input_idx}:a]atrim=start={trim_start:.3f}:duration={effective_duration:.3f},"
                    f"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)
        trim_start = getattr(clip, "start_time", 0) or 0
        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=start={trim_start:.3f}:duration={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