#!/usr/bin/env python3
"""wallet-launcher: make apps use the ONE KDE wallet shared by both Frame sessions.

The Frame runs two sessions: Steam/VR on the outer D-Bus (/run/user/$UID/bus)
and the nested Plasma desktop on a private bus. An app started from the desktop
talks to a second kwalletd6 on the private bus; secrets it stores there are
invisible to the same app in VR, and two daemons writing one wallet file can
overwrite each other's changes.

This writes a launcher override into ~/.local/share/applications that starts
the app on the outer bus (and, for Flatpaks, lets it talk to the wallet), so
every launch, from the desktop menu or the VR "+" menu, uses the same wallet.
Approach from lhns/steam-frame-nix (docs/launchers.md, `keyring.enable`).
What you change in such an entry (its name or icon in KDE's menu editor, an
option added to its command) stays when it's rebuilt or removed (lib/menuedits.py).

  wallet-launcher org.signal.Signal com.anthropic.Claude   # add
  wallet-launcher --electron some.electron.App             # also --password-store=kwallet6
  wallet-launcher --gate org.mozilla.firefox               # also require the wallet password
  wallet-launcher --gate --no-shared android-mirror         # password only, keep the app's own bus
  wallet-launcher --no-gate org.mozilla.firefox            # stop requiring it (re-running without
                                                           # either keeps what the entry has)
  wallet-launcher --remove org.signal.Signal               # restore the original entry
  wallet-launcher --reapply android-mirror                  # rebuild from the original, same flags
  wallet-launcher --list
"""
import argparse
import os
import re
import shlex
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_ = _
import menuedits  # noqa: E402  (your edits to an entry, kept: lib/menuedits.py, installed beside i18n.py)

HOME = Path.home()
APPS = HOME / ".local/share/applications"
SOURCES = [HOME / ".local/share/flatpak/exports/share/applications",
           Path("/var/lib/flatpak/exports/share/applications")]
# System apps (read-only partition): wrapped by an override of the same name in APPS, which
# --remove deletes again. Checked last.
SYSTEM_SOURCES = [Path("/usr/local/share/applications"), Path("/usr/share/applications")]
MARK = "X-FrameTweaks-Wallet"
# Outside APPS: Steam's "+" menu also lists entries in subfolders of APPS (duplicates).
BACKUP_DIR = HOME / ".local/share/frame-tweaks/launcher-originals"
OLD_BACKUP_DIR = APPS / ".wallet-launcher-originals"
# Each wrapper as add() last built it: what the entry in APPS says differently is yours (a new
# name or icon from KDE's menu editor, an option added to the command), kept when it's rebuilt.
BUILT_DIR = HOME / ".local/share/frame-tweaks/launcher-built"
# Frame Tweaks' own entries as it last wrote them (lib/common.sh install_desktop), wrapped or not.
MENU_BUILT = HOME / ".local/share/frame-tweaks/menu-built"
WRAPPER = str(HOME / ".local/bin/frame-outer-bus")
GATE = str(HOME / ".local/bin/wallet-gate")
GUARD = str(HOME / ".local/bin/signal-guard")   # retired (app-timers + --gate); dropped on rebuild
TALK = ["--talk-name=org.kde.kwalletd6", "--talk-name=org.freedesktop.secrets"]
STORE = "--password-store=kwallet6"
ADDED = TALK + [STORE]                          # what rewrite_exec adds to a command's arguments
OWNED = (MARK, "DBusActivatable")               # keys of a wrapper that aren't yours to change


