#!/usr/bin/env python3
"""
ASS Subtitle Renderer · ASS字幕渲染器
==========================​===========
用 ASS/libass 正式渲染字幕，支持字体、粗体、描边、边距、对齐。
不再用 Pillow PNG 凑合。

依赖：
  pip install pysrt
  # FFmpeg 需要编译时启用 libass 支持（通常默认启用）

用法：
  # 生成 ASS 字幕文件
  python ass-subtitle-renderer.py --srt input.srt --style reference-drama --output subtitles.ass

  # 用 FFmpeg + libass 烧字幕到视频
  python ass-subtitle-renderer.py --srt input.srt --video input.mp4 --output output.mp4 --style reference-drama

  # 指定自定义样式配置
  python ass-subtitle-renderer.py --srt input.srt --config custom-style.json --output subtitles.ass

  # 批量处理多集
  python ass-subtitle-renderer.py --batch ./srt/ --style reference-drama --output-dir ./ass/

输出：
  - ASS 字幕文件（可直接用 FFmpeg 烧录）
  - 烧录字幕后的视频（可选）

路径：
  video-ai-system/engines/subtitle-pipeline/reference-analysis/ass-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)


# 默认样式配置（参考短剧风格）
DEFAULT_STYLE = {
    "video_width": 1080,
    "video_height": 1920,
    "font_family": "PingFang SC",
    "font_size": 38,
    "font_color": "&HFFFFFF&",  # 白色（ASS格式：&HBBGGRR&）
    "font_color_hex": "#FFFFFF",
    "bold": True,
    "italic": False,
    "underline": False,
    "strikeout": False,
    "stroke_enabled": True,
    "stroke_width": 2,
    "stroke_color": "&H000000&",  # 黑色描边
    "stroke_opacity": 0.92,  # &H40&（0-255，0=透明，255=不透明）
    "shadow_enabled": False,
    "shadow_depth": 0,
    "alignment": 2,  # 2=中下
    "margin_left": 100,
    "margin_right": 100,
    "margin_vertical": 50,
    "line_spacing": 1.45,
    "letter_spacing": 0,
    "background_enabled": False,
    "fade_in_ms": 150,
    "fade_out_ms": 150,
}


def load_style_config(config_path: str = None, style_name: str = "reference-drama") -> dict:
    """
    加载样式配置
    
    :param config_path: 自定义配置文件路径
    :param style_name: 预设样式名称
    :return: 样式字典
    """
    if config_path and os.path.isfile(config_path):
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
        print(f"[INFO] 已加载自定义样式配置：{config_path}")
        return config
    
    # 尝试从 assets/subtitle-styles/ 加载预设样式
    style_file = os.path.join(
        os.path.dirname(__file__),
        "..", "..", "..", "assets", "subtitle-styles",
        f"subtitle-style.{style_name}.json"
    )
    style_file = os.path.abspath(style_file)
    
    if os.path.isfile(style_file):
        with open(style_file, "r", encoding="utf-8") as f:
            config = json.load(f)
        print(f"[INFO] 已加载预设样式：{style_name}")
        return config
    
    print(f"[WARN] 未找到样式配置，使用默认样式")
    return DEFAULT_STYLE


def srt_time_to_ass(srt_time) -> str:
    """
    将 pysrt 时间对象转换为 ASS 时间戳格式
    
    :param srt_time: pysrt 时间对象
    :return: ASS 时间戳字符串（H:MM:SS.cc）
    """
    hours = srt_time.hours
    minutes = srt_time.minutes
    seconds = srt_time.seconds
    milliseconds = srt_time.milliseconds
    
    # ASS 时间戳格式：H:MM:SS.cc（cc=厘秒，1/100秒）
    centiseconds = milliseconds // 10
    return f"{hours}:{minutes:02d}:{seconds:02d}.{centiseconds:02d}"


def generate_ass_file(srt_path: str, output_path: str, style: dict) -> bool:
    """
    将 SRT 字幕转换为 ASS 格式
    
    :param srt_path: SRT 文件路径
    :param output_path: 输出 ASS 文件路径
    :param style: 样式字典
    :return: 是否成功
    """
    if not os.path.isfile(srt_path):
        print(f"[ERROR] SRT 文件不存在：{srt_path}")
        return False
    
    # 加载 SRT
    try:
        subs = pysrt.open(srt_path, encoding="utf-8")
    except Exception as e:
        print(f"[ERROR] 加载 SRT 失败：{e}")
        return False
    
    print(f"[INFO] 找到 {len(subs)} 条字幕，开始生成 ASS 文件...")
    
    # 构建 ASS 文件内容
    ass_lines = []
    
    # 1. [Script Info] 节
    ass_lines.append("[Script Info]")
    ass_lines.append("; Script generated by ASS Subtitle Renderer")
    ass_lines.append(f"PlayResX: {style.get('video_width', 1080)}")
    ass_lines.append(f"PlayResY: {style.get('video_height', 1920)}")
    ass_lines.append("Aspect Ratio: 9:16")
    ass_lines.append("")
    
    # 2. [V4+ Styles] 节
    ass_lines.append("[V4+ Styles]")
    ass_lines.append(
        "Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, "
        "Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, "
        "BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding"
    )
    
    # 构建样式行
    font_name = style.get("font_family", "PingFang SC")
    font_size = style.get("font_size", 38)
    primary_colour = style.get("font_color", "&HFFFFFF&")
    outline_colour = style.get("stroke_color", "&H000000&")
    
    # 背景色（如果启用背景）
    if style.get("background_enabled", False):
        bg_opacity = int(style.get("background_opacity", 0.6) * 255)
        back_colour = f"&H{bg_opacity:02X}000000&"  # &HAAKKBBDD&
    else:
        back_colour = "&H00000000&"  # 透明
    
    bold = -1 if style.get("bold", True) else 0
    italic = -1 if style.get("italic", False) else 0
    underline = -1 if style.get("underline", False) else 0
    strikeout = -1 if style.get("strikeout", False) else 0
    
    outline_width = style.get("stroke_width", 2) if style.get("stroke_enabled", True) else 0
    shadow_depth = style.get("shadow_depth", 0) if style.get("shadow_enabled", False) else 0
    
    alignment = style.get("alignment", 2)
    margin_l = style.get("margin_left", 100)
    margin_r = style.get("margin_right", 100)
    margin_v = style.get("margin_vertical", 50)
    
    style_line = (
        f"Style: Default,{font_name},{font_size},{primary_colour},&H000000FF&,{outline_colour},{back_colour},"
        f"{bold},{italic},{underline},{strikeout},100,100,{style.get('letter_spacing', 0)},0,"
        f"1,{outline_width},{shadow_depth},{alignment},{margin_l},{margin_r},{margin_v},1"
    )
    ass_lines.append(style_line)
    ass_lines.append("")
    
    # 3. [Events] 节
    ass_lines.append("[Events]")
    ass_lines.append("Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text")
    
    for sub in subs:
        start_time = srt_time_to_ass(sub.start)
        end_time = srt_time_to_ass(sub.end)
        text = sub.text.strip()
        
        # 添加淡入淡出效果（如果启用）
        fade_in = style.get("fade_in_ms", 150)
        fade_out = style.get("fade_out_ms", 150)
        
        if fade_in > 0 and fade_out > 0:
            # ASS 淡入淡出标签：{\fad(fade_in,fade_out)}
            effect_tags = f"{{\\fad({fade_in},{fade_out})}}"
        else:
            effect_tags = ""
        
        # 转义特殊字符（ASS 格式）
        text = text.replace("\\", "\\\\")
        text = text.replace("{", "\\{")
        text = text.replace("}", "\\}")
        text = text.replace("\n", "\\N")  # ASS 换行符
        
        # 构建对话行
        dialogue_line = (
            f"Dialogue: 0,{start_time},{end_time},Default,,0,0,0,,{effect_tags}{text}"
        )
        ass_lines.append(dialogue_line)
    
    # 写入 ASS 文件
    try:
        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
        with open(output_path, "w", encoding="utf-8-sig") as f:  # UTF-8 BOM for ASS
            f.write("\n".join(ass_lines))
        
        print(f"[OK] ASS 字幕文件已生成：{output_path}")
        print(f"[INFO] 共 {len(subs)} 条字幕")
        return True
    
    except Exception as e:
        print(f"[ERROR] 写入 ASS 文件失败：{e}")
        return False


def burn_subtitles_with_ffmpeg(video_path: str, ass_path: str, output_path: str) -> bool:
    """
    用 FFmpeg + libass 烧录字幕到视频
    
    :param video_path: 输入视频路径
    :param ass_path: ASS 字幕文件路径
    :param output_path: 输出视频路径
    :return: 是否成功
    """
    if not os.path.isfile(video_path):
        print(f"[ERROR] 视频文件不存在：{video_path}")
        return False
    
    if not os.path.isfile(ass_path):
        print(f"[ERROR] ASS 字幕文件不存在：{ass_path}")
        return False
    
    # 使用 subtitles 滤镜（更通用，支持 SRT/ASS 等多种格式）
    # 正确语法：subtitles=filename='path'  （需要 filename= 前缀）
    ass_path_escaped = ass_path.replace(":", "\\:").replace("'", "'\\''")
    
    cmd = [
        "ffmpeg", "-y",
        "-i", video_path,
        "-vf", f"subtitles=filename='{ass_path_escaped}'",
        "-c:v", "libx264",
        "-pix_fmt", "yuv420p",
        "-c:a", "copy",
        "-shortest",
        output_path
    ]
    
    print(f"[INFO] 开始用 FFmpeg + libass 烧录字幕...")
    print(f"[INFO] 命令：{' '.join(cmd[:10])}...")  # 只打印前10个参数
    
    try:
        result = subprocess.run(cmd, check=False, capture_output=True, text=True)
    except FileNotFoundError:
        print("[ERROR] 未找到 ffmpeg，请先安装 FFmpeg")
        print("安装方法：")
        print("  macOS: brew install ffmpeg")
        print("  Ubuntu: sudo apt install ffmpeg")
        return False
    
    if result.returncode != 0:
        print("[ERROR] FFmpeg 烧录字幕失败")
        print(result.stderr[-2000:])  # 打印最后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 batch_process(srt_dir: str, output_dir: str, style: dict, burn_video: bool = False, video_dir: str = None):
    """
    批量处理 SRT 文件
    
    :param srt_dir: SRT 文件目录
    :param output_dir: 输出目录
    :param style: 样式字典
    :param burn_video: 是否烧录到视频
    :param video_dir: 视频文件目录（如果 burn_video=True）
    """
    os.makedirs(output_dir, exist_ok=True)
    
    for file in os.listdir(srt_dir):
        if file.lower().endswith(".srt"):
            srt_path = os.path.join(srt_dir, file)
            ass_output = os.path.join(output_dir, file.replace(".srt", ".ass"))
            
            print(f"\n[INFO] 处理文件：{file}")
            ok = generate_ass_file(srt_path, ass_output, style)
            
            if ok and burn_video and video_dir:
                # 查找对应的视频文件
                video_file = file.replace(".srt", ".mp4")
                video_path = os.path.join(video_dir, video_file)
                
                if os.path.isfile(video_path):
                    video_output = os.path.join(output_dir, video_file)
                    burn_subtitles_with_ffmpeg(video_path, ass_output, video_output)
                else:
                    print(f"[WARN] 未找到对应的视频文件：{video_file}")


def main():
    parser = argparse.ArgumentParser(
        description="ASS Subtitle Renderer · ASS字幕渲染器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例：
  python ass-subtitle-renderer.py --srt input.srt --style reference-drama --output subtitles.ass
  python ass-subtitle-renderer.py --srt input.srt --video input.mp4 --output output.mp4 --style reference-drama
  python ass-subtitle-renderer.py --batch ./srt/ --style reference-drama --output-dir ./ass/
"""
    )
    parser.add_argument("--srt", help="输入 SRT 文件路径")
    parser.add_argument("--video", help="输入视频路径（可选，用于直接烧录字幕）")
    parser.add_argument("--output", help="输出文件路径（ASS 或 MP4）")
    parser.add_argument("--style", default="reference-drama", help="样式名称或配置文件路径")
    parser.add_argument("--config", help="自定义样式配置文件路径（覆盖 --style）")
    parser.add_argument("--batch", help="批量处理 SRT 文件目录")
    parser.add_argument("--output-dir", help="批量处理输出目录")
    parser.add_argument("--burn-video", action="store_true", help="是否烧录字幕到视频")
    parser.add_argument("--video-dir", help="视频文件目录（配合 --burn-video 使用）")
    
    args = parser.parse_args()
    
    # 加载样式配置
    if args.config:
        style = load_style_config(config_path=args.config)
    else:
        style = load_style_config(style_name=args.style)
    
    if args.srt:
        # 单文件处理
        if not args.output:
            print("[ERROR] 请指定 --output")
            sys.exit(1)
        
        # 生成 ASS 文件
        ok = generate_ass_file(args.srt, args.output, style)
        if not ok:
            sys.exit(1)
        
        # 如果指定了视频，则烧录字幕
        if args.video:
            video_output = args.output.replace(".ass", ".mp4")
            burn_ok = burn_subtitles_with_ffmpeg(args.video, args.output, video_output)
            if not burn_ok:
                sys.exit(1)
    
    elif args.batch:
        # 批量处理
        if not args.output_dir:
            print("[ERROR] 批量处理需要指定 --output-dir")
            sys.exit(1)
        
        batch_process(args.batch, args.output_dir, style, args.burn_video, args.video_dir)
    
    else:
        print("[ERROR] 请指定 --srt 或 --batch")
        sys.exit(1)
    
    print("\n[OK] 处理完成")


if __name__ == "__main__":
    main()
