# -*- coding: utf-8 -*-
"""
tool_extract_snake_frames.py
从「贪蛇出动」游戏录屏视频中抽帧，用于分析玩法与抠素材。
用法:
  python tool_extract_snake_frames.py                # 默认每 2 秒抽 1 帧
  python tool_extract_snake_frames.py 0.5            # 每 0.5 秒抽 1 帧
  python tool_extract_snake_frames.py 0.2 60 70      # 每 0.2 秒抽 1 帧，只抽 60~70 秒区间
"""
import os
import sys
import cv2

VIDEO = r"D:\Trae_Files\TRAE-Toolbox\temp\a4ef87cf84a3152cb5800eae5e06e6a8.mp4"
OUT_DIR = r"D:\Trae_Files\TRAE-Toolbox\temp\snake_frames"


def main():
    step = float(sys.argv[1]) if len(sys.argv) > 1 else 2.0
    start = float(sys.argv[2]) if len(sys.argv) > 2 else 0.0
    end = float(sys.argv[3]) if len(sys.argv) > 3 else 1e18

    os.makedirs(OUT_DIR, exist_ok=True)
    cap = cv2.VideoCapture(VIDEO)
    if not cap.isOpened():
        print("无法打开视频:", VIDEO)
        sys.exit(1)
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f"fps={fps:.2f} total_frames={total} duration={total / fps:.1f}s step={step}s")

    t = start
    saved = 0
    while t <= end:
        cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000.0)
        ok, frame = cap.read()
        if not ok:
            break
        name = "t%07d.jpg" % int(round(t * 10))  # 文件名单位：0.1 秒
        cv2.imwrite(os.path.join(OUT_DIR, name), frame, [cv2.IMWRITE_JPEG_QUALITY, 88])
        saved += 1
        t += step
    cap.release()
    print(f"saved={saved} -> {OUT_DIR}")


if __name__ == "__main__":
    main()