def source_for(app_id):
    own = APPS / f"{app_id}.desktop"
    backup = BACKUP_DIR / f"{app_id}.desktop"
    # An entry there with nothing of ours is the original (add() backs it up): a host app's own
    # (e.g. Claude Desktop), your edit of a Flatpak's or system app's, or a newer one an app's
    # installer wrote over our wrapper (then the backup is out of date).
    text = own.read_text() if own.exists() else ""
    if text and not any(w in text for w in (f"{MARK}=true", WRAPPER, GATE, GUARD)):
        return own
    if backup.exists():                       # re-run: start from the untouched original
        return backup
    for d in SOURCES:
        if (d / f"{app_id}.desktop").exists():
            return d / f"{app_id}.desktop"
    system = next((d / f"{app_id}.desktop" for d in SYSTEM_SOURCES if (d / f"{app_id}.desktop").exists()), None)
    # Ours without a backup: rebuilt from itself, but an override of a system app from the system entry.
    if own.exists() and not (system and f"{MARK}=true" in text):
        return own
    if system:
        return system
    sys.exit(_("wallet-launcher: no desktop entry found for {app_id}").format(app_id=app_id))


def flatpak_run(args):
    return len(args) > 1 and os.path.basename(args[0]) == "flatpak" and args[1] == "run"


def join_exec(args):
    # Field codes (%U, @@u) must stay unquoted.
    return " ".join(a if re.fullmatch(r"%[a-zA-Z]|@@[a-z]?", a) else shlex.quote(a) for a in args)


class NotWrappable(ValueError):
    pass


def command_parts(cmd):
    """(env, args): cmd's `env VAR=…` in front, if any (KDE's menu editor writes the environment
    variables you give an entry there), and its program and arguments, without the programs
    rewrite_exec puts in front (before or after that env)."""
    args, env = shlex.split(cmd), []
    while args and args[0] in (GATE, WRAPPER, GUARD, "--"):   # (an already wrapped entry)
        args = args[1:]
    if args and os.path.basename(args[0]) == "env":
        n = next((i for i, a in enumerate(args[1:], 1) if "=" not in a), len(args))
        env, args = args[:n], args[n:]
        while args and args[0] in (GATE, WRAPPER, GUARD, "--"):
            args = args[1:]
    return env, args


def rewrite_exec(cmd, app_id, electron, shared=True, gate=False):
    env, args = command_parts(cmd)
    store = [STORE] if electron and STORE not in args else []
    if (shared or electron) and flatpak_run(args):
        idx = next((i for i, a in enumerate(args) if a == app_id), None)
        if idx is None:
            raise NotWrappable(_("wallet-launcher: can't find {app_id} in: {command}").format(app_id=app_id, command=cmd))
        extra = [t for t in TALK if t not in args] if shared else []
        args = args[:idx] + extra + [app_id] + store + args[idx + 1:]
    elif store:                               # (after env VAR=…: the program's own option)
        args = args[:1] + store + args[1:]
    return join_exec(([GATE] if gate else []) + ([WRAPPER] if shared else []) + env + args)


def unwrapped(cmd, original):
    """cmd without what rewrite_exec put in: the programs in front, and the wallet options the
    original command (if any) didn't have itself. None if cmd can't be read."""
    try:
        (env, args), keep = command_parts(cmd), shlex.split(original or "")
    except ValueError:
        return None
    return env + [a for a in args if a not in ADDED or a in keep]


def plain_words(cmd):
    """A command as the app has it, whatever it's wrapped with: to tell whether the app's own
    command changed (then a command you changed gives way to it: menuedits.merge)."""
    env, args = command_parts(cmd)
    return env + [a for a in args if a not in ADDED]


def keep_yours(app_id, built, now, base, original, electron, shared, gate):
    """built (the new wrapper), with what you've changed since base was written (now: the entry
    in the menu). A command you changed is wrapped again as you left it."""
    if base is None:
        return built
    commands = menuedits.keys(original)

    def command(group, yours):
        args = unwrapped(yours, commands.get((group, "Exec")))
        try:
            return None if args is None else rewrite_exec(join_exec(args), app_id, electron, shared, gate)
        except NotWrappable:                  # (not one it can wrap: the new wrapper's command stands)
            return None
    return menuedits.merge(base, built, now, OWNED, command, words=plain_words)


# Electron apps that read the wallet only when told the password store explicitly. Without it,
# Claude Desktop ran with "safeStorage not available": no saved login, secrets in plain text.
KNOWN_ELECTRON = {"com.anthropic.Claude"}


