import os
import subprocess

# Check CUDA installation
print('=== CUDA Installation Check ===')
cuda_path = os.environ.get('CUDA_PATH', 'Not set')
print(f'CUDA_PATH: {cuda_path}')

# Check if CUDA binaries are accessible
try:
    result = subprocess.run(['nvcc', '--version'], capture_output=True, text=True)
    if result.returncode == 0:
        print('✓ nvcc found and working')
        # Extract version from nvcc output
        if "release" in result.stdout:
            version_part = result.stdout.split("release")[1].split(",")[0].strip()
            print(f'nvcc version: {version_part}')
        else:
            print('nvcc version: Unknown')
    else:
        print('✗ nvcc not found or not working')
except Exception as e:
    print(f'✗ nvcc check failed: {e}')

# Check cuDNN installation
print('\n=== cuDNN Installation Check ===')
if cuda_path != 'Not set':
    cudnn_path = os.path.join(cuda_path, 'bin', 'cudnn64_8.dll')
    if os.path.exists(cudnn_path):
        print(f'✓ cuDNN DLL found: {cudnn_path}')
    else:
        print(f'✗ cuDNN DLL not found at: {cudnn_path}')
        print('  Searching for cudnn files...')
        for root, dirs, files in os.walk(cuda_path):
            for file in files:
                if 'cudnn' in file.lower():
                    print(f'    Found: {os.path.join(root, file)}')
else:
    print('✗ CUDA_PATH not set, cannot check cuDNN')

# Check PATH for CUDA
print('\n=== PATH Environment Variable ===')
path_entries = os.environ.get('PATH', '').split(';')
cuda_paths = [p for p in path_entries if 'cuda' in p.lower()]
for path in cuda_paths:
    print(f'  {path}')