#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
Test script to verify NVDA CPU affinity and priority settings
"""

import psutil
import sys

def find_nvda_processes():
    nvda_processes = []
    for proc in psutil.process_iter(['pid', 'name']):
        try:
            name = proc.info['name'].lower()
            if 'nvda' in name:
                nvda_processes.append(proc)
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    return nvda_processes

def get_priority_name(priority_class):
    priority_names = {
        psutil.IDLE_PRIORITY_CLASS: "IDLE",
        psutil.BELOW_NORMAL_PRIORITY_CLASS: "BELOW_NORMAL",
        psutil.NORMAL_PRIORITY_CLASS: "NORMAL",
        psutil.ABOVE_NORMAL_PRIORITY_CLASS: "ABOVE_NORMAL",
        psutil.HIGH_PRIORITY_CLASS: "HIGH",
        psutil.REALTIME_PRIORITY_CLASS: "REALTIME"
    }
    return priority_names.get(priority_class, f"UNKNOWN ({priority_class})")

def main():
    print("NVDA CPU Affinity and Priority Verification")
    print()

    cpu_count = psutil.cpu_count(logical=True)
    print(f"System Information:")
    print(f"  Total CPU cores (logical): {cpu_count}")
    print(f"  Available cores: 0-{cpu_count - 1}")
    print()

    nvda_processes = find_nvda_processes()

    if not nvda_processes:
        print("ERROR: No NVDA processes found!")
        print("Make sure NVDA is running.")
        sys.exit(1)

    print(f"Found {len(nvda_processes)} NVDA process(es):")
    print()

    for proc in nvda_processes:
        try:
            pid = proc.pid
            name = proc.name()

            try:
                affinity = proc.cpu_affinity()
            except (psutil.AccessDenied, AttributeError) as e:
                affinity = f"ERROR: {e}"

            try:
                priority = proc.nice()
                priority_name = get_priority_name(priority)
            except (psutil.AccessDenied, AttributeError) as e:
                priority = f"ERROR: {e}"
                priority_name = ""

            print(f"Process: {name}")
            print(f"  PID: {pid}")
            print(f"  CPU Affinity: {affinity}")
            if isinstance(affinity, list):
                if len(affinity) == cpu_count:
                    print(f"    -> Using ALL cores")
                else:
                    print(f"    -> Restricted to {len(affinity)} core(s)")
            print(f"  Priority: {priority_name}")
            print()

        except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
            print(f"  ERROR: Could not access process: {e}")
            print()



if __name__ == "__main__":
    main()