#!/usr/bin/env python3
"""
Subtitle Preview Contact Sheet · 字幕预览检查图生成器
==========================​====================
每次烧字幕后自动抽帧出检查图。

功能：
  1. 从视频中按字幕时间点抽帧
  2. 生成检查图（contact sheet）
  3. 可选：在检查图上标注字幕框、问题区域
  4. 生成 HTML 预览页面

依赖：
  pip install opencv-python pillow numpy

用法：
  # 从视频抽帧生成检查图
  python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --output preview.png

  # 生成 HTML 预览页面
  python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --generate-html --output preview.html

  # 批量处理多集
  python subtitle-preview-contact-sheet.py --batch ./videos/ --subtitle-dir ./ass/ --output-dir ./previews/

  # 自定义抽帧参数
  python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --max-frames 20 --cols 4

输出：
  - preview.png: 检查图（多帧拼接）
  - preview.html: HTML 预览页面（可选）
  - frames/: 抽出的关键帧（可选）

路径：
  video-ai-system/engines/subtitle-pipeline/preview-tools/subtitle-preview-contact-sheet.py
"""

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

try:
    import cv2
    import numpy as np
    from PIL import Image, ImageDraw, ImageFont
except ImportError as e:
    print(f"[ERROR] 缺少依赖：{e}")
    print("请先安装：pip install opencv-python pillow numpy")
    sys.exit(1)


def extract_frames_by_subtitle(video_path: str, ass_path: str, output_dir: str, max_frames: int = 20) -> list:
    """
    按字幕时间点从视频中抽帧
    
    :param video_path: 视频文件路径
    :param ass_path: ASS 字幕文件路径
    :param output_dir: 输出目录
    :param max_frames: 最大抽帧数量
    :return: 抽出的帧文件路径列表
    """
    # 打开视频
    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))
    
    print(f"[INFO] 视频信息：{total_frames} 帧，{fps} FPS")
    
    # 读取 ASS 字幕文件，获取字幕时间点
    subtitle_times = []
    try:
        with open(ass_path, "r", encoding="utf-8-sig") as f:
            lines = f.readlines()
        
        for line in lines:
            if line.startswith("Dialogue:"):
                # 解析 ASS 对话行
                parts = line.split(",")
                if len(parts) >= 10:
                    start_time = parts[1].strip()  # 开始时间
                    end_time = parts[2].strip()    # 结束时间
                    
                    # 转换为秒
                    start_seconds = ass_time_to_seconds(start_time)
                    end_seconds = ass_time_to_seconds(end_time)
                    
                    # 取字幕中间时间点作为抽帧时间
                    middle_time = (start_seconds + end_seconds) / 2
                    subtitle_times.append(middle_time)
    
    except Exception as e:
        print(f"[ERROR] 读取 ASS 文件失败：{e}")
        return []
    
    if not subtitle_times:
        print("[WARN] 未找到字幕时间点，将按固定间隔抽帧")
        # 按固定间隔抽帧
        interval = total_frames / max_frames
        subtitle_times = [i * interval / fps for i in range(max_frames)]
    
    print(f"[INFO] 找到 {len(subtitle_times)} 个字幕时间点")
    
    # 创建输出目录
    frames_dir = os.path.join(output_dir, "frames")
    os.makedirs(frames_dir, exist_ok=True)
    
    # 抽帧
    frame_paths = []
    for i, timestamp in enumerate(subtitle_times[:max_frames]):
        # 跳转到指定时间
        cap.set(cv2.CAP_PROP_POS_MSEC, timestamp * 1000)
        
        ret, frame = cap.read()
        if not ret:
            print(f"[WARN] 无法读取帧：{timestamp:.2f}s")
            continue
        
        # 保存帧
        frame_path = os.path.join(frames_dir, f"frame_{i:04d}_{timestamp:.2f}s.png")
        cv2.imwrite(frame_path, frame)
        
        frame_paths.append(frame_path)
        print(f"[INFO] 已抽帧：{timestamp:.2f}s -> {frame_path}")
    
    cap.release()
    
    print(f"[OK] 共抽出 {len(frame_paths)} 帧")
    return frame_paths


