import cv2
import os
import argparse
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm # 导入 tqdm

def process_and_save_frame(frame, save_path, resize_dim):
    """
    一个在工作线程中运行的辅助函数，用于调整帧大小并将其保存到磁盘。
    
    参数:
    frame (numpy.ndarray): 要处理的帧 (应该是 .copy() 过的)
    save_path (str): 完整的保存路径
    resize_dim (tuple or None): (width, height) 调整大小的元组
    """
    try:
        frame_to_save = frame
        
        # 调整大小 (CPU密集型, 但OpenCV通常会释放GIL)
        if resize_dim:
            frame_to_save = cv2.resize(frame, resize_dim, interpolation=cv2.INTER_AREA)
        
        # 保存 (I/O 密集型)
        cv2.imwrite(save_path, frame_to_save)
        
        # print(f"已保存: {save_path} (时间戳: {current_msec/1000:.3f}s)") 
        # 注意：在tqdm循环中打印会打乱进度条，故注释掉
        
    except Exception as e:
        print(f"错误: 无法保存帧 {save_path}. 错误: {e}")

def extract_frames(video_path, interval_sec, resize_dim=None, base_output_dir=None):
    """
    从视频中按指定时间间隔提取帧并保存，可选择调整大小和指定输出目录。
    (已更新为使用 ThreadPoolExecutor 和 tqdm)
    """
    
    # 1. 检查视频文件是否存在
    if not os.path.exists(video_path):
        print(f"错误: 视频文件不存在于路径 {video_path}")
        return

    # 2. 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"错误: 无法打开视频文件 {video_path}")
        return

    # --- 获取视频总帧数以用于 tqdm ---
    try:
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if total_frames <= 0:
            print("警告: 无法获取视频总帧数。进度条可能无法正确显示。")
            total_frames = None # 设为 None, tqdm 会显示为未知进度
    except Exception:
        total_frames = None
        
    print(f"正在处理视频: {video_path}")
    if resize_dim:
        print(f"将调整所有输出帧大小为: {resize_dim[0]}x{resize_dim[1]}")

    # 3. 构建输出目录
    video_filename_with_ext = os.path.basename(video_path)
    video_name, _ = os.path.splitext(video_filename_with_ext)
    output_dir = os.path.join(base_output_dir, video_name, "images", "val")

    # 4. 创建目录 (如果不存在)
    try:
        os.makedirs(output_dir, exist_ok=True)
        print(f"帧将保存到: {output_dir}")
    except OSError as e:
        print(f"错误: 无法创建目录 {output_dir}. 错误信息: {e}")
        cap.release()
        return

    # 5. 循环处理视频帧 (使用并行和tqdm)
    save_count = 0
    interval_msec = interval_sec * 1000
    next_save_time_msec = 0.0

    # 使用 ThreadPoolExecutor 来并行处理 I/O 密集型任务 (保存图片)
    # 使用 tqdm 显示进度条
    
    # max_workers=None 会自动使用合理的线程数 (通常是 CPU核心数 * 5)
    with ThreadPoolExecutor(max_workers=None) as executor:
        
        # 使用 tqdm 包裹循环
        with tqdm(total=total_frames, unit="frames", desc=f"Processing {video_name}") as pbar:
            while True:
                ret, frame = cap.read()
                if not ret:
                    break # 视频结束
                
                # 无论是否保存，tqdm 进度条都更新 (表示已读取一帧)
                pbar.update(1) 

                current_msec = cap.get(cv2.CAP_PROP_POS_MSEC)

                if current_msec >= next_save_time_msec:
                    frame_to_save = frame
                    
                    # 构建保存路径
                    msec = current_msec 
                    total_seconds_int = int(msec // 1000)
                    milliseconds = int(msec % 1000)
                    hours = total_seconds_int // 3600
                    minutes = (total_seconds_int % 3600) // 60
                    seconds = total_seconds_int % 60
                    
                    frame_filename = f"{hours:02d}-{minutes:02d}-{seconds:02d}-{milliseconds:03d}.jpg"
                    save_path = os.path.join(output_dir, frame_filename)

                    # --- 并行处理 ---
                    # 将“调整大小和保存”任务提交到线程池
                    # 必须使用 frame.copy()！
                    # 因为 'frame' 是一个可重用缓冲区，主循环会立即用下一帧覆盖它
                    # 如果不复制，所有线程最终可能保存的是同一帧（或损坏的数据）
                    executor.submit(process_and_save_frame, frame.copy(), save_path, resize_dim)
                    # ---------------
                    
                    save_count += 1
                    next_save_time_msec += interval_msec

    # 6. 释放资源
    # 'with' 语句块结束时，executor 会自动调用 shutdown(wait=True)，
    # 确保所有保存任务完成后才继续。
    cap.release()
    print(f"\n处理完成。")
    print(f"总共提交了 {save_count} 帧到目录 {output_dir}")


def main():
    parser = argparse.ArgumentParser(description="视频帧提取脚本 (并行版)")

    parser.add_argument(
        "-v", "--video",
        type=str,
        required=True,
        help="[必需] 输入视频文件的路径 (例如: /home/user/videos/my_test.mp4)"
    )
    
    parser.add_argument(
        "-i", "--interval",
        type=float,
        default=0.5,
        help="[可选] 截取帧的时间间隔 (秒). 默认: 0.5"
    )

    parser.add_argument(
        "-r", "--resize",
        type=str,
        default=None,
        help="[可选] 将输出帧调整为 'WIDTHxHEIGHT' 格式 (例如: '1920x1080')"
    )
    
    default_output_base = os.path.join("..", "DataSet_Public", "5_Predict_Video")
    parser.add_argument(
        "-o", "--output_dir",
        type=str,
        default=default_output_base,
        help=f"[可选] 保存帧的基础目录. "
             f"最终路径将是: <output_dir>/<视频名>/images/val. "
             f"默认: {default_output_base}"
    )
    
    args = parser.parse_args()

    resize_dim = None
    if args.resize:
        try:
            width_str, height_str = args.resize.split('x')
            width = int(width_str)
            height = int(height_str)
            if width <= 0 or height <= 0:
                raise ValueError("宽度和高度必须为正数")
            resize_dim = (width, height)
        except ValueError as e:
            print(f"错误: 无效的 --resize 格式 '{args.resize}'.")
            print("必须使用 'WIDTHxHEIGHT' 格式 (例如: '1920x1080').")
            return

    extract_frames(args.video, args.interval, resize_dim, args.output_dir)

if __name__ == "__main__":
    main()