#!/usr/bin/env python3
"""
Subtitle Renderer · 字幕渲染引擎
=============================
将 SRT 字幕文件渲染为 PNG 序列，再通过 FFmpeg 合成到视频中。

依赖：
  pip install Pillow pysrt

用法：
  # 基本用法（SRT → PNG 序列）
  python subtitle-renderer.py --srt input.srt --output-dir ./subtitles-png/

  # 指定视频尺寸（PNG 宽度匹配视频）
  python subtitle-renderer.py --srt input.srt --output-dir ./subtitles-png/ --width 1080 --height 1920

  # 渲染后直接合成到视频
  python subtitle-renderer.py --srt input.srt --video input.mp4 --output output.mp4

  # 自定义字幕样式
  python subtitle-renderer.py --srt input.srt --font-size 48 --font-color white --style clean-white --position bottom

  # 作为模块导入
  from subtitle_renderer import render_subtitles
  render_subtitles("input.srt", "./subtitles-png/")

字幕样式配置：
  --font-size    : 字体大小（默认 36）
  --font-color   : 字体颜色（默认 white）
  --style        : 字幕样式（reference-drama / short-drama-bold / clean-white / black-box，默认 reference-drama）
  --position     : 位置（top / middle / bottom，默认 bottom）
  --margin-bottom : 底部边距（默认 100px）

路径：
  video-ai-system/engines/subtitle-renderer.py
"""

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

try:
    import pysrt
except ImportError:
    print("[ERROR] 缺少依赖：pysrt")
    print("请先安装：pip install pysrt")
    sys.exit(1)

try:
    from PIL import Image, ImageDraw, ImageFont
except ImportError:
    print("[ERROR] 缺少依赖：Pillow")
    print("请先安装：pip install Pillow")
    sys.exit(1)


# 默认字幕样式
DEFAULT_STYLE = {
    "font_size": 38,
    "font_color": "white",
    "bg_enabled": False,
    "position": "bottom",  # top / middle / bottom
    "margin_bottom": 28,
    "margin_horizontal": 60,
    "stroke_color": "black",
    "stroke_width": 1,
    "bold_weight": 2,
    "video_width": 1080,
    "video_height": 1920,
}

SUBTITLE_STYLE_PRESETS = {
    "reference-drama": {
        "font_color": (255, 255, 255, 255),
        "bg_enabled": False,
        "stroke_width": 2,
        "stroke_color": (0, 0, 0, 235),
        "bold_weight": 1,
        "line_box_ratio": 1.45,
    },
    "short-drama-bold": {
        "font_color": (255, 255, 255, 255),
        "bg_enabled": False,
        "stroke_width": 1,
        "stroke_color": (18, 18, 18, 220),
        "bold_weight": 2,
        "line_box_ratio": 1.55,
    },
    "clean-white": {
        "font_color": (255, 255, 255, 255),
        "bg_enabled": False,
        "stroke_width": 0,
        "bold_weight": 1,
        "line_box_ratio": 1.55,
    },
    "black-box": {
        "font_color": (255, 255, 255, 255),
        "bg_enabled": True,
        "bg_fill": (0, 0, 0, 160),
        "stroke_width": 0,
        "bold_weight": 1,
        "line_box_ratio": 1.55,
    },
    "outlined-white": {
        "font_color": (255, 255, 255, 255),
        "bg_enabled": False,
        "stroke_width": 2,
        "stroke_color": (0, 0, 0, 255),
        "bold_weight": 1,
        "line_box_ratio": 1.55,
    },
}

CHINESE_FONT_CANDIDATES = [
    "/System/Library/Fonts/PingFang.ttc",
    "/System/Library/Fonts/Hiragino Sans GB.ttc",
    "/System/Library/Fonts/STHeiti Medium.ttc",
    "/System/Library/Fonts/Supplemental/Songti.ttc",
    "/System/Library/Fonts/Supplemental/Kaiti.ttc",
    "/Library/Fonts/Arial Unicode.ttf",
    "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
]


def load_font(font_size: int):
    for font_path in CHINESE_FONT_CANDIDATES:
        if not os.path.isfile(font_path):
            continue
        try:
            return ImageFont.truetype(font_path, font_size)
        except Exception:
            continue
    return ImageFont.load_default()


