#!/usr/bin/env python3
import argparse
import subprocess
import ffmpeg
import os
import sys

V4L2_CODEC_NAMES = {
    "h264": "h264_v4l2m2m",
    "hevc": "hevc_v4l2m2m",
    "vp9": "vp9_v4l2m2m",
    "avs2": "avs2_v4l2m2m",
    "mjpeg": "mjpeg_v4l2m2m",
}

CUVID_CODEC_NAMES = {
    "h264": "h264_cuvid",
    "hevc": "hevc_cuvid",
    "vp9": "vp9_cuvid",
    "avs2": "avs2_cuvid",
    "mjpeg": "mjpeg_cuvid",
}

def get_input_codec(input_file):
    """获取输入文件的编码格式（小写）"""
    try:
        probe = ffmpeg.probe(input_file, select_streams='v:0', show_entries='stream=codec_name')
        return probe['streams'][0]['codec_name'].lower()
    except Exception:
        return os.path.splitext(input_file)[1].replace('.', '').lower()

def run_multi_instance_benchmark(input_file, codec_type, threads, log_file_path):
    """运行多实例基准测试"""
    if not os.path.isfile(input_file):
        print(f"Error: Input file '{input_file}' does not exist.", file=sys.stderr)
        sys.exit(1)
    
    input_codec = get_input_codec(input_file)
    
    if codec_type == "cuda":
        stream = ffmpeg.input(input_file, hwaccel='cuda')
    elif codec_type == "v4l2":
        decoder = V4L2_CODEC_NAMES.get(input_codec, input_codec)
        stream = ffmpeg.input(input_file, vcodec=decoder)
    elif codec_type == "cuvid":
        decoder = CUVID_CODEC_NAMES.get(input_codec, input_codec)
        stream = ffmpeg.input(input_file, vcodec=decoder)
    elif codec_type == "vaapi":
        stream = ffmpeg.input(input_file, hwaccel='vaapi')
    else:
        print(f"Unknown codec type: {codec_type}", file=sys.stderr)
        sys.exit(1)
    
    cmd = ffmpeg.output(stream, 'pipe:', format='null').global_args('-benchmark', '-y').compile()
    processes = []
    
    print(f"Starting {threads} concurrent ffmpeg processes...", file=sys.stderr)
    
    with open(log_file_path, 'w') as log_f:
        log_f.write(f"# Multi-instance benchmark log\n")
        log_f.write(f"# Command: {' '.join(cmd)}\n")
        log_f.write(f"# Threads: {threads}\n\n")
        
        for i in range(threads):
            p = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=log_f)
            processes.append(p)
        
        for p in processes:
            p.wait()
            
    print("All processes finished.", file=sys.stderr)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Multi-instance FFmpeg Benchmark")
    parser.add_argument("-i", "--input", required=True, help="Input video file")
    parser.add_argument("-c", "--codec", default="cuda", choices=["cuda", "cuvid", "v4l2", "vaapi"], help="Codec type")
    parser.add_argument("-n", "--threads", type=int, default=6, help="Number of concurrent processes")
    parser.add_argument("-o", "--output", required=True, help="Output log file")
    args = parser.parse_args()
    
    run_multi_instance_benchmark(args.input, args.codec, args.threads, args.output)