#!/usr/bin/env python3
"""Expose an NVIDIA GPU's PCI power state as an AppIndicator tray icon.

The state comes from /sys/bus/pci/devices/<address>/power_state, which is a
plain sysfs read and (unlike nvidia-smi) never wakes a suspended GPU.
"""

import argparse
import os
import sys
from pathlib import Path

PCI_DEVICES = Path("/sys/bus/pci/devices")
PCI_IDS = Path("/usr/share/misc/pci.ids")
NVIDIA_VENDOR = "0x10de"
APP_ID = "nvidia-power-tray"

# Adwaita symbolic icons, themed automatically by the shell.
ICONS = {
    "D0": "power-profile-performance-symbolic",
    "D1": "power-profile-balanced-symbolic",
    "D2": "power-profile-balanced-symbolic",
    "D3hot": "power-profile-power-saver-symbolic",
    "D3cold": "power-profile-power-saver-symbolic",
}
UNKNOWN_ICON = "power-profile-balanced-symbolic"


def is_display_controller(device: Path) -> bool:
    """PCI class 0x03xxxx covers VGA, XGA, 3D and other display controllers."""
    try:
        return device.joinpath("class").read_text().strip().startswith("0x03")
    except OSError:
        return False


def find_gpu() -> str:
    """Return the address of the first NVIDIA display controller found."""
    candidates = []
    for device in sorted(PCI_DEVICES.iterdir()):
        try:
            vendor = device.joinpath("vendor").read_text().strip()
        except OSError:
            continue
        if vendor.lower() != NVIDIA_VENDOR or not is_display_controller(device):
            continue
        # A bound nvidia/nouveau driver is the stronger signal; keep it first.
        driver = read_driver(device.name)
        candidates.append((driver is None, device.name))
    if not candidates:
        sys.exit("no NVIDIA display controller found under /sys/bus/pci/devices")
    candidates.sort()
    return candidates[0][1]


def read_attr(address: str, *parts: str) -> str:
    try:
        return PCI_DEVICES.joinpath(address, *parts).read_text().strip()
    except OSError:
        return "unavailable"


def read_driver(address: str):
    try:
        return os.path.basename(os.readlink(PCI_DEVICES / address / "driver"))
    except OSError:
        return None


def describe(address: str) -> str:
    """Human-readable device name, falling back to the raw PCI IDs.

    Deliberately looks the IDs up in pci.ids rather than shelling out to
    lspci: lspci reads the device's config space, which wakes a GPU that had
    reached D3cold — the exact thing this indicator exists to observe.
    """
    vendor = read_attr(address, "vendor").removeprefix("0x")
    device = read_attr(address, "device").removeprefix("0x")
    name = lookup_pci_ids(vendor, device)
    return name or f"{vendor}:{device}"


def lookup_pci_ids(vendor: str, device: str):
    """Resolve "10de", "2d19" to a device name using the pci.ids database."""
    try:
        with PCI_IDS.open(encoding="utf-8", errors="replace") as handle:
            in_vendor = False
            for line in handle:
                if line.startswith("#") or not line.strip():
                    continue
                if not line.startswith("\t"):
                    # Vendor lines are "10de  NVIDIA Corporation".
                    if in_vendor:
                        return None  # left our vendor block without a match
                    in_vendor = line.split(None, 1)[0].lower() == vendor
                elif in_vendor and not line.startswith("\t\t"):
                    # Device lines are "\t2d19  GB206M [GeForce RTX 5060 ...]".
                    code, _, name = line.strip().partition("  ")
                    if code.lower() == device:
                        return name.strip()
    except OSError:
        pass
    return None