def render_subtitle_png(
    text: str,
    output_path: str,
    width: int = 1080,
    height: int = 200,
    style: dict = None
) -> bool:
    """
    渲染单条字幕为 PNG（带背景）

    :param text: 字幕文本
    :param output_path: 输出 PNG 路径
    :param width: PNG 宽度（匹配视频宽度）
    :param height: PNG 高度
    :param style: 字幕样式字典
    :return: 是否成功
    """
    if style is None:
        style = DEFAULT_STYLE

    try:
        # 创建透明背景 PNG
        img = Image.new("RGBA", (width, height), (0, 0, 0, 0))
        draw = ImageDraw.Draw(img)

        # 字体（使用系统字体）
        font_size = style.get("font_size", 36)
        font = load_font(font_size)

        # 计算文本尺寸
        bbox = draw.textbbox((0, 0), text, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]

        # 居中位置
        x = (width - text_width) // 2
        y = (height - text_height) // 2

        # 可选背景。默认纯白字幕不画背景。
        if style.get("bg_enabled", False):
            bg_padding = 20
            bg_x1 = x - bg_padding
            bg_y1 = y - bg_padding
            bg_x2 = x + text_width + bg_padding
            bg_y2 = y + text_height + bg_padding
            draw.rectangle([bg_x1, bg_y1, bg_x2, bg_y2], fill=style.get("bg_fill", (0, 0, 0, 160)))

        # 可选描边。短剧默认使用极细暗边，不使用投影。
        stroke_width = style.get("stroke_width", 2)
        stroke_color = style.get("stroke_color", "black")
        if stroke_width > 0:
            for offset in range(-stroke_width, stroke_width + 1):
                draw.text((x + offset, y), text, font=font, fill=stroke_color)
                draw.text((x, y + offset), text, font=font, fill=stroke_color)

        # 绘制主文本。bold_weight 用多次微偏移模拟加粗，避免依赖某台机器是否有粗体中文字体。
        font_color = style.get("font_color", "white")
        bold_weight = max(1, int(style.get("bold_weight", 1)))
        offsets = [(0, 0)]
        if bold_weight >= 2:
            offsets += [(-1, 0), (1, 0), (0, -1), (0, 1)]
        if bold_weight >= 3:
            offsets += [(-1, -1), (1, -1), (-1, 1), (1, 1)]
        for dx, dy in offsets:
            draw.text((x + dx, y + dy), text, font=font, fill=font_color)

        # 保存
        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
        img.save(output_path, "PNG")

        return True

    except Exception as e:
        print(f"[ERROR] 渲染字幕失败：{e}")
        return False


def render_subtitles(
    srt_path: str,
    output_dir: str,
    style: dict = None
) -> dict:
    """
    渲染 SRT 字幕为 PNG 序列

    :param srt_path: SRT 文件路径
    :param output_dir: 输出目录
    :param style: 字幕样式字典
    :return: {idx: {"png": png_path, "start": start_time, "end": end_time, "text": text}}
    """
    if not os.path.isfile(srt_path):
        print(f"[ERROR] SRT 文件不存在：{srt_path}")
        return {}

    os.makedirs(output_dir, exist_ok=True)

    # 加载 SRT
    subs = pysrt.open(srt_path, encoding="utf-8")
    print(f"[INFO] 找到 {len(subs)} 条字幕，开始渲染 PNG 序列...")

    results = {}
    for sub in subs:
        idx = str(sub.index).zfill(4)
        text = sub.text.strip()
        start_time = sub.start.ordinal  # 毫秒
        end_time = sub.end.ordinal

        # 渲染 PNG
        png_path = os.path.join(output_dir, f"{idx}.png")
        video_width = style.get("video_width", 1080) if style else 1080
        font_size = style.get("font_size", 36) if style else 36
        line_box_ratio = style.get("line_box_ratio", 1.55) if style else 1.55
        png_height = int(font_size * line_box_ratio)

        ok = render_subtitle_png(text, png_path, width=video_width, height=png_height, style=style)
        if ok:
            results[idx] = {
                "png": png_path,
                "start": start_time,
                "end": end_time,
                "text": text
            }

    print(f"[OK] 字幕 PNG 序列渲染完成：{len(results)}/{len(subs)} 成功")
    return results


