#!/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
import time
from dataclasses import dataclass
from datetime import timedelta

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
ACTIVITY_QUIET = 300  # seconds without SimulateUserActivity that end a logged burst

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'}")


@dataclass
class ActivityBurst:
    first: float
    last: float
    calls: int = 0
    timer: int = 0


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.activity: dict[str, ActivityBurst] = {}  # sender -> SimulateUserActivity burst
        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) -> None:
        for sender in list(self.activity):
            self.end_activity(sender)
        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 simulate_activity(self, sender: str) -> None:
        """Log a sender's calls once per burst: Steam sends thousands per hour."""
        now = time.monotonic()
        burst = self.activity.get(sender)
        if burst is None:
            log(f"SimulateUserActivity from {self.proc_name(sender)} (not implemented)")
            burst = self.activity[sender] = ActivityBurst(first=now, last=now)
        else:
            GLib.source_remove(burst.timer)
        burst.calls += 1
        burst.last = now
        burst.timer = GLib.timeout_add_seconds(ACTIVITY_QUIET, self.on_activity_quiet, sender)

    def on_activity_quiet(self, sender: str) -> bool:
        self.activity[sender].timer = 0
        self.end_activity(sender)
        return GLib.SOURCE_REMOVE

    def end_activity(self, sender: str) -> None:
        burst = self.activity.pop(sender, None)
        if burst is None:
            return
        if burst.timer:
            GLib.source_remove(burst.timer)
        if burst.calls == 1:
            return
        span = timedelta(seconds=round(burst.last - burst.first))
        log(f"SimulateUserActivity from {self.proc_name(sender)}: {burst.calls} calls over {span}")

    def release_sender(self, sender: str) -> None:
        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.end_activity(sender)
        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":
            self.simulate_activity(sender)
            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()
