#!/usr/bin/env python3
"""
为每个帧图片添加对应的字幕文本
字体：50号，红勾边黄字，自动换行
保留原始图片，生成新的带字幕图片
"""
import os
import subprocess
import re
import textwrap

# 导入编码处理模块
from console_helper import (
    safe_print, print_header, print_step, print_info,
    print_success, print_fail, print_done
)


def parse_srt_file(subtitle_path):
    """解析SRT字幕文件，返回每句的字幕文本"""
    subtitles = {}
    try:
        with open(subtitle_path, 'r', encoding='utf-8') as f:
            content = f.read()

        # 按空行分割字幕块
        blocks = content.strip().split('\n\n')

        for block in blocks:
            lines = block.strip().split('\n')
            if len(lines) >= 3:
                # 第一行是序号
                line_number = int(lines[0].strip())
                # 获取字幕文本（可能有多行）
                text = ' '.join(lines[2:])
                subtitles[line_number] = text

        return subtitles
    except Exception as e:
        safe_print(f"      解析字幕文件失败: {e}")
        return {}


def escape_text_for_ffmpeg(text):
    """为FFmpeg的drawtext滤镜转义文本"""
    # 替换特殊字符 - 按照FFmpeg文档的要求
    # 单引号在drawtext中需要转义
    escaped = text.replace("'", "'\\''")  # 单引号转义
    return escaped


def wrap_text(text, max_chars=35):
    """将长文本按指定字符数换行"""
    # 使用textwrap进行换行
    wrapped = textwrap.fill(text, width=max_chars)
    return wrapped


def add_subtitle_to_image(input_image, output_image, subtitle_text, ffmpeg_path="ffmpeg"):
    """为图片添加字幕"""
    # 处理字幕文本，自动换行
    wrapped_text = wrap_text(subtitle_text, max_chars=35)

    # 将文本按行分割
    lines = wrapped_text.split('\n')

    # 构建多个drawtext滤镜，每行一个
    filters = []
    y_position = 50  # 起始Y位置
    line_height = 60  # 每行高度（字体50 + 间距10）

    for i, line in enumerate(lines):
        # 转义特殊字符
        safe_line = escape_text_for_ffmpeg(line)
        current_y = y_position + i * line_height

        filter_str = (
            f"drawtext=text='{safe_line}':"
            f"fontsize=50:"
            f"fontcolor=yellow:"
            f"borderw=3:"
            f"bordercolor=red:"
            f"x=(w-text_w)/2:"
            f"y={current_y}"
        )
        filters.append(filter_str)

    # 组合所有滤镜
    vf_filter = ','.join(filters)

    args = [
        ffmpeg_path,
        '-y',
        '-i', input_image,
        '-vf', vf_filter,
        '-frames:v', '1',
        output_image
    ]

    try:
        result = subprocess.run(args, check=True, capture_output=True, text=True)
        return True
    except subprocess.CalledProcessError as e:
        safe_print(f"      添加字幕失败")
        if e.stderr:
            safe_print(f"      错误: {e.stderr[:200]}")
        return False


def main():
    frames_dir = "end_frames"
    subtitle_path = "subtitles.srt"
    output_dir = "frames_with_subtitle"

    print_header("为帧图片添加字幕")
    safe_print(f"帧图片目录: {frames_dir}")
    safe_print(f"字幕文件: {subtitle_path}")
    safe_print("")

    # 检查目录和文件是否存在
    if not os.path.exists(frames_dir):
        safe_print(f"错误: 帧图片目录 {frames_dir} 不存在")
        return

    if not os.path.exists(subtitle_path):
        safe_print(f"错误: 字幕文件 {subtitle_path} 不存在")
        return

    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 1. 解析字幕文件
    print_step(1, 2, "解析字幕文件...")
    subtitles = parse_srt_file(subtitle_path)
    if not subtitles:
        safe_print("解析字幕失败，退出")
        return

    print_info("字幕行数", str(len(subtitles)))
    safe_print("")

    # 2. 为每个帧图片添加字幕
    print_step(2, 2, "添加字幕...")
    processed_count = 0
    failed_count = 0

    # 获取所有帧图片
    frame_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.jpg')])

    for frame_file in frame_files:
        # 从文件名提取行号
        match = re.match(r'frame_(\d+)\.jpg', frame_file)
        if match:
            line_number = int(match.group(1))

            if line_number in subtitles:
                input_image = os.path.join(frames_dir, frame_file)
                output_image = os.path.join(output_dir, f"frame_{line_number:03d}_with_subtitle.jpg")
                subtitle_text = subtitles[line_number]

                safe_print(f"      处理第{line_number}行...")

                if add_subtitle_to_image(input_image, output_image, subtitle_text):
                    print_info(f"      已保存", f"{output_image}")
                    processed_count += 1
                else:
                    safe_print(f"      处理第{line_number}行失败")
                    failed_count += 1
            else:
                safe_print(f"      警告: 第{line_number}行没有找到对应的字幕文本")

    safe_print("")
    print_info("成功处理", f"{processed_count}/{len(frame_files)} 张图片")
    if failed_count > 0:
        print_info("失败", f"{failed_count} 张图片")
    print_info("输出目录", output_dir)
    print_info("原始图片", f"保留在 {frames_dir} 目录")

    print_done()


if __name__ == '__main__':
    main()