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

"""水印引擎 — 基于 FFmpeg overlay 滤镜的水印叠加.

支持：
- 图片水印（PNG/logo）
- 文字水印（drawtext）
- 9宫格位置 + 边距配置
- 透明度/大小缩放
- 滚动水印（跑马灯）
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any

logger = logging.getLogger(__name__)

# 9宫格位置枚举
WATERMARK_POSITIONS = {
    "top_left": "左上",
    "top_center": "中上",
    "top_right": "右上",
    "center_left": "左中",
    "center": "中心",
    "center_right": "右中",
    "bottom_left": "左下",
    "bottom_center": "中下",
    "bottom_right": "右下",
}


@dataclass
class WatermarkConfig:
    """水印配置.

    mode: "image" 图片水印 | "text" 文字水印
    position: 9宫格位置
    opacity: 透明度 0.0-1.0
    scale: 缩放比例（图片水印），0.1-1.0
    margin: 边距（像素）
    scroll: 是否滚动（跑马灯）
    scroll_speed: 滚动速度（像素/秒）
    """

    mode: str = "text"  # image | text
    position: str = "bottom_right"

    # 图片水印
    image_path: str = ""  # 本地图片路径
    scale: float = 0.2  # 相对输出宽度的比例
    opacity: float = 0.8  # 0.0-1.0

    # 文字水印
    text: str = ""
    font_size: int = 24
    font_color: str = "white"
    font_path: str = ""  # 字体文件路径

    # 边距
    margin_x: int = 20
    margin_y: int = 20

    # 滚动水印
    scroll: bool = False
    scroll_speed: int = 50  # 像素/秒

    @classmethod
    def from_dict(cls, data: dict[str, Any] | None) -> WatermarkConfig | None:
        """从字典构造，空配置返回 None（不加水印）."""
        if not data:
            return None

        enabled = data.get("enabled", False)
        if not enabled:
            return None

        mode = data.get("mode", "text")

        # 图片模式需要 image_path；文字模式需要 text
        if mode == "image":
            image_path = data.get("image_path", "") or data.get("image", "") or ""
            if not image_path:
                logger.warning("图片水印缺少 image_path，跳过水印")
                return None
        elif mode == "text":
            text = data.get("text", "") or ""
            if not text:
                logger.warning("文字水印缺少 text，跳过水印")
                return None

        position = data.get("position", "bottom_right")
        if position not in WATERMARK_POSITIONS:
            position = "bottom_right"

        return cls(
            mode=mode,
            position=position,
            image_path=str(data.get("image_path", data.get("image", "")) or ""),
            scale=float(data.get("scale", 0.2)),
            opacity=float(data.get("opacity", 0.8)),
            text=str(data.get("text", "") or ""),
            font_size=int(data.get("font_size", 24)),
            font_color=str(data.get("font_color", "white")),
            font_path=str(data.get("font_path", "") or ""),
            margin_x=int(data.get("margin_x", 20)),
            margin_y=int(data.get("margin_y", 20)),
            scroll=bool(data.get("scroll", False)),
            scroll_speed=int(data.get("scroll_speed", 50)),
        )

    def validate(self) -> tuple[bool, str]:
        """校验配置是否有效."""
        if self.position not in WATERMARK_POSITIONS:
            return False, f"不支持的位置: {self.position}"

        if not (0.0 <= self.opacity <= 1.0):
            return False, "透明度必须在 0-1 之间"

        if self.mode == "image":
            if not self.image_path:
                return False, "图片水印缺少图片路径"
            if not (0.01 <= self.scale <= 1.0):
                return False, "缩放比例必须在 0.01-1.0 之间"
        elif self.mode == "text":
            if not self.text:
                return False, "文字水印缺少文字内容"
            if self.font_size <= 0:
                return False, "字体大小必须大于 0"
        else:
            return False, f"不支持的水印模式: {self.mode}"

        return True, ""


class WatermarkEngine:
    """水印引擎 — 生成 FFmpeg 水印滤镜."""

    @staticmethod
    def calc_position(
        position: str,
        output_width: int,
        output_height: int,
        wm_width: int,
        wm_height: int,
        margin_x: int,
        margin_y: int,
    ) -> tuple[int, int]:
        """根据9宫格位置计算水印坐标 (x, y).

        坐标系：左上角为 (0, 0)
        """
        if position == "top_left":
            return margin_x, margin_y
        elif position == "top_center":
            return (output_width - wm_width) // 2, margin_y
        elif position == "top_right":
            return output_width - wm_width - margin_x, margin_y
        elif position == "center_left":
            return margin_x, (output_height - wm_height) // 2
        elif position == "center":
            return (output_width - wm_width) // 2, (output_height - wm_height) // 2
        elif position == "center_right":
            return output_width - wm_width - margin_x, (output_height - wm_height) // 2
        elif position == "bottom_left":
            return margin_x, output_height - wm_height - margin_y
        elif position == "bottom_center":
            return (output_width - wm_width) // 2, output_height - wm_height - margin_y
        elif position == "bottom_right":
            return output_width - wm_width - margin_x, output_height - wm_height - margin_y
        else:
            # 默认右下角
            return output_width - wm_width - margin_x, output_height - wm_height - margin_y

    @staticmethod
    def calc_scroll_x(position: str, output_width: int, wm_width: int, speed: int) -> str:
        """生成滚动水印的 x 坐标表达式.

        从右向左滚动（跑马灯效果）
        """
        # x 从 W 到 -wm_width，整个宽度 + wm_width 的距离
        # 使用 overlay 的 enable 表达式
        # x = 'W - (t * speed)' → 不对，应该是持续滚动
        # 标准跑马灯：x = -w + (t * speed) % (W + w)
        # 但 FFmpeg overlay 支持表达式
        return f"mod({output_width}-mod({speed}*t\\,{output_width}+{wm_width})"

    @staticmethod
    def build_image_watermark_filter(
        input_video_label: str,
        wm_image_path: str,
        output_width: int,
        output_height: int,
        output_label: str,
        config: WatermarkConfig,
    ) -> tuple[str, list[str]]:
        """构建图片水印滤镜链.

        Args:
            input_video_label: 输入视频标签，如 "[final_video]"
            wm_image_path: 水印图片本地路径
            output_width: 输出视频宽度
            output_height: 输出视频高度
            output_label: 输出标签
            config: 水印配置

        Returns:
            (filter_complex_str, input_args_list)
            input_args 是 ["-i", wm_image_path] 格式
        """
        # 计算水印尺寸（按输出宽度比例缩放）
        wm_width = int(output_width * config.scale)
        wm_height = -1  # 保持比例
        wm_filter = f"scale={wm_width}:{wm_height}"

        # 透明度处理
        if config.opacity < 1.0:
            wm_filter += f",format=rgba,colorchannelmixer=aa={config.opacity}"

        # 水印预处理标签
        wm_pre_label = "[wm_scaled]"

        # 计算位置
        x, y = WatermarkEngine.calc_position(
            config.position,
            output_width,
            output_height,
            wm_width,
            wm_width,  # 高度未知，先用宽度估算
            config.margin_x,
            config.margin_y,
        )

        # 滚动水印
        if config.scroll:
            # 从右向左滚动：x = W - (t * speed) mod (W + wm_w)
            # 使用 overlay 表达式
            x_expr = f"{output_width}-mod({config.scroll_speed}*t\\,{output_width}+{wm_width}"
            y_expr = str(y)
            overlay_expr = f"x={x_expr}:y={y_expr}"
        else:
            overlay_expr = f"x={x}:y={y}"

        # 构建滤镜
        # 先缩放水印图
        filter_parts = [
            f"[1:v]{wm_filter}{wm_pre_label}",
            f"{input_video_label}{wm_pre_label}overlay={overlay_expr}{output_label}",
        ]

        filter_complex = ";".join(filter_parts)
        input_args = ["-i", wm_image_path]

        return filter_complex, input_args

    @staticmethod
    def build_text_watermark_filter(
        input_video_label: str,
        output_label: str,
        config: WatermarkConfig,
        output_width: int,
        output_height: int,
    ) -> str:
        """构建文字水印滤镜（drawtext）.

        Args:
            input_video_label: 输入视频标签
            output_label: 输出标签
            config: 水印配置
            output_width: 输出宽度
            output_height: 输出高度

        Returns:
            FFmpeg filter 字符串
        """
        # 转义文字中的特殊字符
        text = config.text.replace(":", "\\:").replace("'", "\\'")

        # 字体配置
        font_config = []
        if config.font_path:
            font_path_escaped = config.font_path.replace(":", "\\:").replace("'", "\\'")
            font_config.append(f"fontfile='{font_path_escaped}'")
        font_config.append(f"fontsize={config.font_size}")
        font_config.append(f"fontcolor={config.font_color}@{config.opacity}")

        # 估算文字宽高（粗略估算，用于位置计算）
        # 每个汉字约等于 font_size 宽高
        approx_w = len(config.text) * config.font_size
        approx_h = config.font_size

        # 位置计算
        x, y = WatermarkEngine.calc_position(
            config.position,
            output_width,
            output_height,
            approx_w,
            approx_h,
            config.margin_x,
            config.margin_y,
        )

        # 滚动水印
        if config.scroll:
            x_expr = f"w-mod({config.scroll_speed}*t\\,W+w)"
            pos_config = [f"x={x_expr}", f"y={y}"]
        else:
            pos_config = [f"x={x}", f"y={y}"]

        # 组装 drawtext
        drawtext_parts = [f"text='{text}'"] + font_config + pos_config
        drawtext = "drawtext=" + ":".join(drawtext_parts)

        return f"{input_video_label}{drawtext}{output_label}"