#!/usr/bin/env python3
"""nordvpn-tray: system tray icon for the nordvpn-wg connections, because SteamOS on the
Frame ships no Plasma network applet. Left- or right-click the icon for a menu of
locations (the status is re-read each time it opens); the icon turns green while
connected."""
import os
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, os.environ.get("FRAME_TWEAKS_PYLIB") or str(Path.home() / ".local/share/frame-tweaks/i18n"))   # (lib/common.sh install_i18n)
try:
    from i18n import _, N_, ngettext
except ImportError:                       # (not installed by Frame Tweaks: English)
    def _(text):
        return text

    def ngettext(singular, plural, n):
        return singular if n == 1 else plural
    N_ = _

# Under native Wayland a parentless popup menu can't be positioned; XWayland can.
os.environ.setdefault("QT_QPA_PLATFORM", "xcb")

from PyQt6.QtCore import QProcess, Qt, QTimer  # noqa: E402
from PyQt6.QtCore import QPoint  # noqa: E402
from PyQt6.QtGui import QAction, QActionGroup, QColor, QIcon, QPainter, QPixmap  # noqa: E402
from PyQt6.QtWidgets import QApplication, QMenu, QSystemTrayIcon  # noqa: E402

NWG = os.path.expanduser("~/.local/bin/nordvpn-wg")
PREFIX = "nordvpn-"


def icon(color):
    """A simple shield-ish dot, so we don't depend on the icon theme."""
    pm = QPixmap(64, 64)
    pm.fill(Qt.GlobalColor.transparent)
    p = QPainter(pm)
    p.setRenderHint(QPainter.RenderHint.Antialiasing)
    p.setBrush(QColor(color))
    p.setPen(QColor("#202020"))
    p.drawRoundedRect(8, 6, 48, 52, 14, 14)
    p.setPen(QColor("white"))
    f = p.font()
    f.setPixelSize(26)
    f.setBold(True)
    p.setFont(f)
    p.drawText(pm.rect(), Qt.AlignmentFlag.AlignCenter, "N")
    p.end()
    return QIcon(pm)


def panel_anchor():
    """(edge, thickness, screen width, screen height) of Plasma's first panel, so the
    left-click menu can open next to the tray. (XWayland can't see the pointer while it
    is over the Wayland panel, so QCursor.pos() would be stale.)"""
    script = ('var p=panels()[0]; var g=screenGeometry(p.screen);'
              'print(p.location+" "+p.height+" "+g.width+" "+g.height)')
    try:
        out = subprocess.run(["qdbus6", "org.kde.plasmashell", "/PlasmaShell",
                              "org.kde.PlasmaShell.evaluateScript", script],
                             capture_output=True, text=True, timeout=5).stdout.split()
        return out[0], int(out[1]), int(out[2]), int(out[3])
    except Exception:
        return "bottom", 48, 0, 0


def nm_connections():
    """Our WireGuard profiles' names (sorted); the active one, the one NetworkManager connects at
    startup (its autoconnect) and the one used last (each a name, or None)."""
    out = subprocess.run(["nmcli", "-t", "-f", "NAME,TYPE,DEVICE,AUTOCONNECT,TIMESTAMP", "connection", "show"],
                         capture_output=True, text=True).stdout
    profiles, active, startup, used = [], None, [], {}
    for line in out.splitlines():
        parts = line.replace("\\:", "\0").split(":")
        if len(parts) < 5 or parts[1] != "wireguard":
            continue
        name = parts[0].replace("\0", ":")
        if not name.startswith("NordVPN "):
            continue
        profiles.append(name)
        if parts[2].startswith(PREFIX):
            active = name
        if parts[3] == "yes":
            startup.append(name)
        used[name] = int(parts[4]) if parts[4].isdigit() else 0
    last = max(used, key=used.get) if any(used.values()) else None    # (0: not since it was made)
    return sorted(profiles), active, min(startup, default=None), last


