#!/usr/bin/env python3
"""Small Rofi session menu; native services retain locking and power policy."""
import argparse
from pathlib import Path
import os
import socket
import subprocess
import tomllib

ACTIONS = ("Lock", "Log out", "Suspend", "Restart", "Shut down")
SCREEN_OPERATIONS = ("lock", "away", "suspend", "logout", "restart", "shutdown")


def screen_command(seconds=5, label=None, select=None):
    if type(seconds) is not int or not 1 <= seconds <= 60:
        raise ValueError("session screen seconds must be between 1 and 60")
    label = socket.gethostname() if label is None else label
    if (not isinstance(label, str) or not label.strip() or len(label) > 64
            or not all(character.isprintable() for character in label)):
        raise ValueError("session screen label must be 1 to 64 printable characters")
    if select is not None and select not in SCREEN_OPERATIONS:
        raise ValueError("invalid session screen action")
    command = [str(Path(__file__).resolve().with_name("desktop-suspend")), "menu",
               "--seconds", str(seconds), "--label", label]
    if select is not None:
        command.extend(("--select", select))
    return command


def choose(options, prompt, *, run=subprocess.run):
    result = run(["rofi", "-dmenu", "-i", "-p", prompt, "-selected-row", "0"],
                 input="\n".join(options) + "\n", text=True, capture_output=True)
    if result.returncode != 0:
        return None
    value = result.stdout.rstrip("\n")
    return value if value in options else None


def dispatch(action, *, run=subprocess.run, picker=choose, lock_before_suspend=False, confirmation=None):
    if confirmation is not None and action in {"Lock", "Suspend", "Restart", "Shut down",
                                               "Turn off displays", "Lock and turn off displays"}:
        operation = {"Lock": "lock", "Suspend": "suspend", "Restart": "restart", "Shut down": "shutdown"}.get(action, "away")
        command = [str(Path(__file__).resolve().with_name("desktop-suspend")), "countdown", "--operation", operation,
                   "--seconds", str(confirmation["seconds"]), "--label", confirmation["label"]]
        if operation == "away" and confirmation["lock_on_away"]:
            command.append("--lock")
        run(command, check=True)
        return True
    if action == "Lock":
        command = ["systemctl", "--user", "start", "desktop-lock.service"]
    elif action == "Log out":
        # Niri owns its normal quit confirmation. Never use --skip-confirmation.
        command = ["niri", "msg", "action", "quit"]
    elif action == "Suspend":
        # The same independent unit also makes this idempotent with before-sleep.
        if lock_before_suspend:
            run(["systemctl", "--user", "start", "desktop-lock.service"], check=True)
        command = ["systemctl", "suspend"]
    elif action in {"Restart", "Shut down"}:
        if picker(("Cancel", action), "Confirm " + action.lower()) != action:
            return False
        command = ["systemctl", "reboot" if action == "Restart" else "poweroff"]
    else:
        return False
    run(command, check=True)
    return True


def read_lock_before_suspend(config_home=None):
    config_home = Path(config_home or os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config"))
    policy = config_home / "desktop-session/power.toml"
    try:
        value = tomllib.loads(policy.read_text())["lock_before_suspend"]
    except (OSError, KeyError, tomllib.TOMLDecodeError) as error:
        raise SystemExit(
            "Suspend canceled: set lock_before_suspend as a TOML boolean in "
            f"{policy}"
        ) from error
    if not isinstance(value, bool):
        raise SystemExit(f"Suspend canceled: invalid lock_before_suspend value in {policy}")
    return value


def read_confirmation_policy(config_home=None):
    config_home = Path(config_home or os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config"))
    policy = config_home / "desktop-session/confirmation.toml"
    try:
        data = tomllib.loads(policy.read_text())
    except FileNotFoundError:
        return None  # Optional adoption; retained native confirmation stays in place.
    except (OSError, tomllib.TOMLDecodeError) as error:
        raise SystemExit(f"Session action canceled: cannot read {policy}") from error
    if (type(data.get("seconds")) is not int or not 1 <= data["seconds"] <= 60
            or not isinstance(data.get("label"), str) or not 1 <= len(data["label"]) <= 64
            or not all(character.isprintable() for character in data["label"])
            or not data["label"].strip() or type(data.get("lock_on_away")) is not bool):
        raise SystemExit(f"Session action canceled: invalid confirmation policy in {policy}")
    return data


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--screen", action="store_true", help="Open the native session action screen")
    parser.add_argument("--seconds", type=int, help="Session screen countdown in seconds")
    parser.add_argument("--label", help="Session screen computer label")
    parser.add_argument("--select", choices=SCREEN_OPERATIONS,
                        help="Start this action's cancellable screen countdown")
    args = parser.parse_args(argv)
    if args.screen:
        command = screen_command(args.seconds if args.seconds is not None else 5, args.label, args.select)
        subprocess.run(command, check=True)
        return 0
    if args.seconds is not None or args.label is not None or args.select is not None:
        parser.error("--seconds, --label, and --select require --screen")
    ready = Path(__file__).resolve().with_name("desktop-session-ready")
    if subprocess.run([str(ready)]).returncode != 0:
        raise SystemExit("Session actions require the captured Niri graphical session")
    confirmation = read_confirmation_policy()
    actions = ACTIONS
    if confirmation is not None:
        away = "Lock and turn off displays" if confirmation["lock_on_away"] else "Turn off displays"
        actions = (ACTIONS[0], away, *ACTIONS[1:])
    action = choose(actions, "Session")
    lock_before_suspend = False
    if action == "Suspend" and confirmation is None:
        config_home = Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config"))
        lock_before_suspend = read_lock_before_suspend(config_home)
    dispatch(action, lock_before_suspend=lock_before_suspend, confirmation=confirmation)


if __name__ == "__main__":
    main()