import cv2
import os
import argparse

def extract_frames(video_path, interval_sec, resize_dim=None, base_output_dir=None):
    """
    从视频中按指定时间间隔提取帧并保存，可选择调整大小和指定输出目录。

    参数:
    video_path (str): 输入视频文件的完整路径。
    interval_sec (float): 截取帧的时间间隔（秒）。
    resize_dim (tuple or None): (width, height) 元组，用于调整帧大小。
    base_output_dir (str): 保存帧的基础目录。
    """
    
    # 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

    print(f"正在处理视频: {video_path}")
    if resize_dim:
        print(f"将调整所有输出帧大小为: {resize_dim[0]}x{resize_dim[1]}")

    # 3. 构建输出目录
    # 从路径中获取不带扩展名的文件名, e.g., "my_video"
    video_filename_with_ext = os.path.basename(video_path)
    video_name, _ = os.path.splitext(video_filename_with_ext)

    # --- 更新的目录逻辑 ---
    # 构建最终的完整输出路径
    # 结构: <base_output_dir>/<video_name>/images/val
    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. 循环处理视频帧
    save_count = 0
    interval_msec = interval_sec * 1000
    next_save_time_msec = 0.0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        current_msec = cap.get(cv2.CAP_PROP_POS_MSEC)

        if current_msec >= next_save_time_msec:
            frame_to_save = frame
            
            # 调整大小
            if resize_dim:
                frame_to_save = cv2.resize(frame, resize_dim, interpolation=cv2.INTER_AREA)
            
            # 构建保存路径
            # --- 修改文件名逻辑：使用时间戳 ---
            # 1. 获取当前毫秒时间戳
            msec = current_msec 
            
            # 2. 转换为 时、分、秒、毫秒
            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
            
            # 3. 格式化文件名: 格式为 HH-MM-SS-msmsms.jpg (时-分-秒-毫秒)
            #    例如: 00-00-01-500.jpg (表示 1.5 秒时)
            frame_filename = f"{hours:02d}-{minutes:02d}-{seconds:02d}-{milliseconds:03d}.jpg"
            save_path = os.path.join(output_dir, frame_filename)
            # --- 结束修改 ---

            # 保存
            cv2.imwrite(save_path, frame_to_save)
            
            print(f"已保存: {save_path} (时间戳: {current_msec/1000:.3f}s)")
            save_count += 1
            next_save_time_msec += interval_msec

    # 6. 释放资源
    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')"
    )
    
    # --- 新增参数: --output_dir ---
    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 参数 ---
    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

    # --- 调用主处理函数 ---
    # 传入新参数 args.output_dir
    extract_frames(args.video, args.interval, resize_dim, args.output_dir)

if __name__ == "__main__":
    main()