class Tray:
    def __init__(self, app):
        self.app = app
        self.on, self.off, self.busy = icon("#2e9d4f"), icon("#7a7a7a"), icon("#d9a400")
        self.tray = QSystemTrayIcon(self.off)
        self.menu = QMenu()
        self.tray.setContextMenu(self.menu)
        self.tray.activated.connect(self.on_activated)
        # Re-read the connection state whenever the menu opens (right-click included).
        self.menu.aboutToShow.connect(self.refresh)
        self.proc = None
        self.active = self.startup = self.last = None
        self.profiles = []
        self.anchor = panel_anchor()
        self.refresh()
        self.tray.show()
        QTimer(app, interval=5000, timeout=self.poll).start()

    def poll(self):
        """Periodic update: icon/tooltip always, menu only while it's closed
        (rebuilding an open menu would close it)."""
        self.update_state()
        if not self.menu.isVisible():
            self.build_menu()

    def menu_pos(self):
        edge, thick, w, h = self.anchor
        if not w:
            w, h = self.app.primaryScreen().size().width(), self.app.primaryScreen().size().height()
        size = self.menu.sizeHint()
        x = w - size.width() - 8
        y = thick + 4 if edge == "top" else h - thick - size.height() - 4
        if edge == "left":
            x, y = thick + 4, h - size.height() - 8
        elif edge == "right":
            x, y = w - thick - size.width() - 4, h - size.height() - 8
        return QPoint(max(0, x), max(0, y))

    def on_activated(self, reason):
        if reason == QSystemTrayIcon.ActivationReason.Trigger:   # left click: same menu
            self.refresh()
            self.menu.popup(self.menu_pos())

    def run(self, *args, note=None):
        if self.proc and self.proc.state() != QProcess.ProcessState.NotRunning:
            return
        self.tray.setIcon(self.busy)
        self.proc = QProcess()
        self.proc.finished.connect(lambda *_: self.done(note))
        self.proc.start(NWG, list(args))

    def done(self, note):
        out = bytes(self.proc.readAllStandardOutput()).decode().strip()
        err = bytes(self.proc.readAllStandardError()).decode().strip()
        self.refresh()
        msg = err or (note and out) or ""
        if msg:
            self.tray.showMessage("NordVPN", msg[-300:], self.on if self.active else self.off, 6000)

    def refresh(self):
        self.update_state()
        self.build_menu()

    def update_state(self):
        self.profiles, self.active, self.startup, self.last = nm_connections()
        busy = self.proc and self.proc.state() != QProcess.ProcessState.NotRunning
        self.tray.setIcon(self.busy if busy else self.on if self.active else self.off)
        self.tray.setToolTip(_("NordVPN: {location}").format(location=self.active[8:]) if self.active else
                             _("NordVPN: disconnected"))

    def build_menu(self):
        profiles = self.profiles
        self.menu.clear()
        head = self.menu.addAction(_("Connected: {location}").format(location=self.active[8:]) if self.active
                                   else _("Disconnected"))
        head.setEnabled(False)
        self.menu.addSeparator()
        group = QActionGroup(self.menu)
        for name in profiles:
            a = QAction(name[8:], self.menu, checkable=True, checked=(name == self.active))
            code = self.code_for(name)
            a.triggered.connect(lambda _=False, c=code: self.run("up", c))
            group.addAction(a)
            self.menu.addAction(a)
        if not profiles:
            self.menu.addAction(_("No locations: run 'nordvpn-wg setup'")).setEnabled(False)
        self.menu.addSeparator()
        dis = self.menu.addAction(_("Disconnect"))
        dis.setEnabled(bool(self.active))
        dis.triggered.connect(lambda: self.run("down"))
        self.menu.addAction(_("Show public IP")).triggered.connect(lambda: self.run("status", note=True))
        self.menu.addAction(_("Pick least-loaded servers")).triggered.connect(
            lambda: self.run("refresh", note=True))
        # NetworkManager connects it by itself as the Frame starts, before desktop mode (nordvpn-wg
        # autoconnect). Ticked: the one that does; unticked: the one ticking it would choose.
        start = self.startup or self.active or self.last or (profiles[0] if profiles else None)
        if start:
            self.menu.addSeparator()
            auto = QAction(_("Connect at startup ({location})").format(location=start[8:]), self.menu,
                           checkable=True, checked=bool(self.startup))
            auto.triggered.connect(lambda on, c=self.code_for(start): self.run("autoconnect", c if on else "off", note=True))
            self.menu.addAction(auto)
        self.menu.addSeparator()
        self.menu.addAction(_("Quit tray icon")).triggered.connect(self.app.quit)

    _codes = {}

    def code_for(self, name):
        if name in self._codes:
            return self._codes[name]
        iface = subprocess.run(["nmcli", "-g", "connection.interface-name", "connection", "show", name],
                               capture_output=True, text=True).stdout.strip()
        self._codes[name] = iface[len(PREFIX):].upper()
        return self._codes[name]


def desktop_session_env():
    """Desktop mode's environment (its display, D-Bus and session), from its Plasma shell; None
    when desktop mode isn't running."""
    uid = os.getuid()
    for p in Path("/proc").iterdir():
        if not p.name.isdigit():
            continue
        try:
            if p.stat().st_uid != uid or (p / "comm").read_text().strip() != "plasmashell":
                continue
            raw = (p / "environ").read_bytes().decode(errors="replace").split("\0")
        except OSError:
            continue
        return dict(kv.split("=", 1) for kv in raw if "=" in kv)
    return None


def main():
    # The icon lives in desktop mode's system tray. Started from VR (the Frame Tweaks window, the
    # "+" menu), it would run in the Steam session, which has no tray: no icon, and it would hold
    # the lock below, so desktop mode's own copy quietly quit. So move into desktop mode first.
    if os.environ.get("XDG_CURRENT_DESKTOP") != "KDE":
        env = desktop_session_env()
        if env is None:
            subprocess.run(["notify-send", "-a", "NordVPN", "NordVPN tray",
                            _("Its icon lives in desktop mode's system tray: open desktop mode, and it starts there.")])
            sys.exit(0)
        os.execve(sys.executable, [sys.executable, os.path.abspath(sys.argv[0])], env)
    # One tray only. Autostart starts it, and KDE's session restore used to start a second copy
    # at login; whichever comes second finds the lock held and quietly exits.
    import fcntl
    lock = open(f"/run/user/{os.getuid()}/nordvpn-tray.lock", "w")
    try:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        sys.exit(0)
    main.lock = lock                     # held until the process exits
    app = QApplication(sys.argv)
    app.setApplicationName("nordvpn-tray")
    app.setQuitOnLastWindowClosed(False)
    if not QSystemTrayIcon.isSystemTrayAvailable():
        sys.exit(_("nordvpn-tray: no system tray available"))
    Tray(app)
    sys.exit(app.exec())


if __name__ == "__main__":
    main()