def burn_subtitles_to_video(
    video_path: str,
    subtitles: dict,
    output_path: str,
    video_width: int = 1080,
    video_height: int = 1920
) -> bool:
    """
    将 PNG 字幕序列合成到视频中（使用 FFmpeg overlay 滤镜）

    :param video_path: 输入视频路径
    :param subtitles: render_subtitles 返回的字典
    :param output_path: 输出视频路径
    :param video_width: 视频宽度
    :param video_height: 视频高度
    :return: 是否成功
    """
    if not subtitles:
        print("[ERROR] 没有字幕数据")
        return False

    if not os.path.isfile(video_path):
        print(f"[ERROR] 视频文件不存在：{video_path}")
        return False

    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)

    ordered_subs = [subtitles[k] for k in sorted(subtitles.keys())]
    cmd = ["ffmpeg", "-y", "-i", video_path]
    for sub in ordered_subs:
        cmd.extend(["-i", sub["png"]])

    margin_bottom = DEFAULT_STYLE["margin_bottom"]
    chains = []
    previous = "[0:v]"

    for i, sub in enumerate(ordered_subs, start=1):
        start_sec = sub["start"] / 1000.0
        end_sec = sub["end"] / 1000.0
        scaled = f"[s{i}]"
        out = f"[v{i}]"
        chains.append(f"[{i}:v]scale={video_width}:-1{scaled}")
        chains.append(
            f"{previous}{scaled}"
            f"overlay=x=0:y=H-h-{margin_bottom}:"
            f"enable='between(t,{start_sec:.3f},{end_sec:.3f})'{out}"
        )
        previous = out

    filter_complex = ";".join(chains)
    cmd.extend([
        "-filter_complex", filter_complex,
        "-map", previous,
        "-map", "0:a?",
        "-c:v", "libx264",
        "-pix_fmt", "yuv420p",
        "-c:a", "copy",
        "-shortest",
        output_path,
    ])

    print("[INFO] 开始合成字幕到视频...")
    try:
        result = subprocess.run(cmd, check=False, capture_output=True, text=True)
    except FileNotFoundError:
        print("[ERROR] 未找到 ffmpeg")
        return False

    if result.returncode != 0:
        print("[ERROR] FFmpeg 合成失败")
        print(result.stderr[-2000:])
        return False

    if os.path.isfile(output_path) and os.path.getsize(output_path) > 0:
        print(f"[OK] 字幕已合成到视频：{output_path}")
        return True

    print(f"[ERROR] 输出视频未生成：{output_path}")
    return False


def main():
    parser = argparse.ArgumentParser(
        description="Subtitle Renderer · 字幕渲染引擎",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例：
  python subtitle-renderer.py --srt input.srt --output-dir ./subtitles-png/
  python subtitle-renderer.py --srt input.srt --video input.mp4 --output output.mp4
  python subtitle-renderer.py --srt input.srt --font-size 48 --position bottom
"""
    )
    parser.add_argument("--srt", required=True, help="SRT 字幕文件路径")
    parser.add_argument("--output-dir", help="PNG 序列输出目录")
    parser.add_argument("--video", help="输入视频路径（可选，用于直接合成）")
    parser.add_argument("--output", help="输出视频路径（配合 --video 使用）")
    parser.add_argument("--font-size", type=int, default=36, help="字体大小")
    parser.add_argument("--font-color", default="white", help="字体颜色")
    parser.add_argument("--style", default="reference-drama", choices=sorted(SUBTITLE_STYLE_PRESETS.keys()), help="字幕样式")
    parser.add_argument("--margin-bottom", type=int, default=None, help="字幕距离画面底部的像素边距")
    parser.add_argument("--position", default="bottom", choices=["top", "middle", "bottom"], help="字幕位置")
    parser.add_argument("--video-width", type=int, default=1080, help="视频宽度")
    parser.add_argument("--video-height", type=int, default=1920, help="视频高度")

    args = parser.parse_args()

    # 构建样式字典
    style = {
        **SUBTITLE_STYLE_PRESETS[args.style],
        "font_size": args.font_size,
        "font_color": args.font_color,
        "position": args.position,
        "video_width": args.video_width,
        "video_height": args.video_height,
    }
    if args.margin_bottom is not None:
        style["margin_bottom"] = args.margin_bottom

    # 渲染 PNG 序列
    output_dir = args.output_dir or "./subtitles-png/"
    subtitles = render_subtitles(args.srt, output_dir, style)

    if not subtitles:
        sys.exit(1)

    # 如果指定了视频，则合成
    if args.video and args.output:
        ok = burn_subtitles_to_video(args.video, subtitles, args.output, args.video_width, args.video_height)
        if not ok:
            sys.exit(1)

    print(f"\n[OK] 字幕 PNG 序列已生成：{output_dir}")
    print(f"[INFO] 共 {len(subtitles)} 条字幕")
    sys.exit(0)


if __name__ == "__main__":
    main()