#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import subprocess
import argparse
import sys
import platform
import shutil  # 用于检查 ffprobe 是否存在

# 常见的视频文件扩展名列表 (可以根据需要添加或删除)
VIDEO_EXTENSIONS = {
    '.mp4', '.mkv', '.webm', '.avi', '.mov', '.flv', '.wmv', '.ts', '.m2ts'
}

def check_ffprobe():
    """检查 ffprobe 是否在系统 PATH 中可用"""
    ffprobe_path = shutil.which('ffprobe')
    if not ffprobe_path:
        print("错误: 'ffprobe' 命令未找到。请确保 FFmpeg 已正确安装并添加到系统环境变量 PATH 中。", file=sys.stderr)
        sys.exit(1)
    print(f"找到 ffprobe: {ffprobe_path}")
    return ffprobe_path

def get_video_codec(file_path, ffprobe_path):
    """使用 ffprobe 获取视频文件的第一个视频流的编码名称"""
    command = [
        ffprobe_path,
        '-v', 'error',              # 只输出错误信息
        '-select_streams', 'v:0',   # 选择第一个视频流
        '-show_entries', 'stream=codec_name', # 只显示流的编码名称
        '-of', 'default=noprint_wrappers=1:nokey=1', # 输出格式: 只输出值
        file_path
    ]

    try:
        # 在 Windows 上，隐藏命令行窗口
        startupinfo = None
        if platform.system() == "Windows":
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            startupinfo.wShowWindow = subprocess.SW_HIDE

        result = subprocess.run(
            command,
            capture_output=True, # 捕获标准输出和标准错误
            text=True,           # 以文本模式处理输出 (自动解码)
            check=False,         # 不在 ffprobe 返回非零时抛出异常，手动检查
            startupinfo=startupinfo # 隐藏 Windows 上的窗口
        )

        if result.returncode != 0:
            # ffprobe 执行出错 (可能是文件损坏或非视频文件)
            # print(f"警告: 处理文件 '{file_path}' 时 ffprobe 出错: {result.stderr.strip()}", file=sys.stderr)
            return None # 返回 None 表示无法确定编码

        codec_name = result.stdout.strip()
        return codec_name if codec_name else None # 如果输出为空也返回 None

    except FileNotFoundError:
        print(f"错误: 无法执行 '{ffprobe_path}'。请确保它在您的 PATH 中。", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"错误: 处理文件 '{file_path}' 时发生未知错误: {e}", file=sys.stderr)
        return None

def find_av1_videos(directory, ffprobe_path):
    """在指定目录及其子目录中查找 AV1 编码的视频文件"""
    print(f"\n正在扫描目录: {directory}")
    av1_files_count = 0
    processed_files_count = 0

    for root, _, files in os.walk(directory):
        for filename in files:
            # 检查文件扩展名是否在视频扩展名列表中
            _, ext = os.path.splitext(filename)
            if ext.lower() in VIDEO_EXTENSIONS:
                file_path = os.path.join(root, filename)
                processed_files_count += 1
                print(f"  正在检查: {file_path}", end='\r', flush=True) # 显示当前检查的文件

                codec = get_video_codec(file_path, ffprobe_path)

                # 清除行尾的 "正在检查..."
                print(" " * (len(file_path) + 15), end='\r', flush=True)

                if codec == 'av1':
                    print(f"  [AV1] {file_path}") # 找到 AV1 文件，打印完整路径
                    av1_files_count += 1
                # elif codec: # 可选：打印非 AV1 视频的编码
                #     print(f"  [{codec.upper()}] {file_path}")
                # else: # 可选：打印无法处理的文件
                #     print(f"  [无法识别/错误] {file_path}")


    print(f"\n扫描完成。共处理 {processed_files_count} 个可能的视频文件，找到 {av1_files_count} 个 AV1 编码视频。")

def main():
    parser = argparse.ArgumentParser(
        description="在指定目录及其子目录中查找所有 AV1 编码的视频文件。",
        epilog="需要系统安装 FFmpeg 并将 ffprobe 添加到环境变量 PATH。"
    )
    parser.add_argument(
        "directory",
        help="要扫描的根目录路径。"
    )
    args = parser.parse_args()

    target_directory = args.directory

    # 检查目录是否存在
    if not os.path.isdir(target_directory):
        print(f"错误: 目录 '{target_directory}' 不存在或不是一个有效的目录。", file=sys.stderr)
        sys.exit(1)

    # 检查 ffprobe 是否可用
    ffprobe_executable = check_ffprobe()

    # 开始查找
    find_av1_videos(target_directory, ffprobe_executable)

if __name__ == "__main__":
    main()