#!/usr/bin/env python3
"""
视频切割脚本 - 将视频逐帧切割为图片
支持扫描目录下所有视频文件，多进程并行处理，已切割的自动跳过
"""

import os
import sys
import cv2
import argparse
from multiprocessing import Pool, cpu_count
from pathlib import Path

# 支持的视频扩展名
VIDEO_EXTENSIONS = {'.mp4', '.avi', '.mkv', '.mov', '.flv', '.wmv', '.webm', '.ts', '.m4v'}


def scan_videos(root_path: str) -> list[str]:
    """扫描目录下所有视频文件"""
    videos = []
    root = Path(root_path)
    if root.is_file() and root.suffix.lower() in VIDEO_EXTENSIONS:
        return [str(root)]
    for p in sorted(root.rglob('*')):
        if p.is_file() and p.suffix.lower() in VIDEO_EXTENSIONS:
            videos.append(str(p))
    return videos


def extract_frames(video_path: str, output_dir: str = None) -> dict:
    """
    将单个视频按每秒1帧切割为图片

    Args:
        video_path: 视频文件路径
        output_dir: 输出目录（默认在视频同目录下创建同名文件夹）

    Returns:
        处理结果字典
    """
    video_path = os.path.abspath(video_path)
    video_name = Path(video_path).stem

    if output_dir is None:
        output_dir = os.path.join(os.path.dirname(video_path), video_name + '_frames')

    # 已经切割过则跳过
    if os.path.isdir(output_dir) and any(Path(output_dir).glob('*.jpg')):
        existing = len(list(Path(output_dir).glob('*.jpg')))
        return {
            'video': video_path,
            'status': 'skipped',
            'reason': f'已存在，共 {existing} 张图片',
            'output': output_dir,
        }

    os.makedirs(output_dir, exist_ok=True)

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return {
            'video': video_path,
            'status': 'error',
            'reason': '无法打开视频文件',
            'output': output_dir,
        }

    fps = cap.get(cv2.CAP_PROP_FPS) or 25
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps if fps else 0

    # 每秒取一帧：根据fps算出每隔多少帧取一帧
    frame_interval = int(fps)
    frame_idx = 0
    saved = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_idx % frame_interval == 0:
            # 用视频名+时间戳命名：第01集_01m23s_456.jpg
            sec = frame_idx / fps
            mins = int(sec // 60)
            secs = int(sec % 60)
            ms = int((sec % 1) * 1000)
            filename = os.path.join(output_dir, f'{video_name}_{mins:02d}m{secs:02d}s_{ms:03d}.jpg')
            cv2.imwrite(filename, frame)
            saved += 1
        frame_idx += 1

    cap.release()

    return {
        'video': video_path,
        'status': 'done',
        'frames_total': total_frames,
        'frames_saved': saved,
        'duration': f'{duration:.1f}s',
        'fps': fps,
        'output': output_dir,
    }


def _worker(args):
    """多进程 worker 包装"""
    video_path, output_dir = args
    try:
        return extract_frames(video_path, output_dir)
    except Exception as e:
        return {
            'video': video_path,
            'status': 'error',
            'reason': str(e),
        }


def main(video_path, output=None, interval=1, jobs=None):
    videos = scan_videos(video_path)
    if not videos:
        print(f'未在 {video_path} 中找到视频文件')
        return

    jobs = jobs or cpu_count()
    print(f'找到 {len(videos)} 个视频，使用 {jobs} 个进程并行处理')

    tasks = []
    for v in videos:
        if output:
            video_name = Path(v).stem
            out = os.path.join(output, video_name + '_frames')
        else:
            out = None
        tasks.append((v, out))

    done = 0
    skipped = 0
    errors = 0

    with Pool(jobs) as pool:
        for result in pool.imap_unordered(_worker, tasks):
            status = result['status']
            video = result['video']
            if status == 'done':
                done += 1
                print(f'[完成] {video} -> {result["frames_saved"]}帧 '
                      f'(总{result["frames_total"]}, {result["duration"]})')
            elif status == 'skipped':
                skipped += 1
                print(f'[跳过] {video} ({result["reason"]})')
            else:
                errors += 1
                print(f'[错误] {video} ({result["reason"]})')

    print(f'\n处理完毕: 完成 {done}, 跳过 {skipped}, 错误 {errors}, 共 {len(videos)}')


# ========== 运行设置（改这里）==========
VIDEO_DIR = "./shipin"       # 视频文件或目录路径
OUTPUT_DIR = None            # 输出目录，None 则在视频同目录
INTERVAL = 1                 # 每隔 N 帧取一帧
JOBS = None                  # 并行进程数，None 则用 CPU 核心数
# ======================================


if __name__ == '__main__':
    main(VIDEO_DIR, OUTPUT_DIR, INTERVAL, JOBS)