"""Pause Frametop for VR games, so it leaves the headset's CPU and GPU to the game. The input
relay (input-relay.py) runs this: it's the one part of Frametop that always runs.

Paused:
  - The gaze service stops (frametop-gaze: ft-gazed, ft-gaze, our own eye tracker ft-eyes, the
    gaze panel), and with it every read of SteamVR's eye tracking. Our frame grabber, the root
    service ft-eyegrab, goes idle by itself 3 seconds after ft-eyes stops asking for frames.
  - Hand tracking stops if it runs (frametop-camd, frametop-hands).
  - The desktop, by "pause_desktop": "hide" (the default) hides every screen and floating window
    and slows the desktop down ("pause on" to ft-screens: KWin gets one frame callback a second
    instead of 90, so it and the apps on it hardly draw), and remote desktop stops if it runs
    (session/remote-ctl.sh: krdp, Xvnc, FreeRDP). Windows stay open. "close" closes the desktop
    (desktops.sh stop): its windows close, and resuming starts it again (about 12 seconds), in
    its start profile if it has one.
  - In the relay: the 3D mouse lets go and the mouse is a plain mouse for SteamVR, typing goes to
    Steam, and mapped buttons and key combinations do nothing but PAUSED_ACTIONS and commands.
Resuming starts again only what pausing stopped. The pointer helper and ft-powerd keep running:
they cost little, and the helper is what says a VR game started.

Toggled by:
  - The pause_toggle action: a mouse button, a key combination, or a controller button outside
    games (the pointer helper only reads those outside games).
  - A controller gesture, by default both thumbsticks clicked together twice ("pause_gesture").
    It's read from vrserver's web socket (vrws.py), which works in games and takes nothing from
    them, so the game sees the clicks too.
  - input/ft-pause on|off|toggle|status, which uses the relay's control socket
    ("pause on|off|toggle|?").
  - VR games, with "pause_auto" (on by default): the pointer helper says "vrgame 1|0" as a SteamVR
    scene app starts and ends (and every 5 seconds). A game starting pauses; a pause that starts
    while a game runs ends RESUME_DELAY after the game does, unless another game starts first.
    Toggled back on during a game, Frametop stays on until that game ends. A pause that starts
    outside a game lasts until it's toggled off.

Settings in ~/.config/frametop-input.json (Frametop Input Settings, Game optimization page): "pause_auto"
(bool), "pause_gesture" ({"buttons": one or two controller buttons, "presses": 1 or 2}, or null
for none; one button always takes two presses), "pause_desktop" ("hide" or "close"), and
"pause_sound" (a sound on pause and resume; bool). The state outlives a relay restart in
STATE_PATH.
"""
import json
import os
import queue
import socket
import subprocess
import threading
import time

import vrws

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DESKTOPS = os.path.join(REPO, "desktops.sh")
REMOTE_CTL = os.path.join(REPO, "session", "remote-ctl.sh")
RELAY = "\0frametop_relay"
SCREENS = "\0ft_screens"
STATE_PATH = f"/run/user/{os.getuid()}/frametop-pause.json"
# User services that pausing stops, if they run.
UNITS = ("frametop-gaze.service", "frametop-hands.service", "frametop-camd.service")
# What still works while paused, besides commands ("command:CMD"): the toggle itself, and Steam's menu.
PAUSED_ACTIONS = ("pause_toggle", "steam_menu")
DEFAULT_GESTURE = {"buttons": ["left/thumbstick", "right/thumbstick"], "presses": 2}
DESKTOP_MODES = ("hide", "close")
SOUNDS = {True: "/usr/share/sounds/ocean/stereo/device-removed.oga",
          False: "/usr/share/sounds/ocean/stereo/device-added.oga"}
RESUME_DELAY = 5.0  # after a VR game ends: loading the next one can end one scene app and start another
GAME_STALE = 12.0  # the helper repeats "vrgame" every 5 s; silence this long means no game (or no SteamVR)
CHORD_WINDOW = 0.3  # a two-button gesture's buttons go down within this of each other
DOUBLE_WINDOW = 0.7  # a double press's second press comes within this of the first


def gesture_of(rules, buttons):
    """The rules' pause gesture as (buttons, presses), or None for none. `buttons` are the
    controller buttons there are (the relay's VR_BUTTONS)."""
    spec = rules.get("pause_gesture", DEFAULT_GESTURE)
    if not isinstance(spec, dict) or not isinstance(spec.get("buttons"), list):
        return None
    chosen = tuple(dict.fromkeys(b for b in spec["buttons"] if b in buttons))[:2]
    if not chosen:
        return None
    return chosen, 1 if len(chosen) == 2 and spec.get("presses") == 1 else 2


