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

"""多轨道混音引擎 — 支持多路音频独立音量调节与混合.

基于 FFmpeg amix / amerge 实现：
- 支持任意数量音频轨道（原音、BGM、配音、音效等）
- 每轨独立音量调节
- 每轨独立淡入淡出
- 每轨独立时间偏移（delay）
- 总输出音量归一化补偿

作为 render_audio.py 的增强模块，在 mix_audio 后处理阶段被调用。
与 bgm_mixer.py 的关系：
- bgm_mixer 专注 BGM 单轨道的复杂处理（循环、人声闪避）
- 本模块专注多路轨道的统一音量调节与混合
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

from video_processing.ffmpeg_utils import FFMPEG_BIN, probe_duration, run_ffmpeg

if TYPE_CHECKING:
    from video_processing.render_audio import RenderContext

logger = logging.getLogger(__name__)


# ── 常量 ──────────────────────────────────────────────────────────────────────

TRACK_TYPE_MAIN = "main"        # 原音（视频原声）
TRACK_TYPE_BGM = "bgm"          # 背景音乐
TRACK_TYPE_VOICEOVER = "voiceover"  # 配音（TTS/人声）
TRACK_TYPE_SFX = "sfx"          # 音效
TRACK_TYPE_AMBIENT = "ambient"  # 环境音

# 各轨道默认音量（相对主音频）
DEFAULT_VOLUMES = {
    TRACK_TYPE_MAIN: 1.0,
    TRACK_TYPE_BGM: 0.3,
    TRACK_TYPE_VOICEOVER: 1.0,
    TRACK_TYPE_SFX: 0.7,
    TRACK_TYPE_AMBIENT: 0.2,
}


@dataclass
class AudioTrack:
    """单条音频轨道配置."""

    track_id: str                    # 轨道唯一标识
    track_type: str                  # 轨道类型（main/bgm/voiceover/sfx/ambient）
    audio_path: str                  # 音频文件路径
    volume: float = 1.0              # 音量 0.0 ~ 2.0
    fade_in: float = 0.0             # 淡入时长（秒）
    fade_out: float = 0.0            # 淡出时长（秒）
    start_time: float = 0.0          # 开始时间（相对于视频起点，秒）
    duration: float = 0.0            # 持续时长（0表示到文件末尾）
    enabled: bool = True             # 是否启用

    @classmethod
    def from_dict(cls, track: dict) -> "AudioTrack":
        """从字典创建 AudioTrack，带安全类型转换."""
        track_type = str(track.get("track_type", TRACK_TYPE_SFX))
        default_vol = DEFAULT_VOLUMES.get(track_type, 1.0)

        try:
            volume = float(track.get("volume", default_vol))
        except (TypeError, ValueError):
            volume = default_vol
        volume = max(0.0, min(2.0, volume))

        try:
            fade_in = max(0.0, float(track.get("fade_in", 0.0)))
        except (TypeError, ValueError):
            fade_in = 0.0

        try:
            fade_out = max(0.0, float(track.get("fade_out", 0.0)))
        except (TypeError, ValueError):
            fade_out = 0.0

        try:
            start_time = max(0.0, float(track.get("start_time", 0.0)))
        except (TypeError, ValueError):
            start_time = 0.0

        try:
            duration = max(0.0, float(track.get("duration", 0.0)))
        except (TypeError, ValueError):
            duration = 0.0

        return cls(
            track_id=str(track.get("track_id", "")),
            track_type=track_type,
            audio_path=str(track.get("audio_path", "")),
            volume=volume,
            fade_in=fade_in,
            fade_out=fade_out,
            start_time=start_time,
            duration=duration,
            enabled=bool(track.get("enabled", True)),
        )


@dataclass
class MultiTrackMixConfig:
    """多轨道混音配置."""

    tracks: list[AudioTrack] = field(default_factory=list)
    master_volume: float = 1.0       # 主输出音量
    normalize: bool = True           # 是否自动归一化补偿
    max_output_volume: float = 1.5   # 最大输出音量（防止爆音）

    @classmethod
    def from_config_dict(cls, config: dict | None) -> "MultiTrackMixConfig":
        """从 plan.config.audio_tracks 字典创建配置."""
        if not config or not isinstance(config, dict):
            return cls()

        tracks_raw = config.get("tracks", [])
        tracks: list[AudioTrack] = []

        if isinstance(tracks_raw, list):
            for t in tracks_raw:
                if isinstance(t, dict) and t.get("audio_path"):
                    try:
                        track = AudioTrack.from_dict(t)
                        if track.enabled and track.audio_path:
                            tracks.append(track)
                    except Exception:
                        logger.warning("[multi-track] skip invalid track config: %s", t)
                        continue

        try:
            master_volume = float(config.get("master_volume", 1.0))
            master_volume = max(0.0, min(2.0, master_volume))
        except (TypeError, ValueError):
            master_volume = 1.0

        return cls(
            tracks=tracks,
            master_volume=master_volume,
            normalize=bool(config.get("normalize", True)),
            max_output_volume=float(config.get("max_output_volume", 1.5)),
        )

    @property
    def has_effect(self) -> bool:
        """是否有有效轨道需要混音."""
        return len([t for t in self.tracks if t.enabled and t.audio_path]) > 0


# ── 单轨道预处理 ────────────────────────────────────────────────────────────


def _prepare_single_track(
    ctx: "RenderContext",
    track: AudioTrack,
    target_duration: float,
    output_path: Path,
) -> bool:
    """预处理单条轨道：音量 + 淡入淡出 + 时间偏移 + 截断.

    生成一个精确对齐时间轴的音频文件，后续统一 amix 混音。

    Returns:
        True 表示处理成功，False 表示失败（跳过）
    """
    try:
        audio_dur = probe_duration(track.audio_path)
    except Exception:
        logger.warning("[multi-track] probe failed, skip track: %s", track.track_id)
        return False

    if audio_dur <= 0:
        return False

    # 计算实际有效时长
    effective_start = track.start_time
    if track.duration > 0:
        effective_dur = min(track.duration, audio_dur)
    else:
        effective_dur = audio_dur

    # 如果轨道完全在视频时长之外，跳过
    if effective_start >= target_duration:
        return False
    if effective_start + effective_dur <= 0:
        return False

    # 构建滤镜链
    filter_parts: list[str] = []

    # 1. 先截断到有效范围
    trim_start = 0.0  # 从源文件的哪个位置开始取
    if effective_start < 0:
        trim_start = -effective_start
        effective_start = 0.0

    # 实际需要的源时长
    need_dur = min(effective_dur, target_duration - effective_start)
    if need_dur <= 0:
        return False

    filter_parts.append(f"atrim={trim_start:.3f}:{trim_start + need_dur:.3f}")
    filter_parts.append("asetpts=N/SR/TB")

    # 2. 音量调节
    if abs(track.volume - 1.0) > 0.001:
        filter_parts.append(f"volume={track.volume:.3f}")

    # 3. 淡入
    if track.fade_in > 0 and track.fade_in < need_dur:
        filter_parts.append(f"afade=t=in:st=0:d={track.fade_in:.3f}")

    # 4. 淡出
    if track.fade_out > 0 and track.fade_out < need_dur:
        fade_start = need_dur - track.fade_out
        if fade_start > 0:
            filter_parts.append(f"afade=t=out:st={fade_start:.3f}:d={track.fade_out:.3f}")

    # 5. 时间偏移（用 adelay 实现开头静音填充）
    if effective_start > 0.01:
        delay_ms = int(effective_start * 1000)
        filter_parts.append(f"adelay={delay_ms}|{delay_ms}")

    # 6. 最终截断到目标总时长
    filter_parts.append(f"atrim=0:{target_duration:.3f}")
    filter_parts.append("asetpts=N/SR/TB")

    filter_str = ",".join(filter_parts)

    command = [
        FFMPEG_BIN,
        "-y",
        "-i",
        track.audio_path,
        "-filter:a",
        filter_str,
        "-c:a",
        "aac",
        "-b:a",
        "128k",
        str(output_path),
    ]

    logger.info(
        "[multi-track] prepare track: id=%s type=%s vol=%.2f start=%.2f dur=%.2f",
        track.track_id,
        track.track_type,
        track.volume,
        effective_start,
        need_dur,
    )

    try:
        run_ffmpeg(command)
        return True
    except Exception as e:
        logger.warning("[multi-track] track prepare failed: %s, error=%s", track.track_id, e)
        return False


# ── 多轨道混音主入口 ─────────────────────────────────────────────────────────


def mix_multi_track(
    ctx: "RenderContext",
    main_audio_path: Path,
    config: MultiTrackMixConfig,
    target_duration: float,
) -> Path:
    """多轨道混音：主音频 + 多条附加轨道.

    Args:
        ctx: 渲染上下文
        main_audio_path: 主音频文件路径（原音）
        config: 多轨道混音配置
        target_duration: 目标总时长

    Returns:
        混音后的音频文件路径
    """
    output_path = ctx.work_dir / f"multi_track_mix_{ctx.plan_id}.aac"

    if target_duration <= 0:
        target_duration = 5.0

    # 收集所有有效轨道（已预处理好的）
    prepared_tracks: list[Path] = []

    # 主音频作为第0轨
    prepared_tracks.append(main_audio_path)

    # 预处理每条附加轨道
    for i, track in enumerate(config.tracks):
        if not track.enabled or not track.audio_path:
            continue

        track_out = ctx.work_dir / f"track_{i}_{ctx.plan_id}.aac"
        if _prepare_single_track(ctx, track, target_duration, track_out):
            prepared_tracks.append(track_out)

    # 如果只有主音频，直接返回（无需混音）
    if len(prepared_tracks) <= 1:
        import shutil

        shutil.copy2(main_audio_path, output_path)
        return output_path

    # 使用 amix 混音
    num_inputs = len(prepared_tracks)

    # 构建输入参数
    input_args: list[str] = []
    for tp in prepared_tracks:
        input_args.extend(["-i", str(tp)])

    # amix 的 duration=first 以第一个输入（主音频）时长为准
    # normalize 补偿：amix 会把每路音量除以 N，需要乘回来
    # 但如果所有轨道都同时有声，可能会爆音，所以用 master_volume 控制
    if config.normalize:
        # 经验值：不是所有轨道都同时有声，补偿系数取 N * 0.7
        compensate = num_inputs * 0.7
    else:
        compensate = 1.0

    final_volume = compensate * config.master_volume
    final_volume = min(final_volume, config.max_output_volume)

    # 构建 filter_complex
    inputs_label = "".join(f"[{i}:a]" for i in range(num_inputs))
    filter_complex = (
        f"{inputs_label}amix=inputs={num_inputs}:duration=first:dropout_transition=0[outa];"
        f"[outa]volume={final_volume:.3f}[final]"
    )

    command = [
        FFMPEG_BIN,
        "-y",
        *input_args,
        "-filter_complex",
        filter_complex,
        "-map",
        "[final]",
        "-c:a",
        "aac",
        "-b:a",
        "128k",
        str(output_path),
    ]

    logger.info(
        "[multi-track] mix %d tracks, master_vol=%.2f compensate=%.2f final_vol=%.2f",
        num_inputs,
        config.master_volume,
        compensate,
        final_volume,
    )

    try:
        run_ffmpeg(command)
    except Exception as e:
        logger.error("[multi-track] mix failed, fallback to main audio only: %s", e)
        import shutil

        shutil.copy2(main_audio_path, output_path)

    return output_path


# ── 便捷函数：从 plan.config 快速混音 ───────────────────────────────────────


def mix_audio_tracks_from_config(
    ctx: "RenderContext",
    main_audio_path: Path,
    audio_tracks_config: dict | None,
    target_duration: float,
) -> Path:
    """从 plan.config.audio_tracks 配置执行多轨道混音.

    降级策略：配置无效或混音失败时返回主音频。
    """
    config = MultiTrackMixConfig.from_config_dict(audio_tracks_config)
    if not config.has_effect:
        return main_audio_path

    return mix_multi_track(ctx, main_audio_path, config, target_duration)