# -*- coding: UTF-8 -*-
# CPU Priority Manager for NVDA
# Allows setting CPU affinity and process priority for NVDA processes

import os
import sys

# Add psutil to the path
addonDir = os.path.dirname(__file__)
sys.path.insert(0, addonDir)

try:
	import psutil
except ImportError:
	psutil = None

import globalPluginHandler
import gui
from gui import guiHelper, nvdaControls
from gui.settingsDialogs import SettingsPanel
import config
import wx
import addonHandler
from logHandler import log

addonHandler.initTranslation()

confspec = {
	"enabled": "boolean(default=False)",
	"cpuCores": "string(default='')",
	"priorityLevel": "string(default='HIGH')",
}

config.conf.spec["cpuPriority"] = confspec


class CPUPrioritySettingsPanel(SettingsPanel):
	# Translators: Title of the CPU Priority settings panel
	title = _("CPU Priority Manager")

	def makeSettings(self, settingsSizer):
		sHelper = guiHelper.BoxSizerHelper(self, sizer=settingsSizer)

		# Translators: Label for checkbox to enable/disable CPU priority management
		self.enabledCheckbox = sHelper.addItem(
			wx.CheckBox(self, label=_("&Enable CPU priority management"))
		)
		self.enabledCheckbox.SetValue(config.conf["cpuPriority"]["enabled"])
		self.enabledCheckbox.Bind(wx.EVT_CHECKBOX, self.onEnableToggle)

		# Translators: Label for text field to specify CPU cores
		cpuCoresLabel = _("CPU &cores (comma-separated, e.g., 4,5,6,7):")
		self.cpuCoresEdit = sHelper.addLabeledControl(cpuCoresLabel, wx.TextCtrl)
		self.cpuCoresEdit.SetValue(config.conf["cpuPriority"]["cpuCores"])

		# Translators: Label for priority level dropdown
		priorityLabel = _("Process &priority:")
		self.priorityChoice = sHelper.addLabeledControl(
			priorityLabel,
			wx.Choice,
			choices=[
				# Translators: Normal priority level
				_("Normal"),
				# Translators: Above normal priority level
				_("Above Normal"),
				# Translators: High priority level
				_("High"),
				# Translators: Realtime priority level
				_("Realtime (Use with caution)")
			]
		)

		# Map config values to choice indices
		priorityMap = {
			"NORMAL": 0,
			"ABOVE_NORMAL": 1,
			"HIGH": 2,
			"REALTIME": 3
		}
		currentPriority = config.conf["cpuPriority"]["priorityLevel"]
		self.priorityChoice.SetSelection(priorityMap.get(currentPriority, 2))

		if psutil:
			try:
				cpu_count = psutil.cpu_count(logical=True)
				# Translators: Information about available CPU cores
				infoText = _("Your system has {count} logical CPU cores (0-{max}).").format(
					count=cpu_count,
					max=cpu_count - 1
				)
				sHelper.addItem(wx.StaticText(self, label=infoText))
			except Exception:
				pass

		# Translators: Warning about realtime priority
		warningText = _(
			"Warning: Setting realtime priority may make your system unstable if NVDA uses too much CPU. "
			"Use high priority for most cases. Changes require NVDA restart to take effect."
		)
		warningLabel = wx.StaticText(self, label=warningText)
		warningLabel.Wrap(self.GetSize()[0] - 20)
		sHelper.addItem(warningLabel)

		self.onEnableToggle(None)

	def onEnableToggle(self, evt):
		"""Enable or disable controls based on enabled checkbox"""
		enabled = self.enabledCheckbox.GetValue()
		self.cpuCoresEdit.Enable(enabled)
		self.priorityChoice.Enable(enabled)

	def onSave(self):
		"""Save settings to configuration"""
		config.conf["cpuPriority"]["enabled"] = self.enabledCheckbox.GetValue()
		config.conf["cpuPriority"]["cpuCores"] = self.cpuCoresEdit.GetValue().strip()

		# Map choice indices back to config values
		priorityMap = {
			0: "NORMAL",
			1: "ABOVE_NORMAL",
			2: "HIGH",
			3: "REALTIME"
		}
		config.conf["cpuPriority"]["priorityLevel"] = priorityMap[self.priorityChoice.GetSelection()]


class GlobalPlugin(globalPluginHandler.GlobalPlugin):

	def __init__(self, *args, **kwargs):
		super(GlobalPlugin, self).__init__(*args, **kwargs)

		if psutil is None:
			log.error("CPU Priority Manager: psutil library not available")
			return

		gui.settingsDialogs.NVDASettingsDialog.categoryClasses.append(CPUPrioritySettingsPanel)

		if config.conf["cpuPriority"]["enabled"]:
			self.applyCPUSettings()

	def terminate(self):
		super(GlobalPlugin, self).terminate()

		try:
			gui.settingsDialogs.NVDASettingsDialog.categoryClasses.remove(CPUPrioritySettingsPanel)
		except (ValueError, AttributeError):
			pass

	def applyCPUSettings(self):
		if psutil is None:
			log.error("CPU Priority Manager: psutil not available, cannot apply settings")
			return

		try:
			current_process = psutil.Process()

			self.applyToProcess(current_process)

			try:
				children = current_process.children(recursive=True)
				for child in children:
					try:
						self.applyToProcess(child)
					except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
						log.debug(f"CPU Priority Manager: Could not apply settings to child process {child.pid}: {e}")
			except Exception as e:
				log.error(f"CPU Priority Manager: Error getting child processes: {e}")

			log.info("CPU Priority Manager: Settings applied successfully")

		except Exception as e:
			log.error(f"CPU Priority Manager: Error applying settings: {e}")

	def applyToProcess(self, process):
		cpuCoresStr = config.conf["cpuPriority"]["cpuCores"]
		if cpuCoresStr:
			try:
				# Parse CPU cores from comma-separated string
				cpuCores = [int(core.strip()) for core in cpuCoresStr.split(',') if core.strip()]
				if cpuCores:
					process.cpu_affinity(cpuCores)
					log.debug(f"CPU Priority Manager: Set CPU affinity for process {process.pid} to cores {cpuCores}")
			except ValueError as e:
				log.error(f"CPU Priority Manager: Invalid CPU cores format: {e}")
			except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
				log.debug(f"CPU Priority Manager: Could not set CPU affinity for process {process.pid}: {e}")
			except Exception as e:
				log.error(f"CPU Priority Manager: Error setting CPU affinity: {e}")

		priorityLevel = config.conf["cpuPriority"]["priorityLevel"]
		try:
			priorityClass = {
				"NORMAL": psutil.NORMAL_PRIORITY_CLASS,
				"ABOVE_NORMAL": psutil.ABOVE_NORMAL_PRIORITY_CLASS,
				"HIGH": psutil.HIGH_PRIORITY_CLASS,
				"REALTIME": psutil.REALTIME_PRIORITY_CLASS
			}.get(priorityLevel, psutil.HIGH_PRIORITY_CLASS)

			process.nice(priorityClass)
			log.debug(f"CPU Priority Manager: Set priority for process {process.pid} to {priorityLevel}")
		except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
			log.debug(f"CPU Priority Manager: Could not set priority for process {process.pid}: {e}")
		except Exception as e:
			log.error(f"CPU Priority Manager: Error setting priority: {e}")