def build(original, app_id, electron, shared, gate):
    """The wrapper for app_id: its original entry (text), with these flags."""
    out = []
    for line in original.splitlines():
        if line.startswith("Exec="):
            line = "Exec=" + rewrite_exec(line[5:], app_id, electron, shared, gate)
        elif line.startswith("DBusActivatable=") and (shared or gate):
            line = "DBusActivatable=false"   # otherwise the desktop starts it over D-Bus, not Exec=
        elif line.startswith(MARK + "="):
            continue
        out.append(line)
        if line == "[Desktop Entry]":
            out.append(f"{MARK}=true")
    return "\n".join(out) + "\n"


def add(app_id, electron, shared=True, gate=None):
    """gate: True/False to set it, None to keep what the current entry has (re-wrapping after an
    app update, or another tweak adding its app, must not silently drop the password gate)."""
    src = source_for(app_id)
    original = src.read_text()
    current = APPS / f"{app_id}.desktop"
    now = current.read_text() if current.exists() else ""
    # What the entry says differently from what Frame Tweaks wrote there is yours: from the wrapper
    # as it was built, or, for one of Frame Tweaks' own entries wrapped now, from that entry as it
    # was written (then it's the original: your edits go on top, not into it).
    base, written = BUILT_DIR / f"{app_id}.desktop", MENU_BUILT / f"{app_id}.desktop"
    if f"{MARK}=true" in now:
        base = base.read_text() if base.exists() else None
    elif src == current and written.exists():
        base = original = written.read_text()
    else:
        base = None
    lines = original.splitlines()
    if gate is None:
        gate = f"{MARK}=true" in now and GATE in now
    # Re-applying (e.g. after an app update) keeps --password-store, even if not asked again.
    electron = electron or app_id in KNOWN_ELECTRON or STORE in now
    if src.parent == APPS and src.name == f"{app_id}.desktop" and f"{MARK}=true" not in original:
        # (never the wrapper itself as the original: Remove then put a wrapper back, running
        # wallet-gate after the wallet's uninstall had deleted it)
        BACKUP_DIR.mkdir(parents=True, exist_ok=True)
        (BACKUP_DIR / src.name).write_text("\n".join(lines) + "\n")
    try:
        built = build(original, app_id, electron, shared, gate)
    except NotWrappable as e:
        sys.exit(str(e))
    APPS.mkdir(parents=True, exist_ok=True)
    current.write_text(keep_yours(app_id, built, now, base, original, electron, shared, gate))
    BUILT_DIR.mkdir(parents=True, exist_ok=True)
    (BUILT_DIR / f"{app_id}.desktop").write_text(built)
    what = [w for w, on in ((_("shared wallet"), shared), (_("password required"), gate)) if on]
    print(f"{app_id}: {', '.join(what)}")


def flags_of(app_id):
    """(shared, gate) as currently applied to app_id's entry."""
    text = (APPS / f"{app_id}.desktop").read_text()
    return WRAPPER in text, GATE in text


def migrate_backups():
    """Move backups from the old in-APPS folder (which Steam listed as extra entries)."""
    if OLD_BACKUP_DIR.is_dir():
        BACKUP_DIR.mkdir(parents=True, exist_ok=True)
        for f in OLD_BACKUP_DIR.glob("*.desktop"):
            if not (BACKUP_DIR / f.name).exists():
                f.replace(BACKUP_DIR / f.name)
            else:
                f.unlink()
        OLD_BACKUP_DIR.rmdir()


def record_built():
    """The copy of a wrapper made before launcher-built/ was kept, if it's still as add() builds
    it: what changes in it from then on is yours. (One that isn't can't be told from an app's
    update: it goes on as before, until it's built again.)"""
    for f in APPS.glob("*.desktop"):
        base = BUILT_DIR / f.name
        try:
            text = f.read_text()
            if base.exists() or f"{MARK}=true" not in text:
                continue
            built = build(source_for(f.stem).read_text(), f.stem, f.stem in KNOWN_ELECTRON or STORE in text,
                          WRAPPER in text, GATE in text)
        except (OSError, ValueError, SystemExit):    # (an app gone, a command it can't read)
            continue
        if menuedits.keys(built) == menuedits.keys(text):
            BUILT_DIR.mkdir(parents=True, exist_ok=True)
            base.write_text(built)


