#!/usr/bin/env python3
"""Track session ScreenSaver inhibitors and expose combined idle blockers as TSV.

logind exposes idle inhibitors through systemd-inhibit. Browser video wake locks
use org.freedesktop.ScreenSaver instead, which has no list API, so this monitor
observes Inhibit/UnInhibit calls and keeps their active state in a small JSON
file for the Noctalia widget and panel.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import signal
import subprocess
import sys
from pathlib import Path
from tempfile import NamedTemporaryFile

STATE_FILE = Path.home() / ".local/state/noctalia/idle-inhibitors.json"
SCREEN_SAVER = "org.freedesktop.ScreenSaver"


def write_state(inhibitors: dict[tuple[str, int], dict[str, str]]) -> None:
    STATE_FILE.parent.mkdir(parents=True, exist_ok=True)
    payload = {"inhibitors": list(inhibitors.values())}
    with NamedTemporaryFile("w", dir=STATE_FILE.parent, delete=False) as file:
        json.dump(payload, file, separators=(",", ":"))
        file.write("\n")
        temporary_name = file.name
    os.replace(temporary_name, STATE_FILE)


def read_screensaver_inhibitors() -> list[dict[str, str]]:
    try:
        payload = json.loads(STATE_FILE.read_text())
    except (FileNotFoundError, json.JSONDecodeError):
        return []
    inhibitors = payload.get("inhibitors", [])
    return [item for item in inhibitors if isinstance(item, dict)]


def logind_inhibitors() -> list[dict[str, str]]:
    result = subprocess.run(
        ["systemd-inhibit", "--list", "--no-pager"],
        check=False,
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        return []

    inhibitors = []
    for line in result.stdout.splitlines()[1:]:
        fields = line.split(maxsplit=6)
        if len(fields) < 7 or "idle" not in fields[5].split(":"):
            continue
        reason_and_mode = fields[6].rsplit(maxsplit=1)
        inhibitors.append(
            {
                "source": "logind",
                "who": fields[0],
                "user": fields[2],
                "pid": fields[3],
                "comm": fields[4],
                "mode": reason_and_mode[-1],
                "why": reason_and_mode[0] if len(reason_and_mode) == 2 else "No reason supplied",
            }
        )
    return inhibitors


def combined_inhibitors() -> list[dict[str, str]]:
    return logind_inhibitors() + read_screensaver_inhibitors()


def print_inhibitors() -> None:
    for inhibitor in combined_inhibitors():
        values = [
            inhibitor.get("source", "screensaver"),
            inhibitor.get("who", ""),
            inhibitor.get("user", ""),
            inhibitor.get("pid", ""),
            inhibitor.get("comm", ""),
            inhibitor.get("mode", ""),
            inhibitor.get("why", ""),
        ]
        print("\t".join(value.replace("\t", " ").replace("\n", " ") for value in values))


def monitor() -> None:
    active: dict[tuple[str, int], dict[str, str]] = {}
    pending: dict[tuple[str, int], dict[str, str]] = {}
    write_state(active)

    process = subprocess.Popen(
        ["stdbuf", "-oL", "busctl", "--user", "--json=short", "monitor"],
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        text=True,
        bufsize=1,
    )

    def stop(_signal: int, _frame: object) -> None:
        process.terminate()
        write_state({})
        raise SystemExit(0)

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)

    assert process.stdout is not None
    for line in process.stdout:
        try:
            message = json.loads(line)
        except json.JSONDecodeError:
            continue

        message_type = message.get("type")
        sender = message.get("sender", "")
        payload = message.get("payload", {}).get("data", [])

        if (
            message_type == "method_call"
            and message.get("interface") == SCREEN_SAVER
            and message.get("member") == "Inhibit"
            and len(payload) >= 2
        ):
            pending[(sender, int(message.get("cookie", 0)))] = {
                "source": "screensaver",
                "who": str(payload[0]),
                "user": "session",
                "pid": "",
                "comm": str(payload[0]),
                "mode": "block",
                "why": str(payload[1]),
            }
        elif (
            message_type == "method_return"
            and payload
            and isinstance(payload[0], int)
        ):
            request = (message.get("destination", ""), int(message.get("reply_cookie", 0)))
            inhibitor = pending.pop(request, None)
            if inhibitor is not None:
                active[(request[0], int(payload[0]))] = inhibitor
                write_state(active)
        elif (
            message_type == "method_call"
            and message.get("interface") == SCREEN_SAVER
            and message.get("member") == "UnInhibit"
            and payload
            and isinstance(payload[0], int)
        ):
            cookie = int(payload[0])
            active.pop((sender, cookie), None)
            # Some clients reconnect before releasing a cookie; the cookie is
            # globally unique for Noctalia, so clean up that fallback as well.
            for key in [key for key in active if key[1] == cookie]:
                active.pop(key, None)
            write_state(active)
        elif (
            message_type == "signal"
            and message.get("interface") == "org.freedesktop.DBus"
            and message.get("member") == "NameOwnerChanged"
            and len(payload) == 3
            and payload[2] == ""
        ):
            vanished = str(payload[0])
            for key in [key for key in active if key[0] == vanished]:
                active.pop(key, None)
            for key in [key for key in pending if key[0] == vanished]:
                pending.pop(key, None)
            write_state(active)


parser = argparse.ArgumentParser()
parser.add_argument("--count", action="store_true", help="print combined inhibitor count")
parser.add_argument("--list", action="store_true", help="print combined inhibitors as TSV")
arguments = parser.parse_args()

if arguments.count:
    print(len(combined_inhibitors()))
elif arguments.list:
    print_inhibitors()
else:
    monitor()