def settings_of(rules, buttons):
    return {"auto": rules.get("pause_auto", True) is not False,
            "gesture": gesture_of(rules, buttons),
            "desktop": rules.get("pause_desktop") if rules.get("pause_desktop") in DESKTOP_MODES else "hide",
            "sound": rules.get("pause_sound", True) is not False}


class Gesture:
    """The pause gesture, from button presses and releases. A press of a two-button gesture is
    both buttons going down within CHORD_WINDOW of each other, so a button held for a while (to
    sprint, say) and the other clicked doesn't count. Two presses within DOUBLE_WINDOW make a
    double press."""

    def __init__(self, buttons, presses):
        self.buttons, self.presses = tuple(buttons), presses
        self.down = {}  # button -> when it went down
        self.first = None  # when the first press of a double press was
        self.fired = None

    def reset(self):
        self.down.clear()
        self.first = None

    def feed(self, button, pressed, now):
        """A button's state; True when it completes the gesture."""
        if button not in self.buttons:
            return False
        if not pressed:
            self.down.pop(button, None)
            return False
        if button in self.down:
            return False  # still down
        self.down[button] = now
        if len(self.down) < len(self.buttons) or now - min(self.down.values()) > CHORD_WINDOW:
            return False
        if self.presses == 2 and (self.first is None or now - self.first > DOUBLE_WINDOW):
            self.first = now
            return False
        self.first = None
        if self.fired is not None and now - self.fired < 1.0:
            return False  # a triple press is still one
        self.fired = now
        return True


class ControllerWatch(threading.Thread):
    """Follows the Frame controllers on vrserver's web socket for the pause gesture, and sends
    the relay "pause toggle gesture" when it's made. Reconnects every 5 seconds while SteamVR
    is away. With no gesture set, it doesn't connect."""

    def __init__(self, log):
        super().__init__(name="pause-gesture", daemon=True)
        self.log = log
        self.lock = threading.Lock()
        self.gesture = None
        self.changed = threading.Event()
        self.connected = False
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)

    def set_gesture(self, spec):
        """(buttons, presses), or None for no gesture."""
        with self.lock:
            g = self.gesture
            if spec is None:
                self.gesture = None
            elif g is None or (g.buttons, g.presses) != spec:
                self.gesture = Gesture(*spec)
        self.changed.set()

    def current(self):
        with self.lock:
            return self.gesture

    def run(self):
        while True:
            if self.current() is None:
                if self.connected:
                    self.log("pause gesture: none set, stopped reading the controllers")
                self.connected = False
                self.changed.wait()
                self.changed.clear()
                continue
            try:
                self.session()
            except (OSError, ValueError) as e:  # urllib's errors are OSErrors
                if self.connected:
                    self.log(f"pause gesture: SteamVR's web socket went away ({e}); retrying")
                self.connected = False
            self.changed.wait(5)
            self.changed.clear()

    def session(self):
        ws = vrws.VrSocket()
        try:
            ws.open(f"frametop_pause_{os.getpid()}")
            sides = {}  # root path -> side
            next_poll = 0.0
            while True:
                g = self.current()
                if g is None:
                    return
                now = time.monotonic()
                if now >= next_poll:
                    # A controller connects later, or changes its root path when the 3D mouse's
                    # virtual controller takes or gives back its hand role.
                    next_poll = now + 3.0
                    found = vrws.controllers()
                    if found != sides:
                        for path in sides.keys() - found.keys():
                            ws.subscribe(path, on=False)
                        for path in found.keys() - sides.keys():
                            ws.subscribe(path)
                        sides = found
                        g.reset()
                    if not self.connected:
                        self.connected = True
                        self.log(f"pause gesture: following {' and '.join(sorted(sides.values())) or 'no'} "
                                 f"controller{'s' if len(sides) != 1 else ''} on SteamVR's web socket")
                wanted = {b: (b.split("/", 1)[0], f"/input/{b.split('/', 1)[1]}/click") for b in g.buttons}
                for text in ws.messages(timeout=1.0):
                    # About 160 messages a second, nearly all capacitive sensing: parse only the
                    # few that carry one of the gesture's buttons.
                    if not any(c in text for _, c in wanted.values()):
                        continue
                    try:
                        msg = json.loads(text)
                    except ValueError:
                        continue
                    side = sides.get(msg.get("device"))
                    components = msg.get("components") or {}
                    for button, (s, component) in wanted.items():
                        if s == side and component in components:
                            if g.feed(button, bool(components[component]), time.monotonic()):
                                self.log("pause gesture")
                                try:
                                    self.sock.sendto(b"pause toggle gesture", RELAY)
                                except OSError:
                                    pass
        finally:
            ws.close()


