#!/usr/bin/env python3
#
# process.py
# Shows processes using the Nvidia GPU
#
# Copyright (c) 2026 Erich Eickmeyer
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import glob
import os
import re
import signal
import subprocess
import sys
import xml.etree.ElementTree as ET

from PyQt6.QtWidgets import (
    QApplication, QDialog, QVBoxLayout, QLabel,
    QTreeWidget, QTreeWidgetItem, QDialogButtonBox,
    QMessageBox, QHeaderView,
)
from PyQt6.QtCore import Qt


def _nvidia_smi_processes():
    """Parse nvidia-smi XML output for running processes."""
    try:
        result = subprocess.run(
            ["nvidia-smi", "-q", "-x"],
            capture_output=True, text=True,
        )
    except FileNotFoundError:
        return []

    if result.returncode != 0:
        return []

    processes = []
    try:
        root = ET.fromstring(result.stdout)
        for gpu in root.iter("gpu"):
            for proc_info in gpu.iter("process_info"):
                pid_el = proc_info.find("pid")
                name_el = proc_info.find("process_name")
                mem_el = proc_info.find("used_memory")
                if pid_el is not None:
                    processes.append((
                        pid_el.text.strip() if pid_el.text else "",
                        name_el.text.strip() if name_el is not None and name_el.text else "",
                        mem_el.text.strip() if mem_el is not None and mem_el.text else "",
                    ))
    except ET.ParseError:
        pass

    return processes


def _find_nvidia_render_nodes():
    """Return DRI render-node paths whose device belongs to NVIDIA (0x10de)."""
    nodes = []
    for path in sorted(glob.glob("/dev/dri/renderD*")):
        name = os.path.basename(path)
        vendor_file = f"/sys/class/drm/{name}/device/vendor"
        try:
            vendor = open(vendor_file).read().strip()
        except OSError:
            continue
        if vendor == "0x10de":
            nodes.append(path)
    return nodes


def _cmdline_for_pid(pid):
    """Return the process name for a PID via /proc."""
    try:
        cmdline = open(f"/proc/{pid}/cmdline", "rb").read()
        if cmdline:
            return cmdline.split(b"\x00")[0].decode(errors="replace")
        return open(f"/proc/{pid}/comm").read().strip()
    except OSError:
        return ""


def _device_node_processes():
    """Detect processes using NVIDIA GPU via device-node file descriptors.

    On modern Wayland/DRI systems nvidia-smi may not report GPU clients.
    Fall back to scanning /proc/*/fd for open handles on NVIDIA device
    nodes (/dev/nvidia*, /dev/dri/renderD* with vendor 0x10de).
    """
    nvidia_devs = set(glob.glob("/dev/nvidia[0-9]*"))
    nvidia_devs |= set(glob.glob("/dev/nvidiactl"))
    nvidia_devs |= set(_find_nvidia_render_nodes())

    if not nvidia_devs:
        return []

    # Resolve to real paths for comparison
    nvidia_inodes = {}
    for dev in nvidia_devs:
        try:
            st = os.stat(dev)
            nvidia_inodes[st.st_ino] = dev
        except OSError:
            continue

    seen_pids = set()
    processes = []

    for fd_dir in glob.glob("/proc/[0-9]*/fd"):
        pid_str = fd_dir.split("/")[2]
        if pid_str in seen_pids:
            continue
        try:
            for entry in os.listdir(fd_dir):
                try:
                    target = os.readlink(f"{fd_dir}/{entry}")
                except OSError:
                    continue
                if target in nvidia_devs:
                    seen_pids.add(pid_str)
                    name = _cmdline_for_pid(pid_str)
                    processes.append((pid_str, name, ""))
                    break
        except PermissionError:
            continue

    return processes


def get_gpu_processes():
    """Return list of (pid, name, memory) tuples for NVIDIA GPU processes."""
    processes = _nvidia_smi_processes()
    if processes:
        return processes
    return _device_node_processes()


class ProcessDialog(QDialog):
    def __init__(self, processes):
        super().__init__()
        self.setWindowTitle("Plasma Optimus — GPU Processes")
        self.resize(600, 400)

        layout = QVBoxLayout(self)
        layout.addWidget(QLabel("Running Processes:"))

        self.tree = QTreeWidget()
        self.tree.setHeaderLabels(["PID", "Process", "Memory"])
        self.tree.header().setSectionResizeMode(
            QHeaderView.ResizeMode.ResizeToContents
        )
        self.tree.setRootIsDecorated(False)

        for pid, name, mem in processes:
            QTreeWidgetItem(self.tree, [pid, name, mem])

        layout.addWidget(self.tree)

        buttons = QDialogButtonBox()
        self.kill_btn = buttons.addButton(
            "Kill Selected", QDialogButtonBox.ButtonRole.ActionRole
        )
        self.reload_btn = buttons.addButton(
            "Reload Selected", QDialogButtonBox.ButtonRole.ActionRole
        )
        self.close_btn = buttons.addButton(QDialogButtonBox.StandardButton.Close)

        self.kill_btn.clicked.connect(self.kill_process)
        self.reload_btn.clicked.connect(self.reload_process)
        self.close_btn.clicked.connect(self.accept)

        layout.addWidget(buttons)

    def _selected_pid(self):
        items = self.tree.selectedItems()
        if not items:
            QMessageBox.information(self, "Plasma Optimus", "No process selected.")
            return None
        return int(items[0].text(0))

    def kill_process(self):
        pid = self._selected_pid()
        if pid is None:
            return
        reply = QMessageBox.warning(
            self, "Kill Process",
            f"Kill PID {pid}? This can cause issues in the application.",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No,
        )
        if reply == QMessageBox.StandardButton.Yes:
            os.kill(pid, signal.SIGKILL)
            self.accept()

    def reload_process(self):
        pid = self._selected_pid()
        if pid is None:
            return
        os.kill(pid, signal.SIGHUP)
        self.accept()


def main():
    app = QApplication(sys.argv)
    processes = get_gpu_processes()

    if not processes:
        QMessageBox.information(
            None, "Plasma Optimus",
            "No processes are currently using the Nvidia GPU.",
        )
        sys.exit(0)

    dialog = ProcessDialog(processes)
    dialog.exec()
    sys.exit(0)


if __name__ == "__main__":
    main()
