import subprocess
import sys

def check_gpu():
    print("--- Checking for NVIDIA GPU with nvidia-smi ---")
    try:
        result_smi = subprocess.run(
            ["nvidia-smi"], 
            capture_output=True, 
            text=True, 
            check=True
        )
        print("SUCCESS: nvidia-smi found.")
        print("nvidia-smi stdout:")
        print(result_smi.stdout)
    except FileNotFoundError:
        print("ERROR: nvidia-smi command not found. Make sure NVIDIA drivers are installed and nvidia-smi is in your system's PATH.")
    except subprocess.CalledProcessError as e:
        print(f"ERROR: nvidia-smi failed with exit code {e.returncode}.")
        print("nvidia-smi stderr:")
        print(e.stderr)

    print("\n--- Checking ffmpeg for CUDA support ---")
    try:
        result_ffmpeg = subprocess.run(
            ["ffmpeg", "-hwaccels"], 
            capture_output=True, 
            text=True, 
            check=True
        )
        print("SUCCESS: ffmpeg found.")
        print("ffmpeg stdout:")
        print(result_ffmpeg.stdout)
        if "cuda" in result_ffmpeg.stdout:
            print("\nCONCLUSION: ffmpeg appears to have CUDA support.")
        else:
            print("\nCONCLUSION: ffmpeg does NOT appear to have CUDA support. It needs to be compiled with CUDA enabled.")
    except FileNotFoundError:
        print("ERROR: ffmpeg command not found. Make sure ffmpeg is installed and is in your system's PATH.")
    except subprocess.CalledProcessError as e:
        print(f"ERROR: ffmpeg failed with exit code {e.returncode}.")
        print("ffmpeg stderr:")
        print(e.stderr)

if __name__ == "__main__":
    check_gpu()