def run(cmd, timeout=30):
    """A command's exit status (None if it couldn't run), its output going nowhere."""
    try:
        return subprocess.run(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
                              stderr=subprocess.DEVNULL, timeout=timeout).returncode
    except (OSError, subprocess.TimeoutExpired):
        return None


def output(cmd, timeout=10):
    try:
        return subprocess.run(cmd, stdin=subprocess.DEVNULL, capture_output=True, text=True,
                              timeout=timeout).stdout
    except (OSError, subprocess.TimeoutExpired):
        return ""


def active(unit):
    return run(["systemctl", "--user", "-q", "is-active", unit], timeout=10) == 0


def desktop_running():
    return run(["pgrep", "-x", "ft-screens"], timeout=5) == 0


class Worker(threading.Thread):
    """Does the slow part of pausing and resuming, in order, one job at a time: stopping and
    starting services, remote desktop, and closing or starting the desktop. `stopped` is what
    the last pause stopped, for the resume to start again."""

    def __init__(self, log, save):
        super().__init__(name="pause-worker", daemon=True)
        self.log, self.save = log, save
        self.jobs = queue.Queue()
        self.lock = threading.Lock()
        self.stopped = {"units": [], "remote": False, "desktop": None}
        self.busy = False

    def submit(self, job, *args):
        self.busy = True
        self.jobs.put((job, args))

    def run(self):
        while True:
            job, args = self.jobs.get()
            try:
                job(*args)
            except Exception as e:  # one failure mustn't stop the next pause or resume
                self.log(f"pause: {job.__name__} failed: {e!r}")
            self.busy = not self.jobs.empty()
            self.save()

    def snapshot(self):
        with self.lock:
            return {"units": list(self.stopped["units"]), "remote": self.stopped["remote"],
                    "desktop": self.stopped["desktop"]}

    def restore(self, stopped):
        with self.lock:
            self.stopped = {"units": [u for u in stopped.get("units", []) if u in UNITS],
                            "remote": bool(stopped.get("remote")),
                            "desktop": stopped.get("desktop") if stopped.get("desktop") in ("hidden", "closed") else None}

    def stop_units(self):
        """Stop the UNITS that run, adding them to what the resume starts again."""
        units = [u for u in UNITS if active(u)]
        if units:
            run(["systemctl", "--user", "stop", *units], timeout=30)
            with self.lock:
                self.stopped["units"] += [u for u in units if u not in self.stopped["units"]]
            self.log(f"pause: stopped {', '.join(units)}")

    def pause(self, desktop):
        if desktop == "close":
            if desktop_running():
                run([DESKTOPS, "stop"], timeout=40)  # remote desktop goes with it
                with self.lock:
                    self.stopped["desktop"] = "closed"
                self.log("pause: desktop closed")
        else:
            if "running=1" in output([REMOTE_CTL, "status"]):
                run([REMOTE_CTL, "stop"])
                with self.lock:
                    self.stopped["remote"] = True
                self.log("pause: remote desktop stopped")
            with self.lock:
                self.stopped["desktop"] = "hidden"  # ft-screens was told already (or isn't running)
        self.stop_units()

    def enforce(self):
        """While paused: SteamVR restarted, and started the gaze service with it."""
        self.stop_units()

    def resume(self):
        stopped = self.snapshot()
        units = stopped["units"]
        if units:
            # They need SteamVR (Requisite=steamvr.service): without it they start with it next time.
            if run(["systemctl", "--user", "start", *units], timeout=30) == 0:
                self.log(f"resume: started {', '.join(units)}")
            else:
                self.log(f"resume: couldn't start {', '.join(units)} (SteamVR not running?)")
        if stopped["remote"]:
            # In a scope of its own, not this service's, so a relay restart doesn't end it.
            run(["systemd-run", "--user", "--scope", "--collect", "--quiet", REMOTE_CTL, "start"])
            self.log("resume: remote desktop started")
        if stopped["desktop"] == "closed":
            if run([DESKTOPS, "start"], timeout=60) == 0:
                self.log("resume: desktop started")
            else:
                self.log("resume: the desktop didn't start (see /tmp/frametop-session.log)")
        with self.lock:
            self.stopped = {"units": [], "remote": False, "desktop": None}


