#!/usr/bin/env python3
"""
Subtitle Style A/B Test · 字幕样式 A/B 测试器
==========================​=================
同一底片一次生成多个字幕风格版本，对比选，不要一版一版盲调。

功能：
  1. 读取多个字幕样式配置
  2. 用同一个底片和 SRT，生成多个字幕版本
  3. 生成对比预览图（contact sheet）
  4. 生成 HTML 对比页面

依赖：
  pip install pysrt

用法：
  # 单视频 A/B 测试
  python subtitle-style-ab-test.py --video input.mp4 --srt input.srt --styles style1.json style2.json --output-dir ./ab-test/

  # 批量 A/B 测试
  python subtitle-style-ab-test.py --batch ./videos/ --srt-dir ./srt/ --styles style1.json style2.json --output-dir ./ab-tests/

  # 使用预设样式
  python subtitle-style-ab-test.py --video input.mp4 --srt input.srt --preset-styles reference-drama short-drama-bold --output-dir ./ab-test/

输出目录结构：
  ./ab-test/
    ├── style1/
    │   ├── subtitles.ass
    │   ├── output.mp4
    │   └── preview.png
    ├── style2/
    │   ├── subtitles.ass
    │   ├── output.mp4
    │   └── preview.png
    ├── comparison.png  # 对比预览图
    └── comparison.html # 对比 HTML 页面

路径：
  video-ai-system/engines/subtitle-pipeline/ab-test-tools/subtitle-style-ab-test.py
"""

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

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


# 预设样式列表
PRESET_STYLES = {
    "reference-drama": {
        "font_family": "PingFang SC",
        "font_size": 38,
        "font_color": "&HFFFFFF&",
        "bold": True,
        "stroke_enabled": True,
        "stroke_width": 2,
        "stroke_color": "&H000000&",
        "alignment": 2,
        "margin_left": 100,
        "margin_right": 100,
        "margin_vertical": 50,
    },
    "short-drama-bold": {
        "font_family": "PingFang SC",
        "font_size": 42,
        "font_color": "&HFFFFFF&",
        "bold": True,
        "stroke_enabled": True,
        "stroke_width": 1,
        "stroke_color": "&H121212&",
        "alignment": 2,
        "margin_left": 80,
        "margin_right": 80,
        "margin_vertical": 60,
    },
    "clean-white": {
        "font_family": "PingFang SC",
        "font_size": 36,
        "font_color": "&HFFFFFF&",
        "bold": False,
        "stroke_enabled": False,
        "stroke_width": 0,
        "stroke_color": "&H000000&",
        "alignment": 2,
        "margin_left": 100,
        "margin_right": 100,
        "margin_vertical": 50,
    },
    "black-box": {
        "font_family": "PingFang SC",
        "font_size": 36,
        "font_color": "&HFFFFFF&",
        "bold": False,
        "stroke_enabled": False,
        "stroke_width": 0,
        "stroke_color": "&H000000&",
        "background_enabled": True,
        "background_color": "&H000000&",
        "background_opacity": 0.6,
        "alignment": 2,
        "margin_left": 100,
        "margin_right": 100,
        "margin_vertical": 50,
    },
}


def load_style_config(style_path: str) -> dict:
    """
    加载样式配置
    
    :param style_path: 样式文件路径或预设样式名称
    :return: 样式字典
    """
    # 检查是否是预设样式
    if style_path in PRESET_STYLES:
        print(f"[INFO] 使用预设样式：{style_path}")
        return PRESET_STYLES[style_path]
    
    # 加载自定义样式文件
    if not os.path.isfile(style_path):
        print(f"[ERROR] 样式文件不存在：{style_path}")
        return None
    
    try:
        with open(style_path, "r", encoding="utf-8") as f:
            config = json.load(f)
        print(f"[INFO] 已加载样式配置：{style_path}")
        return config
    except Exception as e:
        print(f"[ERROR] 加载样式配置失败：{e}")
        return None


