#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
视频时间段提取工具
功能：从视频文件中提取指定时间段的内容，并保存为图片帧序列
作者：工程师
"""

import cv2
import os
import argparse
import sys
from pathlib import Path
import time


def parse_time_to_seconds(time_str):
    """
    将时间字符串转换为秒数
    支持格式：
    - 秒数："30" 或 "30.5"
    - 分:秒："1:30" 或 "1:30.5"
    - 时:分:秒："0:1:30" 或 "0:1:30.5"
    """
    try:
        if ':' not in time_str:
            return float(time_str)
        
        parts = time_str.split(':')
        if len(parts) == 2:  # MM:SS
            minutes, seconds = parts
            return int(minutes) * 60 + float(seconds)
        elif len(parts) == 3:  # HH:MM:SS
            hours, minutes, seconds = parts
            return int(hours) * 3600 + int(minutes) * 60 + float(seconds)
        else:
            raise ValueError("时间格式不正确")
    except ValueError as e:
        raise ValueError(f"无法解析时间格式 '{time_str}': {e}")


def extract_video_segment(video_path, start_time, end_time, output_dir, fps=None, prefix="frame"):
    """
    从视频中提取指定时间段的帧
    
    Args:
        video_path (str): 输入视频文件路径
        start_time (str): 开始时间（秒或时:分:秒格式）
        end_time (str): 结束时间（秒或时:分:秒格式）
        output_dir (str): 输出目录路径
        fps (float): 提取帧率，None表示使用原视频帧率
        prefix (str): 输出文件名前缀
    
    Returns:
        tuple: (成功提取的帧数, 总处理时间)
    """
    # 检查输入文件
    if not os.path.exists(video_path):
        raise FileNotFoundError(f"视频文件不存在: {video_path}")
    
    # 创建输出目录
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 解析时间
    start_seconds = parse_time_to_seconds(start_time)
    end_seconds = parse_time_to_seconds(end_time)
    
    if start_seconds >= end_seconds:
        raise ValueError("开始时间必须小于结束时间")
    
    print(f"正在处理视频: {video_path}")
    print(f"提取时间段: {start_seconds:.2f}s - {end_seconds:.2f}s")
    print(f"输出目录: {output_dir}")
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"无法打开视频文件: {video_path}")
    
    try:
        # 获取视频信息
        video_fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        video_duration = total_frames / video_fps
        
        print(f"视频信息: {video_fps:.2f} FPS, {total_frames} 帧, {video_duration:.2f}s")
        
        # 检查时间范围
        if end_seconds > video_duration:
            print(f"警告: 结束时间 {end_seconds:.2f}s 超过视频长度 {video_duration:.2f}s，将调整为视频结束时间")
            end_seconds = video_duration
        
        # 设置提取帧率
        extract_fps = fps if fps is not None else video_fps
        frame_interval = video_fps / extract_fps if extract_fps <= video_fps else 1
        
        print(f"提取帧率: {extract_fps:.2f} FPS (每 {frame_interval:.2f} 帧提取一帧)")
        
        # 定位到开始时间
        start_frame = int(start_seconds * video_fps)
        end_frame = int(end_seconds * video_fps)
        cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
        
        extracted_count = 0
        current_frame = start_frame
        next_extract_frame = start_frame
        start_time_process = time.time()
        
        print(f"开始提取帧 ({start_frame} - {end_frame})...")
        
        while current_frame <= end_frame:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 检查是否需要提取当前帧
            if current_frame >= next_extract_frame:
                # 生成输出文件名
                output_filename = f"{prefix}_{extracted_count:06d}.jpg"
                output_filepath = output_path / output_filename
                
                # 保存帧
                success = cv2.imwrite(str(output_filepath), frame)
                if success:
                    extracted_count += 1
                    if extracted_count % 100 == 0:
                        progress = (current_frame - start_frame) / (end_frame - start_frame) * 100
                        print(f"已提取 {extracted_count} 帧 ({progress:.1f}%)")
                else:
                    print(f"警告: 无法保存帧到 {output_filepath}")
                
                # 计算下一个要提取的帧
                next_extract_frame += frame_interval
            
            current_frame += 1
        
        process_time = time.time() - start_time_process
        
        print(f"\n提取完成!")
        print(f"成功提取 {extracted_count} 帧")
        print(f"处理时间: {process_time:.2f}s")
        print(f"输出目录: {output_dir}")
        
        return extracted_count, process_time
        
    finally:
        cap.release()


def main():
    parser = argparse.ArgumentParser(
        description="从视频文件中提取指定时间段的帧",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
时间格式示例:
  30        - 30秒
  1:30      - 1分30秒
  0:1:30    - 1分30秒
  1:23:45   - 1小时23分45秒
  
使用示例:
  python video_segment_extractor.py input.mp4 10 60 -o frames
  python video_segment_extractor.py input.mp4 0:10 1:00 -o frames --fps 10
        """
    )
    
    parser.add_argument("video", help="输入视频文件路径")
    parser.add_argument("start_time", help="开始时间 (秒数或时:分:秒格式)")
    parser.add_argument("end_time", help="结束时间 (秒数或时:分:秒格式)")
    parser.add_argument("-o", "--output", default="extracted_frames", 
                       help="输出目录 (默认: extracted_frames)")
    parser.add_argument("--fps", type=float, 
                       help="提取帧率 (默认使用原视频帧率)")
    parser.add_argument("--prefix", default="frame", 
                       help="输出文件名前缀 (默认: frame)")
    
    args = parser.parse_args()
    
    try:
        extracted_count, process_time = extract_video_segment(
            args.video,
            args.start_time,
            args.end_time,
            args.output,
            args.fps,
            args.prefix
        )
        
        print(f"\n任务完成! 共提取 {extracted_count} 帧，耗时 {process_time:.2f}s")
        
    except Exception as e:
        print(f"错误: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()