class GamePause:
    """The pause state, its automatic side, and the relay's way in. `on_change(paused)` is
    called (from the relay's thread) when the state flips, for the relay's own part."""

    def __init__(self, log, on_change, buttons):
        self.log, self.on_change, self.buttons = log, on_change, buttons
        self.lock = threading.Lock()  # the state file, written from both threads
        self.paused = False
        self.reason = ""  # what paused it: gesture, button, key, controller, command, game
        self.since = 0.0  # wall clock
        self.ends_with_game = False  # resume when the VR game ends
        self.game = False
        self.game_seen = 0.0  # last "vrgame" from the helper
        self.resume_at = None
        self.settings = settings_of({}, buttons)
        self.screens = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM | socket.SOCK_NONBLOCK)
        self.worker = Worker(log, self.save)
        self.watch = ControllerWatch(log)
        self.worker.start()
        self.watch.start()
        self.load()

    # --- state file ---
    def save(self):
        with self.lock:
            state = {"paused": self.paused, "reason": self.reason, "since": self.since,
                     "ends_with_game": self.ends_with_game, "stopped": self.worker.snapshot()}
            try:
                tmp = STATE_PATH + ".tmp"
                with open(tmp, "w") as f:
                    json.dump(state, f)
                os.replace(tmp, STATE_PATH)
            except OSError as e:
                self.log(f"pause: can't save {STATE_PATH}: {e}")

    def load(self):
        """A relay restarted while paused stays paused (the relay's part is applied by its first
        configure)."""
        try:
            with open(STATE_PATH) as f:
                state = json.load(f)
        except (OSError, ValueError):
            return
        if not isinstance(state, dict) or not state.get("paused"):
            return
        self.paused = True
        self.reason = str(state.get("reason", ""))
        self.since = float(state.get("since", 0) or 0)
        self.ends_with_game = bool(state.get("ends_with_game"))
        self.worker.restore(state.get("stopped") or {})
        self.log(f"paused since before this relay started ({self.reason})")
        self.worker.submit(self.worker.enforce)

    # --- settings ---
    def configure(self, rules):
        self.settings = settings_of(rules, self.buttons)
        self.watch.set_gesture(self.settings["gesture"])
        if not self.settings["auto"]:
            self.ends_with_game, self.resume_at = False, None

    # --- pausing and resuming ---
    def sound(self, paused):
        if self.settings["sound"] and os.path.exists(SOUNDS[paused]):
            try:
                subprocess.Popen(["pw-play", SOUNDS[paused]], stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
                                 stderr=subprocess.DEVNULL, start_new_session=True)
            except OSError:
                pass

    def tell_screens(self, paused):
        try:
            self.screens.sendto(b"pause on" if paused else b"pause off", SCREENS)
        except OSError:
            pass  # the desktop isn't running

    def pause(self, reason, now):
        if self.paused:
            return
        self.paused, self.reason, self.since = True, reason, time.time()
        self.ends_with_game = self.game and self.settings["auto"]
        self.resume_at = None
        self.log(f"Frametop paused ({reason}{', until the game ends' if self.ends_with_game else ''})")
        self.sound(True)
        self.tell_screens(True)  # at once, even when the desktop is closing: that takes a while
        self.on_change(True)
        self.worker.submit(self.worker.pause, self.settings["desktop"])
        self.save()

    def resume(self, reason, now):
        if not self.paused:
            return
        self.paused, self.resume_at, self.ends_with_game = False, None, False
        self.log(f"Frametop resumed ({reason})")
        self.sound(False)
        self.tell_screens(False)
        self.on_change(False)
        self.worker.submit(self.worker.resume)
        self.save()

    def toggle(self, reason, now):
        """Pausing only happens by itself as a game starts, so a game resumed during stays on."""
        (self.resume if self.paused else self.pause)(reason, now)

    def set(self, on, reason, now):
        if on != self.paused:
            self.toggle(reason, now)

    # --- VR games ---
    def game_state(self, running, now):
        """From the pointer helper: a VR game runs or not (on a change and every 5 s)."""
        self.game_seen = now
        if running == self.game:
            return
        self.game = running
        self.log(f"a VR game {'started' if running else 'ended'}")
        if running:
            self.resume_at = None  # the next game came before the resume
            if self.settings["auto"]:
                self.pause("game", now)
        elif self.paused and self.ends_with_game:
            self.resume_at = now + RESUME_DELAY

    def tick(self, now):
        if self.game and now - self.game_seen > GAME_STALE:
            self.game_state(False, now)  # the helper went quiet: SteamVR (or the helper) stopped
        if self.resume_at is not None and now >= self.resume_at:
            self.resume("the game ended", now)

    def timeout(self, now):
        """How long the relay may wait before tick() has something to do."""
        if self.resume_at is not None:
            return max(0.05, self.resume_at - now)
        return 1.0 if self.game else 3600.0

    def helper_started(self):
        """The pointer helper (re)started, so SteamVR did too, maybe with the gaze service."""
        if self.paused:
            self.worker.submit(self.worker.enforce)

    def status(self):
        return {"t": "pause", "paused": self.paused, "reason": self.reason, "since": self.since,
                "ends_with_game": self.ends_with_game, "game": self.game,
                "busy": self.worker.busy, "stopped": self.worker.snapshot(),
                "gesture_reader": self.watch.connected, "auto": self.settings["auto"],
                "desktop": self.settings["desktop"]}