#!/usr/bin/env python3
"""Default-deny gate for desktop notifications.

Installed as /usr/bin/notify-send (dpkg-divert keeps the real binary at
notify-send.real). A notification is forwarded only when an entry in the
approved registry matches BOTH the emitting program and the normalized
summary; anything else is swallowed and recorded as pending.

Invoked as `notify-send`  -> emit gate.
Invoked with a subcommand -> registry CLI (approve/revoke are TTY-gated by
the notif-approve wrapper, never by this file alone).
"""
import fcntl
import json
import os
import re
import sys
import time

REAL = "/usr/bin/notify-send.real"
STATE = os.path.join(
    os.environ.get("XDG_STATE_HOME") or os.path.expanduser("~/.local/state"),
    "notif-gate",
)
APPROVED = os.path.join(STATE, "approved.json")
PENDING = os.path.join(STATE, "pending.json")
LOCK = os.path.join(STATE, ".lock")
LOG = os.path.join(STATE, "notif-log.jsonl")
LOG_ROTATED = LOG + ".1"
# Retention: at most two generations of this size, so the log is bounded at 8 MiB.
LOG_MAX_BYTES = 4 * 1024 * 1024

NO_VALUE_OPTS = {
    "-p", "--print-id", "-e", "--transient", "-w", "--wait",
    "-v", "--version", "-?", "--help",
}
VALUE_OPTS = {
    "-u", "--urgency", "-t", "--expire-time", "-i", "--icon", "-c", "--category",
    "-h", "--hint", "-a", "--app-name", "-r", "--replace-id", "-A", "--action",
}
PASSTHROUGH_OPTS = {"-v", "--version", "-?", "--help"}

INTERPRETERS = {
    "python", "python2", "python3", "bash", "sh", "dash", "zsh", "perl", "ruby",
    "node", "bun", "env", "systemd-run", "nice", "ionice", "timeout", "stdbuf",
}


def normalize(summary):
    """Digits are the only free variable in an approved summary."""
    return re.sub(r"\d+", "#", summary or "").strip()


def slug(text):
    s = re.sub(r"[^a-z0-9#]+", "-", (text or "").lower()).strip("-")
    return s[:60] or "empty"


def parse_argv(argv):
    """Return (summary, body, passthrough) mirroring notify-send's own parsing."""
    positional = []
    passthrough = False
    i = 0
    while i < len(argv):
        a = argv[i]
        if a == "--":
            positional.extend(argv[i + 1:])
            break
        if a.startswith("-") and a != "-":
            head = a.split("=", 1)[0]
            if head in PASSTHROUGH_OPTS:
                passthrough = True
            if "=" in a and head.startswith("--"):
                i += 1
                continue
            if head in VALUE_OPTS:
                i += 2
                continue
            if head in NO_VALUE_OPTS:
                i += 1
                continue
            i += 1
            continue
        positional.append(a)
        i += 1
    summary = positional[0] if positional else ""
    body = positional[1] if len(positional) > 1 else ""
    return summary, body, passthrough


def systemd_unit():
    """Leaf unit of this process's cgroup — the outer ones are user@N.service and slices."""
    try:
        with open("/proc/self/cgroup") as f:
            found = re.findall(r"/([^/\s]+\.(?:service|scope|timer))", f.read())
    except OSError:
        return ""
    return found[-1] if found else ""


def emitter_program():
    """Absolute path of the program that called notify-send.

    A unit whose ExecStart is notify-send itself has systemd as its parent; the
    unit name is then the only identity that distinguishes one such unit from
    another.
    """
    ppid = os.getppid()
    # os.path.realpath returns the input unchanged when /proc/N/exe is unreadable,
    # so the link is resolved with readlink, which fails loudly instead.
    try:
        exe = os.readlink(f"/proc/{ppid}/exe")
    except OSError:
        exe = ""
    try:
        with open(f"/proc/{ppid}/cmdline", "rb") as f:
            parts = [p.decode("utf-8", "replace") for p in f.read().split(b"\0") if p]
    except OSError:
        parts = []
    if "systemd" in (os.path.basename(exe), os.path.basename(parts[0] if parts else "")):
        unit = systemd_unit()
        if unit:
            return f"unit:{unit}"
    for tok in parts:
        base = os.path.basename(tok)
        if base in INTERPRETERS or tok.startswith("-"):
            continue
        cand = os.path.realpath(tok)
        if os.path.isfile(cand):
            return cand
    return exe or "unknown"


class Registry:
    def __init__(self):
        os.makedirs(STATE, mode=0o700, exist_ok=True)
        self._fh = None

    def __enter__(self):
        self._fh = open(LOCK, "a+")
        fcntl.flock(self._fh, fcntl.LOCK_EX)
        return self

    def __exit__(self, *exc):
        fcntl.flock(self._fh, fcntl.LOCK_UN)
        self._fh.close()
        self._fh = None

    @staticmethod
    def _load(path, default):
        try:
            with open(path) as f:
                return json.load(f)
        except (OSError, ValueError):
            return default

    @staticmethod
    def _store(path, data):
        tmp = f"{path}.tmp.{os.getpid()}"
        with open(tmp, "w") as f:
            json.dump(data, f, indent=1, ensure_ascii=False, sort_keys=True)
            f.write("\n")
        os.replace(tmp, path)

    def approved(self):
        return self._load(APPROVED, {"version": 1, "entries": {}})

    def pending(self):
        return self._load(PENDING, {"version": 1, "entries": {}})

    def store_approved(self, data):
        self._store(APPROVED, data)

    def store_pending(self, data):
        self._store(PENDING, data)


