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

def run_ffmpeg_command(cmd):
    """运行命令并返回 stderr 日志"""
    try:
        result = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
        return result.stderr
    except Exception as e:
        return f"Error: {e}"

def main():
    input_file = "./sample_1920x1080_30fps_8bit.h264"
    if not os.path.exists(input_file):
        print(f"Error: File {input_file} not found")
        sys.exit(1)

    print(f"--- Starting CUDA Benchmark for {input_file} ---\n")

    # 1. CUDA 解码测试 (hwaccel cuda)
    print("[1] Testing CUDA Decode...")
    try:
        stream = ffmpeg.input(input_file, hwaccel='cuda')
        cmd = ffmpeg.output(stream, 'pipe:', format='null').global_args('-benchmark', '-y').compile()
        log_content = run_ffmpeg_command(cmd)
        print(f"    Decode Success\n")
    except Exception as e:
        print(f"    Decode Failed: {e}\n")

    # 2. CUDA 编码测试 (h264_nvenc)
    print("[2] Testing CUDA Encode (H264 NVENC)...")
    output_file = "output_cuda_temp.mp4"
    try:
        stream = ffmpeg.input(input_file)
        cmd = ffmpeg.output(stream, output_file, vcodec='h264_nvenc').global_args('-benchmark', '-y').compile()
        log_content = run_ffmpeg_command(cmd)
        print(f"    Encode Success\n")
    except Exception as e:
        print(f"    Encode Failed: {e}\n")
    finally:
        # 清理临时文件
        if os.path.exists(output_file):
            os.remove(output_file)

    print("--- Benchmark Finished ---")

if __name__ == "__main__":
    main()