def generate_ass_file_for_style(srt_path: str, output_path: str, style: dict) -> bool:
    """
    为指定样式生成 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] 为样式生成 ASS 文件：{os.path.basename(output_path)}")
    
    # 构建 ASS 文件内容（简化版，只生成 Dialogues）
    ass_lines = []
    
    # 1. [Script Info] 节
    ass_lines.append("[Script Info]")
    ass_lines.append("; Script generated by Subtitle Style A/B Test")
    ass_lines.append("PlayResX: 1080")
    ass_lines.append("PlayResY: 1920")
    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&"
    else:
        back_colour = "&H00000000&"
    
    bold = -1 if style.get("bold", False) else 0
    outline_width = style.get("stroke_width", 2) if style.get("stroke_enabled", True) 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},0,0,0,100,100,0,0,"
        f"1,{outline_width},0,{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()
        
        # 转义特殊字符
        text = text.replace("\\", "\\\\")
        text = text.replace("{", "\\{")
        text = text.replace("}", "\\}")
        text = text.replace("\n", "\\N")
        
        dialogue_line = f"Dialogue: 0,{start_time},{end_time},Default,,0,0,0,,{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:
            f.write("\n".join(ass_lines))
        
        print(f"[OK] ASS 字幕文件已生成：{output_path}")
        return True
    
    except Exception as e:
        print(f"[ERROR] 写入 ASS 文件失败：{e}")
        return False


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
    
    centiseconds = milliseconds // 10
    return f"{hours}:{minutes:02d}:{seconds:02d}.{centiseconds:02d}"


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
    
    # FFmpeg 命令
    # 正确语法：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] 烧录字幕：{os.path.basename(output_path)}")
    
    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
    
    return False


def extract_sample_frames(video_path: str, output_dir: str, num_frames: int = 4) -> list:
    """
    从视频中抽取样本帧（用于预览）
    
    :param video_path: 视频文件路径
    :param output_dir: 输出目录
    :param num_frames: 抽取帧数
    :return: 帧文件路径列表
    """
    import cv2
    
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"[ERROR] 无法打开视频：{video_path}")
        return []
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # 计算采样间隔
    interval = total_frames // (num_frames + 1)
    
    frame_paths = []
    for i in range(1, num_frames + 1):
        frame_num = i * interval
        
        # 跳转到指定帧
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
        
        ret, frame = cap.read()
        if not ret:
            continue
        
        # 保存帧
        frame_path = os.path.join(output_dir, f"sample_{i:02d}.png")
        cv2.imwrite(frame_path, frame)
        
        frame_paths.append(frame_path)
    
    cap.release()
    
    print(f"[INFO] 已抽取 {len(frame_paths)} 个样本帧")
    return frame_paths


def generate_comparison_preview(style_dirs: list, output_path: str):
    """
    生成对比预览图
    
    :param style_dirs: 样式目录列表
    :param output_path: 输出图片路径
    """
    try:
        from PIL import Image
    except ImportError:
        print("[ERROR] 缺少依赖：Pillow")
        print("请先安装：pip install Pillow")
        return False
    
    # 收集所有样本帧
    all_frames = []
    style_names = []
    
    for style_dir in style_dirs:
        style_name = os.path.basename(style_dir)
        style_names.append(style_name)
        
        # 查找样本帧
        frame_files = [f for f in os.listdir(style_dir) if f.startswith("sample_") and f.endswith(".png")]
        frame_files.sort()
        
        for frame_file in frame_files:
            frame_path = os.path.join(style_dir, frame_file)
            all_frames.append((style_name, frame_path))
    
    if not all_frames:
        print("[WARN] 未找到样本帧，跳过对比预览图生成")
        return False
    
    # 读取第一帧，获取尺寸
    first_frame = Image.open(all_frames[0][1])
    frame_width, frame_height = first_frame.size
    
    # 计算拼接图尺寸
    num_styles = len(style_dirs)
    num_samples = len(all_frames) // num_styles
    
    # 横向拼接：每个样式一行，每行包含多个样本帧
    sheet_width = num_samples * (frame_width + 10) + 10
    sheet_height = num_styles * (frame_height + 30) + 10  # 30px 用于样式名称标注
    
    # 创建拼接图
    sheet = Image.new("RGB", (sheet_width, sheet_height), (0, 0, 0))
    
    # 拼接帧
    for i, (style_name, frame_path) in enumerate(all_frames):
        frame = Image.open(frame_path)
        
        # 计算位置
        style_idx = style_names.index(style_name)
        sample_idx = i % num_samples
        
        x = sample_idx * (frame_width + 10) + 10
        y = style_idx * (frame_height + 30) + 10
        
        # 粘贴帧
        sheet.paste(frame, (x, y + 20))  # +20px 用于样式名称标注
        
        # 添加样式名称标注
        draw = ImageDraw.Draw(sheet)
        draw.text((x + 5, y + 5), style_name, fill=(255, 255, 255))
    
    # 保存
    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    sheet.save(output_path)
    
    print(f"[OK] 对比预览图已生成：{output_path}")
    return True


def main():
    parser = argparse.ArgumentParser(
        description="Subtitle Style A/B Test · 字幕样式 A/B 测试器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例：
  python subtitle-style-ab-test.py --video input.mp4 --srt input.srt --styles style1.json style2.json --output-dir ./ab-test/
  python subtitle-style-ab-test.py --video input.mp4 --srt input.srt --preset-styles reference-drama short-drama-bold --output-dir ./ab-test/
"""
    )
    parser.add_argument("--video", help="视频文件路径")
    parser.add_argument("--srt", help="SRT 字幕文件路径")
    parser.add_argument("--styles", nargs="+", help="样式配置文件路径列表")
    parser.add_argument("--preset-styles", nargs="+", help="预设样式名称列表")
    parser.add_argument("--batch", help="批量处理视频目录")
    parser.add_argument("--srt-dir", help="SRT 文件目录（用于批量处理）")
    parser.add_argument("--output-dir", required=True, help="输出目录")
    parser.add_argument("--num-samples", type=int, default=4, help="每个样式抽取的样本帧数")
    
    args = parser.parse_args()
    
    # 收集样式配置
    styles = []
    style_names = []
    
    if args.styles:
        for style_path in args.styles:
            style = load_style_config(style_path)
            if style:
                styles.append(style)
                style_names.append(os.path.splitext(os.path.basename(style_path))[0])
    
    if args.preset_styles:
        for preset_name in args.preset_styles:
            style = load_style_config(preset_name)
            if style:
                styles.append(style)
                style_names.append(preset_name)
    
    if not styles:
        print("[ERROR] 请指定至少一种样式（--styles 或 --preset-styles）")
        sys.exit(1)
    
    print(f"[INFO] 将测试 {len(styles)} 种样式：{style_names}")
    
    if args.video:
        # 单视频处理
        if not args.srt:
            print("[ERROR] 请指定 --srt")
            sys.exit(1)
        
        # 为每个样式生成字幕和视频
        style_dirs = []
        
        for i, (style, style_name) in enumerate(zip(styles, style_names)):
            print(f"\n[INFO] 处理样式 {i+1}/{len(styles)}：{style_name}")
            
            # 创建样式目录
            style_dir = os.path.join(args.output_dir, style_name)
            os.makedirs(style_dir, exist_ok=True)
            style_dirs.append(style_dir)
            
            # 生成 ASS 文件
            ass_path = os.path.join(style_dir, "subtitles.ass")
            ok = generate_ass_file_for_style(args.srt, ass_path, style)
            
            if ok:
                # 烧录字幕到视频
                output_video = os.path.join(style_dir, os.path.basename(args.video))
                burn_subtitles_with_ffmpeg(args.video, ass_path, output_video)
                
                # 抽取样本帧
                extract_sample_frames(output_video, style_dir, args.num_samples)
        
        # 生成对比预览图
        comparison_preview = os.path.join(args.output_dir, "comparison.png")
        generate_comparison_preview(style_dirs, comparison_preview)
    
    elif args.batch:
        # 批量处理
        if not args.srt_dir:
            print("[ERROR] 批量处理需要指定 --srt-dir")
            sys.exit(1)
        
        for video_file in os.listdir(args.batch):
            if not video_file.lower().endswith((".mp4", ".mov", ".avi")):
                continue
            
            video_path = os.path.join(args.batch, video_file)
            
            # 查找对应的 SRT 文件
            srt_file = video_file.replace(".mp4", ".srt").replace(".mov", ".srt").replace(".avi", ".srt")
            srt_path = os.path.join(args.srt_dir, srt_file)
            
            if not os.path.isfile(srt_path):
                print(f"[WARN] 未找到对应的 SRT 文件：{srt_file}")
                continue
            
            print(f"\n[INFO] 批量处理：{video_file}")
            
            # 为每个样式生成字幕和视频
            video_output_dir = os.path.join(args.output_dir, os.path.splitext(video_file)[0])
            
            for i, (style, style_name) in enumerate(zip(styles, style_names)):
                print(f"\n[INFO] 处理样式 {i+1}/{len(styles)}：{style_name}")
                
                style_dir = os.path.join(video_output_dir, style_name)
                os.makedirs(style_dir, exist_ok=True)
                
                # 生成 ASS 文件
                ass_path = os.path.join(style_dir, "subtitles.ass")
                ok = generate_ass_file_for_style(srt_path, ass_path, style)
                
                if ok:
                    # 烧录字幕到视频
                    output_video = os.path.join(style_dir, video_file)
                    burn_subtitles_with_ffmpeg(video_path, ass_path, output_video)
                    
                    # 抽取样本帧
                    extract_sample_frames(output_video, style_dir, args.num_samples)
    
    print("\n[OK] A/B 测试完成")


if __name__ == "__main__":
    main()