class Tray:
    def __init__(self, address: str, interval: int):
        self.address = address
        self.interval = interval
        self.state = None

        from gi.repository import GLib, Gtk

        self.GLib = GLib
        self.Gtk = Gtk
        self.indicator = self.make_indicator()
        self.details = {}
        self.indicator.set_menu(self.make_menu())
        self.refresh()
        GLib.timeout_add_seconds(interval, self.refresh)

    def make_indicator(self):
        import gi

        for namespace in ("AyatanaAppIndicator3", "AppIndicator3"):
            try:
                gi.require_version(namespace, "0.1")
                module = __import__("gi.repository", fromlist=[namespace])
                appindicator = getattr(module, namespace)
                break
            except (ValueError, ImportError):
                continue
        else:
            sys.exit(
                "no AppIndicator GIR found; install gir1.2-ayatanaappindicator3-0.1"
            )

        indicator = appindicator.Indicator.new(
            APP_ID, UNKNOWN_ICON, appindicator.IndicatorCategory.HARDWARE
        )
        indicator.set_status(appindicator.IndicatorStatus.ACTIVE)
        return indicator

    def make_menu(self):
        menu = self.Gtk.Menu()
        for key, label in (
            ("device", "Device"),
            ("address", "Address"),
            ("driver", "Driver"),
            ("power_state", "Power state"),
            ("runtime_status", "Runtime status"),
            ("runtime_pm", "Runtime PM"),
        ):
            item = self.Gtk.MenuItem(label=f"{label}: …")
            item.set_sensitive(False)
            item.show()
            menu.append(item)
            self.details[key] = (label, item)

        separator = self.Gtk.SeparatorMenuItem()
        separator.show()
        menu.append(separator)

        refresh = self.Gtk.MenuItem(label="Refresh now")
        refresh.connect("activate", lambda _item: self.refresh())
        refresh.show()
        menu.append(refresh)

        quit_item = self.Gtk.MenuItem(label="Quit")
        quit_item.connect("activate", lambda _item: self.Gtk.main_quit())
        quit_item.show()
        menu.append(quit_item)
        return menu

    def set_detail(self, key: str, value: str) -> None:
        label, item = self.details[key]
        item.set_label(f"{label}: {value}")

    def refresh(self) -> bool:
        state = read_attr(self.address, "power_state")
        runtime = read_attr(self.address, "power", "runtime_status")
        control = read_attr(self.address, "power", "control")
        driver = read_driver(self.address) or "none"

        if state != self.state:
            self.state = state
            self.indicator.set_icon_full(
                ICONS.get(state, UNKNOWN_ICON), f"GPU power state {state}"
            )
        # The second argument is a sizing hint, not a fallback value.
        self.indicator.set_label(state, "D3cold")
        self.indicator.set_title(f"GPU {self.address}: {state}")

        self.set_detail("address", self.address)
        self.set_detail("driver", driver)
        self.set_detail("power_state", state)
        self.set_detail("runtime_status", runtime)
        self.set_detail("runtime_pm", control)
        return True  # keep the GLib timeout alive

    def run(self) -> None:
        self.set_detail("device", describe(self.address))
        self.Gtk.main()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "-a",
        "--address",
        default=os.environ.get("NVIDIA_POWER_TRAY_ADDRESS"),
        help="PCI address (e.g. 0000:01:00.0); autodetected when omitted",
    )
    parser.add_argument(
        "-i",
        "--interval",
        type=int,
        default=int(os.environ.get("NVIDIA_POWER_TRAY_INTERVAL", "5")),
        help="seconds between sysfs polls (default: 5)",
    )
    parser.add_argument(
        "-p",
        "--print",
        action="store_true",
        dest="print_once",
        help="print the current state and exit, without starting the tray",
    )
    args = parser.parse_args()

    address = args.address or find_gpu()
    if not PCI_DEVICES.joinpath(address, "power_state").exists():
        sys.exit(f"no power_state attribute for PCI device {address}")

    if args.print_once:
        print(f"{address} {read_attr(address, 'power_state')}")
        return

    import gi

    gi.require_version("Gtk", "3.0")
    Tray(address, max(1, args.interval)).run()


if __name__ == "__main__":
    main()
