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

"""字幕生成器 — 将字幕时间轴转换为 ASS 字幕文件。

与 render_subtitles.py 的区别：
- render_subtitles.py 处理静态整段标题/字幕
- 本模块处理带时间轴的多段 ASR 字幕

两者最终都输出 ASS 文件，供 FFmpeg 烧录。
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

from packages.domain.subtitle import SubtitleTimeline

logger = logging.getLogger(__name__)


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

DEFAULT_MAX_CHARS_PER_LINE = 20  # 每行最多字符数
DEFAULT_MIN_CHARS_PER_SEGMENT = 8  # 每段最少字符数


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


def _hex_to_ass_color(hex_color: str) -> str:
    """将 HEX 颜色（#RRGGBB）转换为 ASS &HBBGGRR 格式。"""
    hex_color = hex_color.lstrip("#")
    if len(hex_color) != 6:
        return "&H00FFFFFF"
    r, g, b = hex_color[0:2], hex_color[2:4], hex_color[4:6]
    return f"&H{b.upper()}{g.upper()}{r.upper()}"


def _position_to_ass_alignment(position: str) -> int:
    """将文字位置映射为 ASS \\an 对齐编号。"""
    mapping = {
        "top": 8,
        "center": 5,
        "bottom": 2,
    }
    return mapping.get(position, 2)


def _format_ass_time(seconds: float) -> str:
    """将秒数格式化为 ASS 时间格式 H:MM:SS.cc。"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = seconds % 60
    return f"{hours}:{minutes:02d}:{secs:05.2f}"


def _escape_ass_text(text: str) -> str:
    """转义 ASS 文本中的特殊字符。"""
    text = text.replace("\r\n", "\\N").replace("\n", "\\N").replace("\r", "\\N")
    text = text.replace("{", "(").replace("}", ")")
    return text


def _wrap_text(text: str, max_chars: int) -> list[str]:
    """将长文本按字数换行。

    优先在标点处换行，没有合适标点时硬切。
    """
    if len(text) <= max_chars:
        return [text]

    lines: list[str] = []
    remaining = text

    while len(remaining) > max_chars:
        # 在前 max_chars 个字符中找标点断开
        break_point = max_chars
        punctuations = "，。！？、；：,.;:!?"

        for i in range(max_chars, max_chars // 2, -1):
            if i < len(remaining) and remaining[i] in punctuations:
                break_point = i + 1
                break

        lines.append(remaining[:break_point])
        remaining = remaining[break_point:]

    if remaining:
        lines.append(remaining)

    return lines


# ── 主生成器 ─────────────────────────────────────────────────────────────────


def generate_ass_from_timeline(
    output_path: Path,
    timeline: SubtitleTimeline,
    *,
    video_width: int,
    video_height: int,
    subtitle_config: dict[str, Any] | None = None,
) -> Path:
    """从字幕时间轴生成 ASS 字幕文件。

    Args:
        output_path: 输出 ASS 文件路径
        timeline: 字幕时间轴
        video_width: 视频宽度
        video_height: 视频高度
        subtitle_config: 字幕样式配置（同 SubtitleConfig dict）

    Returns:
        生成的 ASS 文件路径
    """
    subtitle_config = subtitle_config or {}

    if not timeline.segments:
        output_path.write_text("", encoding="utf-8")
        return output_path

    # 样式参数
    font_name = subtitle_config.get("font", "思源黑体")
    font_size = int(subtitle_config.get("size", 24))
    color = _hex_to_ass_color(subtitle_config.get("color", "#ffffff"))
    position = subtitle_config.get("position", "bottom")
    alignment = _position_to_ass_alignment(position)
    max_chars_per_line = int(subtitle_config.get("max_chars_per_line", DEFAULT_MAX_CHARS_PER_LINE))

    # 描边（默认黑色描边，保证可读性）
    outline_color = "&H00000000"
    outline_width = 1.5

    # 边距
    margin_v = 60 if position == "bottom" else 60
    margin_l = 40
    margin_r = 40

    # 生成样式行
    style_line = (
        f"Style: Default,{font_name},{font_size},{color},"
        f"&H000000FF,{outline_color},&H00000000,"
        f"-1,0,0,0,100,100,0,0,"
        f"1,{outline_width},0,{alignment},"
        f"{margin_l},{margin_r},{margin_v},1"
    )

    # 生成事件行
    events: list[str] = []
    for seg in timeline.segments:
        start_time = _format_ass_time(seg.start)
        end_time = _format_ass_time(seg.end)

        # 自动换行
        lines = _wrap_text(seg.text, max_chars_per_line)
        display_text = "\\N".join(lines)

        safe_text = _escape_ass_text(display_text)

        events.append(f"Dialogue: 0,{start_time},{end_time},Default,,0,0,0,,{safe_text}")

    # 组装 ASS 文件
    ass_content = f"""[Script Info]
ScriptType: v4.00+
PlayResX: {video_width}
PlayResY: {video_height}
ScaledBorderAndShadow: yes
WrapStyle: 2
Encoding: UTF-8

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
{style_line}

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
{chr(10).join(events)}
"""

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(ass_content, encoding="utf-8")
    return output_path