def make_id(program, summary):
    return f"{os.path.basename(program)}:{slug(normalize(summary))}"


def record_pending(nid, program, summary, body, channel="notify-send"):
    now = time.time()
    with Registry() as reg:
        data = reg.pending()
        e = data["entries"].get(nid)
        if e is None:
            e = {
                "program": program,
                "summary_template": normalize(summary),
                "summary_sample": summary,
                "body_sample": (body or "")[:200],
                "channel": channel,
                "count": 0,
                "first": now,
                "last": now,
            }
            data["entries"][nid] = e
        else:
            e.pop("seeded", None)
        # An id is basename-derived, so a program path that moves is the same source
        # relocating; a stale path would make is_approved() reject it forever.
        e["program"] = program
        e["channel"] = channel
        e["count"] += 1
        e["last"] = now
        e["summary_sample"] = summary
        if body:
            e["body_sample"] = body[:200]
        reg.store_pending(data)


def is_approved(nid, program, summary):
    with Registry() as reg:
        e = reg.approved()["entries"].get(nid)
    if not e:
        return False
    return e.get("program") == program and e.get("summary_template") == normalize(summary)


def log_attempt(verdict, nid, program, summary, body, channel):
    """Append one record per attempt. Never raises, never blocks a notification.

    Written with a single O_APPEND write so concurrent emitters cannot interleave
    a partial line, and without the registry lock so a slow reader cannot delay
    delivery.
    """
    try:
        timestamp = time.time()
        record = {
            # These are the durable event-store fields. Keep the older aliases
            # below while existing collector deployments drain their old logs.
            "timestamp": timestamp,
            "source_app": program,
            "verdict": verdict,
            "title": summary or "",
            "body": body or "",
            "ts": timestamp,
            "decision": "delivered" if verdict == "allowed" else "suppressed",
            "id": nid,
            "program": program,
            "unit": systemd_unit(),
            "channel": channel,
            "summary": summary or "",
            "pid": os.getpid(),
        }
        line = json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n"
        os.makedirs(STATE, mode=0o700, exist_ok=True)
        fd = os.open(LOG, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
        try:
            if os.fstat(fd).st_size >= LOG_MAX_BYTES:
                os.close(fd)
                os.replace(LOG, LOG_ROTATED)
                fd = os.open(LOG, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
            os.write(fd, line.encode("utf-8"))
        finally:
            os.close(fd)
    except Exception:
        pass


def forward(real, argv):
    if not os.path.isfile(real):
        sys.stderr.write(f"notif-gate: real emitter missing at {real}\n")
        return 1
    os.execv(real, [real] + argv)


def emit(argv):
    real = os.environ.get("NOTIF_GATE_REAL", REAL)
    summary, body, passthrough = parse_argv(argv)
    if passthrough or not summary:
        return forward(real, argv)
    program = emitter_program()
    nid = make_id(program, summary)
    if is_approved(nid, program, summary):
        # forward() is execv: this process is replaced, so the record is written first.
        log_attempt("allowed", nid, program, summary, body, "notify-send")
        return forward(real, argv)
    log_attempt("blocked", nid, program, summary, body, "notify-send")
    try:
        record_pending(nid, program, summary, body)
    except OSError:
        pass
    # notify-send's caller contract: success, and an id on stdout under -p.
    if "-p" in argv or "--print-id" in argv:
        sys.stdout.write("0\n")
    return 0


GATED_CHANNELS = {"notify-send", "dbus-gated"}


def _fmt_age(ts):
    return time.strftime("%Y-%m-%d %H:%M", time.localtime(ts)) if ts else "-"


SAFETY_MARKERS = ("mem-pressure", "disk-fill", "disk-check", "pids-guard",
                  "stall-guard", "agent-guard", "agent_guard", "system-monitor",
                  "reaper")


def _is_safety(entry):
    return any(m in (entry.get("program") or "") for m in SAFETY_MARKERS)


def cli_pending(_args):
    with Registry() as reg:
        entries = reg.pending()["entries"]
    if not entries:
        print("no pending notification sources")
        return 0
    order = sorted(entries.items(),
                   key=lambda kv: (not _is_safety(kv[1]), -kv[1].get("count", 0), kv[0]))
    print(f"{len(order)} pending notification source(s) — none of these can reach your screen.")
    print("SAFETY-RELEVANT ones are listed first.\n")
    for nid, e in order:
        print(f"  id     {nid}{'   [SAFETY-RELEVANT]' if _is_safety(e) else ''}")
        print(f"  from   {e.get('program')}")
        print(f"  says   {e.get('summary_sample')!r}")
        if e.get("body_sample"):
            print(f"  body   {e['body_sample'][:120]!r}")
        print(f"  tried  {e.get('count')}x   first {_fmt_age(e.get('first'))}   last {_fmt_age(e.get('last'))}")
        if e.get("channel", "notify-send") not in GATED_CHANNELS:
            print(f"  NOTE   emits over {e['channel']} (direct D-Bus) — the gate CANNOT swallow this one")
        print()
    print("approve:  notif-approve approve <id> [<id> ...]")
    return 0


def cli_approved(_args):
    with Registry() as reg:
        entries = reg.approved()["entries"]
    if not entries:
        print("no approved notification sources (default-deny: nothing can notify you)")
        return 0
    for nid, e in sorted(entries.items()):
        print(f"{nid}\t{e.get('program')}\t{e.get('summary_template')!r}\tapproved {_fmt_age(e.get('approved_at'))}")
    return 0


def cli_approve(args):
    if not args:
        print("usage: notif-approve approve <id> [<id> ...]", file=sys.stderr)
        return 2
    rc = 0
    with Registry() as reg:
        pend = reg.pending()
        appr = reg.approved()
        for nid in args:
            e = pend["entries"].get(nid)
            if e is None:
                print(f"unknown pending id: {nid}", file=sys.stderr)
                rc = 1
                continue
            appr["entries"][nid] = {
                "program": e["program"],
                "summary_template": e["summary_template"],
                "approved_at": time.time(),
            }
            del pend["entries"][nid]
            print(f"approved {nid}  ({e['program']})")
        reg.store_approved(appr)
        reg.store_pending(pend)
    return rc


def cli_revoke(args):
    if not args:
        print("usage: notif-approve revoke <id> [<id> ...]", file=sys.stderr)
        return 2
    rc = 0
    with Registry() as reg:
        appr = reg.approved()
        for nid in args:
            if nid not in appr["entries"]:
                print(f"not approved: {nid}", file=sys.stderr)
                rc = 1
                continue
            del appr["entries"][nid]
            print(f"revoked {nid}")
        reg.store_approved(appr)
    return rc


def cli_forget(args):
    """Drop a pending record. Grants nothing — the source reappears if it tries again."""
    if not args:
        print("usage: notif-approve forget <id> [<id> ...]", file=sys.stderr)
        return 2
    rc = 0
    with Registry() as reg:
        pend = reg.pending()
        for nid in args:
            if nid not in pend["entries"]:
                print(f"not pending: {nid}", file=sys.stderr)
                rc = 1
                continue
            del pend["entries"][nid]
            print(f"forgot {nid}")
        reg.store_pending(pend)
    return rc


def cli_register(args):
    """Register a source as PENDING without it having fired. Seeding only."""
    if len(args) < 2:
        print("usage: notif-approve register <program> <summary> [body] [channel]", file=sys.stderr)
        return 2
    program, summary = args[0], args[1]
    body = args[2] if len(args) > 2 else ""
    channel = args[3] if len(args) > 3 else "notify-send"
    nid = make_id(program, summary)
    with Registry() as reg:
        if nid in reg.approved()["entries"]:
            print(f"already approved: {nid}")
            return 0
    record_pending(nid, program, summary, body, channel)
    print(nid)
    return 0


def cli_check(args):
    """Emit decision for a caller that cannot route through the notify-send binary.

    exit 0 = approved, the caller may emit.
    exit 1 = not approved; recorded as pending, the caller must emit nothing.
    """
    if len(args) < 2:
        print("usage: notif-approve check <program> <summary> [body] [channel]", file=sys.stderr)
        return 2
    program, summary = args[0], args[1]
    body = args[2] if len(args) > 2 else ""
    channel = args[3] if len(args) > 3 else "dbus-gated"
    nid = make_id(program, summary)
    if is_approved(nid, program, summary):
        log_attempt("allowed", nid, program, summary, body, channel)
        return 0
    log_attempt("blocked", nid, program, summary, body, channel)
    try:
        record_pending(nid, program, summary, body, channel)
    except OSError:
        pass
    return 1


def cli_id(args):
    if len(args) < 2:
        print("usage: notif-approve id <program> <summary>", file=sys.stderr)
        return 2
    print(make_id(args[0], args[1]))
    return 0


COMMANDS = {
    "pending": cli_pending,
    "approved": cli_approved,
    "approve": cli_approve,
    "revoke": cli_revoke,
    "forget": cli_forget,
    "register": cli_register,
    "check": cli_check,
    "id": cli_id,
}


def main(argv):
    if os.path.basename(argv[0] if argv else "") == "notify-send":
        return emit(argv[1:])
    args = argv[1:]
    if not args or args[0] not in COMMANDS:
        print("usage: notif-approve {pending|approved|approve|revoke|forget|register|check|id} [args]",
          file=sys.stderr)
        return 2
    return COMMANDS[args[0]](args[1:])


if __name__ == "__main__":
    sys.exit(main(sys.argv))