#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
视频画面提取工具
从指定视频文件的特定时间段提取画面帧并保存为图像文件
"""

import cv2
import os
import sys
from pathlib import Path


def extract_frames_from_video(video_path, start_time_seconds, end_time_seconds, output_dir, target_frame_count=500, image_extension="jpg", image_params=None):
    """
    从视频中提取指定时间段的画面帧
    
    参数:
        video_path (str): 视频文件路径
        start_time_seconds (int): 开始时间（秒）
        end_time_seconds (int): 结束时间（秒）
        output_dir (str): 输出目录路径
        target_frame_count (int): 目标提取帧数，默认500帧
    
    返回:
        bool: 提取是否成功
    """
    
    # 检查视频文件是否存在
    if not os.path.exists(video_path):
        print(f"错误：视频文件不存在 - {video_path}")
        return False
    
    # 创建输出目录
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    
    if not cap.isOpened():
        print(f"错误：无法打开视频文件 - {video_path}")
        return False
    
    # 获取视频基本信息
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps
    
    print(f"视频信息:")
    print(f"  帧率: {fps:.2f} FPS")
    print(f"  总帧数: {total_frames}")
    print(f"  总时长: {duration:.2f} 秒")
    
    # 计算开始和结束帧号
    start_frame = int(start_time_seconds * fps)
    end_frame = int(end_time_seconds * fps)
    
    # 检查时间范围是否有效
    if start_frame >= total_frames:
        print(f"错误：开始时间超出视频长度")
        cap.release()
        return False
    
    if end_frame > total_frames:
        print(f"警告：结束时间超出视频长度，将提取到视频结尾")
        end_frame = total_frames
    
    # 计算实际提取的帧数范围
    frame_range = end_frame - start_frame
    
    if frame_range <= 0:
        print(f"错误：无效的时间范围")
        cap.release()
        return False
    
    # 计算帧间隔以获得目标帧数
    if frame_range <= target_frame_count:
        # 如果总帧数少于目标帧数，提取所有帧
        frame_interval = 1
        actual_frame_count = frame_range
    else:
        # 计算间隔以获得约500帧
        frame_interval = frame_range // target_frame_count
        actual_frame_count = frame_range // frame_interval
    
    print(f"提取参数:")
    print(f"  开始时间: {start_time_seconds}秒 (第{start_frame}帧)")
    print(f"  结束时间: {end_time_seconds}秒 (第{end_frame}帧)")
    print(f"  帧间隔: {frame_interval}")
    print(f"  预计提取帧数: {actual_frame_count}")
    
    # 跳转到开始帧
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
    
    extracted_count = 0
    current_frame = start_frame
    
    print(f"开始提取画面...")
    
    while current_frame < end_frame:
        ret, frame = cap.read()
        
        if not ret:
            print(f"警告：读取帧失败，停止提取")
            break
        
        # 只保存指定间隔的帧
        if (current_frame - start_frame) % frame_interval == 0:
            # 计算时间戳
            timestamp_seconds = current_frame / fps
            minutes = int(timestamp_seconds // 60)
            seconds = int(timestamp_seconds % 60)
            milliseconds = int((timestamp_seconds % 1) * 1000)
            
            # 生成文件名
            filename = f"frame_{extracted_count:04d}_time_{minutes:02d}m{seconds:02d}s{milliseconds:03d}ms.{image_extension}"
            output_file = output_path / filename

            # 保存图像
            success = cv2.imwrite(str(output_file), frame, image_params if image_params is not None else [])
            
            if success:
                extracted_count += 1
                if extracted_count % 50 == 0:
                    print(f"  已提取 {extracted_count} 帧...")
            else:
                print(f"警告：保存图像失败 - {output_file}")
        
        current_frame += 1
    
    cap.release()
    
    print(f"提取完成！")
    print(f"  成功提取 {extracted_count} 帧")
    print(f"  输出目录: {output_dir}")
    
    return True


def main():
    """主函数"""
    
    # 视频文件路径
    video_path = r"E:\20251204\D15_20251130235724.mp4"
    
    # 时间设置（28分钟到30分钟）
    start_time_minutes = 0
    end_time_minutes = 15
    start_time_seconds = start_time_minutes * 60  # 28分钟 = 1680秒
    end_time_seconds = end_time_minutes * 60      # 30分钟 = 1800秒
    
    # 输出目录
    output_dir = r"D:\data\20251204\OUT5"
    
    # 目标提取帧数
    target_frame_count = 100
    
    print(f"视频画面提取工具")
    print(f"=" * 50)
    print(f"输入视频: {video_path}")
    print(f"提取时间段: {start_time_minutes}:00 - {end_time_minutes}:00")
    print(f"目标帧数: {target_frame_count}")
    print(f"输出目录: {output_dir}")
    print(f"=" * 50)
    
    # 执行提取
    success = extract_frames_from_video(
        video_path=video_path,
        start_time_seconds=start_time_seconds,
        end_time_seconds=end_time_seconds,
        output_dir=output_dir,
        target_frame_count=target_frame_count
    )
    
    if success:
        print(f"\n✓ 画面提取成功完成！")
    else:
        print(f"\n✗ 画面提取失败！")
        sys.exit(1)


if __name__ == "__main__":
    main()