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.ass_subtitle_builder import (
    TITLE_MARGIN_SIDE,
    TITLE_MARGIN_TOP,
    _wrap_title_text,
    build_ass_style,
    escape_ass_text,
    format_ass_time,
    hex_to_ass_color,
    position_to_ass_alignment,
)
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,
    video_duration: float = 0.0,
    subtitle_config: dict[str, Any] | None = None,
    title_text: str = "",
    title_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}")

    # ── 标题样式与事件（叠加在 ASR 字幕之上）───────────────────────────
    title_cfg = title_config or {}
    if not isinstance(title_cfg, dict):
        title_cfg = {}
    title_enabled = title_cfg.get("enabled", True) and bool(title_text.strip())

    title_style_line = ""
    title_event_line = ""

    if title_enabled:
        # 兼容 boolean stroke/shadow → dict
        _stroke_val = title_cfg.get("stroke")
        if isinstance(_stroke_val, bool):
            title_cfg["stroke"] = (
                {"enabled": _stroke_val, "color": "#000000", "width": 2} if _stroke_val else {"enabled": False}
            )
        _shadow_val = title_cfg.get("shadow")
        if isinstance(_shadow_val, bool):
            title_cfg["shadow"] = (
                {"enabled": _shadow_val, "color": "#000000", "blur": 4, "offset_x": 2, "offset_y": 2}
                if _shadow_val
                else {"enabled": False}
            )

        # 字段名归一化: font_size→size, font_color→color
        if "font_size" in title_cfg and "size" not in title_cfg:
            title_cfg["size"] = title_cfg["font_size"]
        if "font_color" in title_cfg and "color" not in title_cfg:
            title_cfg["color"] = title_cfg["font_color"]

        t_color = hex_to_ass_color(title_cfg.get("color", "#ffffff"))
        t_stroke = title_cfg.get("stroke", {}) or {}
        t_shadow = title_cfg.get("shadow", {}) or {}
        s_color = hex_to_ass_color(t_stroke.get("color", "#000000"))
        s_width = float(t_stroke.get("width", 2)) if t_stroke.get("enabled", False) else 0.0
        sh_blur = float(t_shadow.get("blur", 4)) if t_shadow.get("enabled", False) else 0.0
        sh_offset = (
            t_shadow.get("offset_x", 2) if t_shadow.get("enabled", False) else 0,
            t_shadow.get("offset_y", 2) if t_shadow.get("enabled", False) else 0,
        )
        t_alignment = position_to_ass_alignment(title_cfg.get("position", "top"))

        title_style_line = build_ass_style(
            "TitleStyle",
            font_name=title_cfg.get("font", "思源黑体"),
            font_size=min(int(title_cfg.get("size", 36)), 36),
            primary_color=t_color,
            outline_color=s_color,
            outline_width=s_width,
            shadow_blur=sh_blur,
            shadow_offset=sh_offset,
            bold=bool(title_cfg.get("bold", True)),
            italic=bool(title_cfg.get("italic", False)),
            alignment=t_alignment,
            margin_v=TITLE_MARGIN_TOP,
            margin_l=TITLE_MARGIN_SIDE,
            margin_r=TITLE_MARGIN_SIDE,
        )

        t_font_size = min(int(title_cfg.get("size", 36)), 36)
        safe_raw = escape_ass_text(title_text.strip())
        safe_wrapped = _wrap_title_text(safe_raw, video_width, t_font_size)

        if video_duration > 0:
            t_end_time = format_ass_time(video_duration)
        else:
            t_end_time = format_ass_time((timeline.segments[-1].end + 5.0) if timeline.segments else 60.0)
        title_event_line = f"Dialogue: 0,0:00:00.00,{t_end_time},TitleStyle,,0,0,0,,{safe_wrapped}"

        # 组装 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  # noqa: E501
{chr(10).join(filter(None, [title_style_line, style_line]))}

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

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