def ass_time_to_seconds(ass_time: str) -> float:
    """
    将 ASS 时间戳转换为秒
    
    :param ass_time: ASS 时间戳（H:MM:SS.cc）
    :return: 秒数
    """
    parts = ass_time.split(":")
    
    hours = int(parts[0])
    minutes = int(parts[1])
    seconds_parts = parts[2].split(".")
    seconds = int(seconds_parts[0])
    centiseconds = int(seconds_parts[1]) if len(seconds_parts) > 1 else 0
    
    total_seconds = hours * 3600 + minutes * 60 + seconds + centiseconds / 100
    return total_seconds


def generate_contact_sheet(frame_paths: list, output_path: str, cols: int = 4, padding: int = 10):
    """
    生成检查图（多帧拼接）
    
    :param frame_paths: 帧文件路径列表
    :param output_path: 输出图片路径
    :param cols: 列数
    :param padding: 内边距（像素）
    """
    if not frame_paths:
        print("[ERROR] 没有帧文件")
        return False
    
    # 读取第一帧，获取尺寸
    first_frame = Image.open(frame_paths[0])
    frame_width, frame_height = first_frame.size
    
    # 计算拼接图尺寸
    rows = (len(frame_paths) + cols - 1) // cols
    sheet_width = cols * (frame_width + padding) + padding
    sheet_height = rows * (frame_height + padding) + padding
    
    # 创建拼接图
    sheet = Image.new("RGB", (sheet_width, sheet_height), (0, 0, 0))
    
    # 拼接帧
    for i, frame_path in enumerate(frame_paths):
        frame = Image.open(frame_path)
        
        # 计算位置
        row = i // cols
        col = i % cols
        x = col * (frame_width + padding) + padding
        y = row * (frame_height + padding) + padding
        
        # 粘贴帧
        sheet.paste(frame, (x, y))
        
        # 添加时间戳标注
        draw = ImageDraw.Draw(sheet)
        timestamp = os.path.basename(frame_path).split("_")[1].replace(".png", "")
        draw.text((x + 5, y + 5), f"{timestamp}", 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 generate_html_preview(frame_paths: list, output_path: str, video_path: str = None):
    """
    生成 HTML 预览页面
    
    :param frame_paths: 帧文件路径列表
    :param output_path: 输出 HTML 文件路径
    :param video_path: 视频文件路径（可选）
    """
    html_lines = []
    html_lines.append("<!DOCTYPE html>")
    html_lines.append("<html>")
    html_lines.append("<head>")
    html_lines.append("  <meta charset='utf-8'>")
    html_lines.append("  <title>字幕预览检查图</title>")
    html_lines.append("  <style>")
    html_lines.append("    body { font-family: Arial, sans-serif; margin: 20px; }")
    html_lines.append("    .frame { margin: 10px; border: 1px solid #ccc; display: inline-block; }")
    html_lines.append("    .timestamp { font-size: 12px; color: #666; text-align: center; }")
    html_lines.append("  </style>")
    html_lines.append("</head>")
    html_lines.append("<body>")
    html_lines.append("  <h1>字幕预览检查图</h1>")
    
    if video_path:
        html_lines.append(f"  <p>视频：{os.path.basename(video_path)}</p>")
    
    html_lines.append(f"  <p>共 {len(frame_paths)} 帧</p>")
    html_lines.append("  <div>")
    
    for frame_path in frame_paths:
        timestamp = os.path.basename(frame_path).split("_")[1].replace(".png", "")
        rel_path = os.path.relpath(frame_path, os.path.dirname(output_path))
        
        html_lines.append("    <div class='frame'>")
        html_lines.append(f"      <img src='{rel_path}' width='320'>")
        html_lines.append(f"      <div class='timestamp'>{timestamp}</div>")
        html_lines.append("    </div>")
    
    html_lines.append("  </div>")
    html_lines.append("</body>")
    html_lines.append("</html>")
    
    # 写入文件
    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("\n".join(html_lines))
    
    print(f"[OK] HTML 预览页面已生成：{output_path}")


def main():
    parser = argparse.ArgumentParser(
        description="Subtitle Preview Contact Sheet · 字幕预览检查图生成器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例：
  python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --output preview.png
  python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --generate-html --output preview.html
  python subtitle-preview-contact-sheet.py --batch ./videos/ --subtitle-dir ./ass/ --output-dir ./previews/
"""
    )
    parser.add_argument("--video", help="视频文件路径")
    parser.add_argument("--subtitle-ass", help="ASS 字幕文件路径")
    parser.add_argument("--output", help="输出文件路径（PNG 或 HTML）")
    parser.add_argument("--generate-html", action="store_true", help="生成 HTML 预览页面")
    parser.add_argument("--batch", help="批量处理视频目录")
    parser.add_argument("--subtitle-dir", help="字幕文件目录（用于批量处理）")
    parser.add_argument("--output-dir", help="输出目录（用于批量处理）")
    parser.add_argument("--max-frames", type=int, default=20, help="最大抽帧数量")
    parser.add_argument("--cols", type=int, default=4, help="检查图列数")
    parser.add_argument("--padding", type=int, default=10, help="内边距（像素）")
    
    args = parser.parse_args()
    
    if args.video:
        # 单文件处理
        if not args.subtitle_ass:
            print("[ERROR] 请指定 --subtitle-ass")
            sys.exit(1)
        
        if not args.output:
            print("[ERROR] 请指定 --output")
            sys.exit(1)
        
        # 抽帧
        output_dir = os.path.dirname(args.output)
        if not output_dir:
            output_dir = "."
        
        frame_paths = extract_frames_by_subtitle(args.video, args.subtitle_ass, output_dir, args.max_frames)
        
        if not frame_paths:
            sys.exit(1)
        
        # 生成检查图或 HTML
        if args.generate_html:
            generate_html_preview(frame_paths, args.output, args.video)
        else:
            generate_contact_sheet(frame_paths, args.output, args.cols, args.padding)
    
    elif args.batch:
        # 批量处理
        if not args.subtitle_dir:
            print("[ERROR] 批量处理需要指定 --subtitle-dir")
            sys.exit(1)
        
        if not args.output_dir:
            print("[ERROR] 批量处理需要指定 --output-dir")
            sys.exit(1)
        
        os.makedirs(args.output_dir, exist_ok=True)
        
        for video_file in os.listdir(args.batch):
            if video_file.lower().endswith((".mp4", ".mov", ".avi")):
                video_path = os.path.join(args.batch, video_file)
                
                # 查找对应的字幕文件
                subtitle_file = video_file.replace(".mp4", ".ass").replace(".mov", ".ass").replace(".avi", ".ass")
                subtitle_path = os.path.join(args.subtitle_dir, subtitle_file)
                
                if not os.path.isfile(subtitle_path):
                    print(f"[WARN] 未找到对应的字幕文件：{subtitle_file}")
                    continue
                
                print(f"\n[INFO] 处理视频：{video_file}")
                
                # 抽帧
                frame_paths = extract_frames_by_subtitle(video_path, subtitle_path, args.output_dir, args.max_frames)
                
                if frame_paths:
                    # 生成检查图
                    output_path = os.path.join(args.output_dir, video_file.replace(".mp4", "_preview.png"))
                    generate_contact_sheet(frame_paths, output_path, args.cols, args.padding)
                    
                    # 生成 HTML
                    html_path = os.path.join(args.output_dir, video_file.replace(".mp4", "_preview.html"))
                    generate_html_preview(frame_paths, html_path, video_path)
    
    else:
        print("[ERROR] 请指定 --video 或 --batch")
        sys.exit(1)
    
    print("\n[OK] 处理完成")


if __name__ == "__main__":
    main()