#!/usr/bin/env python3
"""
按照字幕时间轴的结束时间，从视频中提取每一句对应的帧并保存为图片
"""
import os
import subprocess
import re

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

def parse_srt_time(time_str):
    """解析SRT时间格式 (00:00:01,129) 转换为秒"""
    # 匹配格式: 00:00:01,129 或 00:00:01.129
    match = re.match(r'(\d{2}):(\d{2}):(\d{2})[,.](\d{3})', time_str)
    if match:
        hours, minutes, seconds, milliseconds = match.groups()
        return int(hours) * 3600 + int(minutes) * 60 + int(seconds) + int(milliseconds) / 1000
    return 0

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 = lines[0].strip()
                # 第二行是时间轴
                time_line = lines[1].strip()
                # 解析时间轴
                if '-->' in time_line:
                    start_time_str, end_time_str = time_line.split('-->')
                    end_time = parse_srt_time(end_time_str.strip())
                    # 获取字幕文本（可能有多行）
                    text = ' '.join(lines[2:])
                    subtitles.append({
                        'line_number': int(line_number),
                        'end_time': end_time,
                        'text': text
                    })
        
        return subtitles
    except Exception as e:
        safe_print(f"      解析字幕文件失败: {e}")
        return []

def extract_frame_at_time(video_path, time_seconds, output_image, ffmpeg_path="ffmpeg"):
    """从视频中提取指定时间的帧"""
    args = [
        ffmpeg_path,
        '-y',
        '-ss', str(time_seconds),
        '-i', video_path,
        '-vframes', '1',
        '-q:v', '2',
        output_image
    ]
    
    try:
        subprocess.run(args, check=True, capture_output=True)
        return True
    except subprocess.CalledProcessError as e:
        safe_print(f"      提取帧失败: {e}")
        return False

def main():
    video_path = "input.mp4"
    subtitle_path = "subtitles.srt"
    output_dir = "end_frames"
    
    print_header("提取结束时间帧")
    safe_print(f"视频文件: {video_path}")
    safe_print(f"字幕文件: {subtitle_path}")
    safe_print("")
    
    # 检查文件是否存在
    if not os.path.exists(video_path):
        safe_print(f"错误: 视频文件 {video_path} 不存在")
        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, "提取帧...")
    extracted_count = 0
    
    for subtitle in subtitles:
        line_number = subtitle['line_number']
        end_time = subtitle['end_time']
        text = subtitle['text']
        
        output_image = os.path.join(output_dir, f"frame_{line_number:03d}.jpg")
        
        safe_print(f"      提取第{line_number}行帧 (时间: {end_time:.3f}秒)...")
        
        if extract_frame_at_time(video_path, end_time, output_image):
            print_info(f"      已保存", f"{output_image}")
            extracted_count += 1
        else:
            safe_print(f"      提取第{line_number}行帧失败")
    
    safe_print("")
    print_info("成功提取", f"{extracted_count}/{len(subtitles)} 帧")
    print_info("输出目录", output_dir)
    
    print_done()

if __name__ == '__main__':
    main()