def remove(app_id):
    target, backup, base = (d / f"{app_id}.desktop" for d in (APPS, BACKUP_DIR, BUILT_DIR))
    now = target.read_text() if target.exists() else ""
    if not backup.exists() and f"{MARK}=true" not in now:
        print(_("{app_id}: no wallet launcher").format(app_id=app_id)); return
    # The original goes back with what you'd changed in the wrapper, as KDE's menu editor saves an
    # edit: for a Flatpak's or system app's entry, an override of the same name.
    original = backup if backup.exists() else next(
        (d / f"{app_id}.desktop" for d in SOURCES + SYSTEM_SOURCES if (d / f"{app_id}.desktop").exists()), None)
    text = mine = original.read_text() if original else ""
    if f"{MARK}=true" in now and base.exists() and original:
        commands = menuedits.keys(text)

        def command(group, yours):
            args, cmd = unwrapped(yours, commands.get((group, "Exec"))), commands.get((group, "Exec"), "")
            return None if args is None or args == shlex.split(cmd) else join_exec(args)
        mine = menuedits.merge(base.read_text(), text, now, OWNED, command, words=plain_words)
    if menuedits.keys(mine) != menuedits.keys(text):
        target.write_text(mine)
        backup.unlink(missing_ok=True)
    elif backup.exists():
        backup.replace(target)
    else:
        target.unlink()
    base.unlink(missing_ok=True)
    print(_("{app_id}: original launcher restored").format(app_id=app_id))


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("apps", nargs="*")
    ap.add_argument("--electron", action="store_true", help=_("also pass --password-store=kwallet6"))
    ap.add_argument("--no-gate", action="store_true", help=_("stop requiring the wallet password"))
    ap.add_argument("--gate", action="store_true",
                    help=_("require the wallet password (wallet-unlock) before starting the app"))
    ap.add_argument("--gate-if-new", action="store_true",
                    help=_("require it for apps not protected yet; the others keep their choice (installers)"))
    ap.add_argument("--no-shared", action="store_true",
                    help=_("don't move the app to the outer bus (for apps that don't use the wallet)"))
    ap.add_argument("--remove", action="store_true")
    ap.add_argument("--list", action="store_true")
    ap.add_argument("--reapply", action="store_true",
                    help=_("rebuild the wrapper from the (updated) original, keeping its flags"))
    a = ap.parse_args()
    migrate_backups()
    record_built()
    if a.list:
        for f in sorted(APPS.glob("*.desktop")):
            try:
                text = f.read_text(errors="replace")
            except OSError:                   # (a link to an entry that's gone, say)
                continue
            if f"{MARK}=true" in text:
                flags = [w for w, t in (("shared", WRAPPER), ("gate", GATE)) if t in text]
                print(f"{f.stem:28} {' '.join(flags)}")
        return
    if not a.apps:
        ap.error(_("give at least one app ID"))
    for app in a.apps:
        if a.remove:
            remove(app)
        elif a.reapply:
            current = APPS / f"{app}.desktop"
            if not (current.exists() and f"{MARK}=true" in current.read_text()):
                remove(app)          # no wrapper left to rebuild (its app was removed): the original goes back
                continue
            shared, gate = flags_of(app)
            add(app, STORE in current.read_text(), shared, gate)
        else:
            gate = True if a.gate else False if a.no_gate else None
            if a.gate_if_new:          # (an update runs the installers again: Password first unticked stays so)
                current = APPS / f"{app}.desktop"
                gate = None if current.exists() and f"{MARK}=true" in current.read_text() else True
            add(app, a.electron, not a.no_shared, gate)
    os.system("update-desktop-database -q ~/.local/share/applications 2>/dev/null")
    os.utime(APPS)   # Steam's "+" menu rescans when this directory's mtime changes


if __name__ == "__main__":
    main()
