#!/usr/bin/env python3
"""Own org.freedesktop.ScreenSaver and turn inhibit requests into xautolock -disable/-enable.

While an inhibit is active, X screen blanking is also paused (like mpv does).
"""

import argparse
import ctypes
import itertools
import signal
import subprocess
import sys

from gi.repository import Gio, GLib

try:
    from gi.repository import GLibUnix  # GLib >= 2.80
    signal_add = GLibUnix.signal_add
except ImportError:
    signal_add = GLib.unix_signal_add

BUS_NAME = "org.freedesktop.ScreenSaver"
OBJECT_PATHS = ("/org/freedesktop/ScreenSaver", "/ScreenSaver")  # Vivaldi uses the first, Zen the second
EXIT_NAME_TAKEN = 2

INTROSPECTION = """
<node>
  <interface name="org.freedesktop.ScreenSaver">
    <method name="Inhibit">
      <arg name="application_name" type="s" direction="in"/>
      <arg name="reason_for_inhibit" type="s" direction="in"/>
      <arg name="cookie" type="u" direction="out"/>
    </method>
    <method name="UnInhibit">
      <arg name="cookie" type="u" direction="in"/>
    </method>
    <method name="SimulateUserActivity"/>
    <method name="Lock"/>
    <method name="GetActive">
      <arg type="b" direction="out"/>
    </method>
    <method name="SetActive">
      <arg type="b" direction="in"/>
      <arg type="b" direction="out"/>
    </method>
    <method name="GetActiveTime">
      <arg type="u" direction="out"/>
    </method>
    <method name="GetSessionIdleTime">
      <arg type="u" direction="out"/>
    </method>
  </interface>
</node>
"""


def log(msg):
    print(msg, flush=True)


def xautolock(arg):
    try:
        subprocess.run(["xautolock", arg], check=True, timeout=10)
        log(f"xautolock {arg}")
    except (OSError, subprocess.SubprocessError) as e:
        log(f"xautolock {arg} failed: {e}")


class XBlanking:
    """Pause X screen blanking with XScreenSaverSuspend (MIT-SCREEN-SAVER extension).

    Unlike `xset s off -dpms` this changes no global setting: the X server ends the
    pause by itself when this client releases it or its connection closes (exit, crash).
    """

    def __init__(self):
        self.dpy = None
        try:
            self.x11 = ctypes.CDLL("libX11.so.6")
            self.xss = ctypes.CDLL("libXss.so.1")
        except OSError as e:
            log(f"Cannot load X11 libraries ({e}); X blanking will not be paused")
            return
        self.x11.XOpenDisplay.argtypes = [ctypes.c_char_p]
        self.x11.XOpenDisplay.restype = ctypes.c_void_p
        self.x11.XFlush.argtypes = [ctypes.c_void_p]
        self.xss.XScreenSaverQueryExtension.argtypes = [
            ctypes.c_void_p, ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_int)]
        self.xss.XScreenSaverSuspend.argtypes = [ctypes.c_void_p, ctypes.c_int]

        dpy = self.x11.XOpenDisplay(None)
        if not dpy:
            log("Cannot open the X display; X blanking will not be paused")
            return
        event_base, error_base = ctypes.c_int(), ctypes.c_int()
        if not self.xss.XScreenSaverQueryExtension(dpy, ctypes.byref(event_base), ctypes.byref(error_base)):
            log("X server lacks the MIT-SCREEN-SAVER extension; X blanking will not be paused")
            return
        self.dpy = dpy

    def suspend(self, on):
        if self.dpy:
            self.xss.XScreenSaverSuspend(self.dpy, int(on))
            self.x11.XFlush(self.dpy)
            log(f"X blanking {'paused' if on else 'resumed'}")


