import cv2
import os


def v2i(video_path,
        output_dir,
        frame_rate: int,
        frame_nums: int):
    
    video_capture = cv2.VideoCapture(video_path)

    if not video_capture.isOpened():
        print("Error: Cannot open video file.")
        exit()
    
    # 创建文件夹及timestamps.txt
    os.makedirs(os.path.join(output_dir, "images"), exist_ok=True)
    timestamps_path = os.path.join(output_dir, "timestamps.txt")
    with open(timestamps_path, "w") as f:
        pass
    
    video_fps = video_capture.get(cv2.CAP_PROP_FPS)
    print(f"Video FPS: {video_fps}")

    frame_interval = video_fps // frame_rate
    
    frame_count = 0
    saved_frame_cout = 0
    timestamps = []
    
    while True:
        ret, frame = video_capture.read()
        
        if frame_count > frame_nums:
            break
        
        if ret:
            if frame_count % frame_interval == 0:
                
                timestamp = video_capture.get(cv2.CAP_PROP_POS_MSEC) / 1000.0
                
                if frame_count == 0:
                    timestamps.append(0/1000.0)
                
                else:
                    timestamps.append(timestamp)
                
                frame_filename = os.path.join(output_dir, "images", f"image_{saved_frame_cout:06d}.png")

                cv2.imwrite(frame_filename, frame)

                print(f"Saved {frame_filename}, timestamp: {timestamp}")

                saved_frame_cout += 1
        
            frame_count += 1
        
        else:
            break
    
    video_capture.release()
    
    with open(timestamps_path, "w") as f:
        for idx, t in enumerate(timestamps):
            if idx + 1 < len(timestamps):
                f.write(str(t)+'\n')
            else:
                f.write(str(t))
    
    print("Finished extracting frames.")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser()
    
    parser.add_argument("--video_path", type=str)
    
    parser.add_argument("--output_dir", type=str)
    
    parser.add_argument("--frame_rate", type=int, default=25)
    
    parser.add_argument("--frame_nums", type=int, default=90)
    
    args = parser.parse_args()
    
    v2i(args.video_path, args.output_dir, args.frame_rate, args.frame_nums)