class Inhibitor:
    def __init__(self, ignored_reasons, release_delay, blanking):
        self.ignored_reasons = ignored_reasons
        self.blanking = blanking
        self.release_delay = release_delay
        self.cookies = {}  # cookie -> (sender, application_name, reason, counts)
        self.next_cookie = itertools.count(1)
        self.proc_names = {}  # sender -> "comm[pid]", for logging only
        self.disabled = False
        self.release_timer = 0
        self.conn = None

    # --- lock state -------------------------------------------------------

    def active_count(self):
        return sum(1 for *_, counts in self.cookies.values() if counts)

    def update(self):
        """Run xautolock only on 0->1 and 1->0 transitions; delay re-enabling."""
        active = self.active_count()
        if active and self.release_timer:
            GLib.source_remove(self.release_timer)
            self.release_timer = 0
        if active and not self.disabled:
            self.block()
        elif not active and self.disabled and not self.release_timer:
            # Browsers send quick Inhibit/UnInhibit bursts when a video starts.
            self.release_timer = GLib.timeout_add_seconds(self.release_delay, self.on_release_timeout)

    def on_release_timeout(self):
        self.release_timer = 0
        if not self.active_count() and self.disabled:
            self.unblock()
        return GLib.SOURCE_REMOVE

    def block(self):
        xautolock("-disable")
        self.blanking.suspend(True)
        self.disabled = True

    def unblock(self):
        xautolock("-enable")
        self.blanking.suspend(False)
        self.disabled = False

    def shutdown(self):
        if self.release_timer:
            GLib.source_remove(self.release_timer)
            self.release_timer = 0
        if self.disabled:
            self.unblock()

    # --- D-Bus ------------------------------------------------------------

    def proc_name(self, sender):
        if sender not in self.proc_names:
            try:
                pid = self.conn.call_sync(
                    "org.freedesktop.DBus", "/org/freedesktop/DBus", "org.freedesktop.DBus",
                    "GetConnectionUnixProcessID", GLib.Variant("(s)", (sender,)),
                    GLib.VariantType("(u)"), Gio.DBusCallFlags.NONE, 1000, None,
                ).unpack()[0]
                with open(f"/proc/{pid}/comm") as f:
                    self.proc_names[sender] = f"{f.read().strip()}[{pid}]"
            except (GLib.Error, OSError):
                self.proc_names[sender] = "?"
        return f"{sender} {self.proc_names[sender]}"

    def inhibit(self, sender, app, reason):
        cookie = next(self.next_cookie)
        counts = reason not in self.ignored_reasons
        self.cookies[cookie] = (sender, app, reason, counts)
        note = "" if counts else " (reason ignored, does not block locking)"
        log(f"Inhibit {cookie}: '{app}' '{reason}' from {self.proc_name(sender)}{note}")
        self.update()
        return cookie

    def uninhibit(self, sender, cookie):
        entry = self.cookies.get(cookie)
        if entry is None or entry[0] != sender:
            log(f"UnInhibit {cookie} from {self.proc_name(sender)}: unknown cookie, ignored")
            return
        del self.cookies[cookie]
        _, app, reason, _ = entry
        log(f"UnInhibit {cookie}: '{app}' '{reason}'")
        self.update()

    def release_sender(self, sender):
        gone = [c for c, (s, *_) in self.cookies.items() if s == sender]
        for c in gone:
            _, app, reason, _ = self.cookies.pop(c)
            log(f"Released {c}: '{app}' '{reason}' ({sender} left the bus)")
        self.proc_names.pop(sender, None)
        if gone:
            self.update()

    def on_method_call(self, conn, sender, path, iface, method, params, invocation):
        if method == "Inhibit":
            app, reason = params.unpack()
            invocation.return_value(GLib.Variant("(u)", (self.inhibit(sender, app, reason),)))
        elif method == "UnInhibit":
            self.uninhibit(sender, params.unpack()[0])
            invocation.return_value(None)
        elif method == "SimulateUserActivity":
            log(f"SimulateUserActivity from {self.proc_name(sender)} (not implemented)")
            invocation.return_value(None)
        elif method == "Lock":
            log(f"Lock from {self.proc_name(sender)}")
            xautolock("-locknow")
            invocation.return_value(None)
        elif method in ("GetActive", "SetActive"):
            invocation.return_value(GLib.Variant("(b)", (False,)))
        elif method in ("GetActiveTime", "GetSessionIdleTime"):
            invocation.return_value(GLib.Variant("(u)", (0,)))

    def on_name_owner_changed(self, conn, sender, path, iface, signal_name, params):
        name, _old, new = params.unpack()
        if not new and name.startswith(":"):
            self.release_sender(name)

    def on_bus_acquired(self, conn, name):
        self.conn = conn
        iface = Gio.DBusNodeInfo.new_for_xml(INTROSPECTION).interfaces[0]
        for path in OBJECT_PATHS:
            conn.register_object(path, iface, self.on_method_call, None, None)
        conn.signal_subscribe(
            "org.freedesktop.DBus", "org.freedesktop.DBus", "NameOwnerChanged",
            "/org/freedesktop/DBus", None, Gio.DBusSignalFlags.NONE, self.on_name_owner_changed,
        )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "-i", "--ignore-reason", action="append", metavar="REASON",
        help="inhibit reason that does not block locking (repeatable; default: 'Playing audio')",
    )
    parser.add_argument(
        "--count-all", action="store_true",
        help="let every inhibit block locking, ignoring no reasons",
    )
    parser.add_argument(
        "-d", "--release-delay", type=int, default=5, metavar="SECONDS",
        help="wait this long after the last release before re-enabling xautolock (default: 5)",
    )
    args = parser.parse_args()

    if args.count_all:
        ignored = set()
    else:
        ignored = set(args.ignore_reason or ["Playing audio"])
    inhibitor = Inhibitor(ignored, max(args.release_delay, 0), XBlanking())
    loop = GLib.MainLoop()
    exit_code = 0
    acquired = False

    def on_name_acquired(conn, name):
        nonlocal acquired
        acquired = True
        log(f"Owning {name}; ignored reasons: {sorted(ignored) or 'none'}")
        # Start from a known state, in case a previous instance died while disabled.
        # Only after owning the name, so a second instance can't undo an active inhibit.
        xautolock("-enable")

    def on_name_lost(conn, name):
        nonlocal exit_code
        if conn is None:
            log("Could not connect to the session bus")
            exit_code = 1
        elif acquired:
            log(f"Lost {name} (bus connection closed?)")
            exit_code = 1
        else:
            log(f"{name} is owned by another process (a desktop environment?); exiting")
            exit_code = EXIT_NAME_TAKEN
        loop.quit()

    def on_signal():
        log("Signal received, exiting")
        loop.quit()
        return GLib.SOURCE_REMOVE

    for sig in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
        signal_add(GLib.PRIORITY_HIGH, sig, on_signal)

    owner_id = Gio.bus_own_name(
        Gio.BusType.SESSION, BUS_NAME, Gio.BusNameOwnerFlags.DO_NOT_QUEUE,
        inhibitor.on_bus_acquired, on_name_acquired, on_name_lost,
    )
    try:
        loop.run()
    finally:
        inhibitor.shutdown()
        Gio.bus_unown_name(owner_id)
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
