#!/usr/bin/env python3
"""SteamVR dashboard patches (Frame Tweaks): the shared part of Curved world windows,
Back to the previous window, and any later dashboard tweak.

  sft-dashboard arm       before SteamVR starts: turn on its dashboard DevTools port (127.0.0.1:8087)
  sft-dashboard restore   after SteamVR stops: put that setting back as it was
  sft-dashboard watch     while SteamVR runs: inject every patch in patches/ (again after a
                          dashboard reload); a patch whose file is removed is stopped
  sft-dashboard status [NAME]   is the port open, is NAME (or every patch) injected?
  sft-dashboard check     is SteamVR's dashboard code still what the patches expect?

A patch's settings, if any, are ~/.config/frame-tweaks/NAME.json, handed to the page as
window.__sftPatchConfig[NAME]. Patches can ask for actions the page can't do itself through
bindings this service registers on the page (sftScroll: scrolling; see scroll() below).

A patch is patches/NAME.js: an expression that registers itself as window.__sftPatches[NAME]
= {version, stop}. Header lines "// requires: what | exact text" name pieces of SteamVR's
dashboard code it relies on; if an update changes any of them, that patch isn't injected (and
status says so) instead of poking at code it doesn't understand.
SteamVR rewrites steamvr.vrsettings from memory, so the setting is changed only while it's stopped.
"""
import base64, hashlib, json, os, re, socket, struct, subprocess, sys, threading, time, urllib.request
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:
    import i18n
    from i18n import _, N_, ngettext
except ImportError:                       # (not installed by Frame Tweaks: English)
    i18n = None
    def _(text):
        return text

    def ngettext(singular, plural, n):
        return singular if n == 1 else plural
    N_ = _

SETTINGS = Path.home() / ".config/openvr/config/steamvr.vrsettings"
STATE = Path.home() / ".local/state/frame-tweaks/vr-debugger.prev"
OLD_STATE = Path.home() / ".local/state/steam-frame-tweaks/vr-debugger.prev"   # before 0.5
PATCHES = Path(__file__).resolve().parent / "patches"
LIB = Path(__file__).resolve().parent / "lib.js"         # helpers the patches share, injected first
PORT = 8087
DASHBOARD = Path("/opt/steamvr/resources/webinterface/dashboard")
CHECK_FILE = Path(f"/run/user/{os.getuid()}/sft-dashboard.status")
CONF_DIR = Path.home() / ".config/frame-tweaks"      # patch settings: NAME.json (set in the installer window)
BASE = {"window.FrameStore": "window.FrameStore="}     # every patch starts from FrameStore


def load():
    return json.loads(SETTINGS.read_text()) if SETTINGS.exists() else {}


def write(d):
    tmp = SETTINGS.with_suffix(".sft-tmp")
    tmp.write_text(json.dumps(d, indent=3) + "\n")
    os.replace(tmp, SETTINGS)


def arm():
    d = load()
    cur = d.get("VRWebHelper", {}).get("DebuggerEnabled")
    if cur is True:
        return                          # already on (by you or another tool): leave it alone
    if not STATE.exists():
        STATE.parent.mkdir(parents=True, exist_ok=True)
        STATE.write_text(json.dumps(cur))   # null = the key wasn't there
    d.setdefault("VRWebHelper", {})["DebuggerEnabled"] = True
    write(d)


def restore():
    if not STATE.exists() and OLD_STATE.exists():      # armed by an older version
        STATE.parent.mkdir(parents=True, exist_ok=True)
        OLD_STATE.replace(STATE)
    if not STATE.exists():
        return
    try:
        prev = json.loads(STATE.read_text())
    except ValueError:                  # (empty or garbled, say cut short by a crash: as if the key
        prev = None                     # wasn't there, so the port goes off, SteamVR's default)
    d = load()
    sec = d.setdefault("VRWebHelper", {})
    if prev is None:
        sec.pop("DebuggerEnabled", None)
        if not sec:
            d.pop("VRWebHelper")
    else:
        sec["DebuggerEnabled"] = prev
    write(d)
    STATE.unlink()


# ---- minimal DevTools (WebSocket) client, stdlib only --------------------------------------
class Ws:
    """One WebSocket connection to a DevTools target (text frames, client-side masking)."""

    def __init__(self, url, timeout=5):
        host, path = url.split("://", 1)[1].split("/", 1)
        h, p = host.split(":")
        self.s = socket.create_connection((h, int(p)), timeout=timeout)
        key = base64.b64encode(os.urandom(16)).decode()
        self.s.sendall(f"GET /{path} HTTP/1.1\r\nHost: {host}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"
                       f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n".encode())
        buf = b""
        while b"\r\n\r\n" not in buf:
            chunk = self.s.recv(4096)
            if not chunk:
                raise OSError("handshake closed")
            buf += chunk
        if b" 101 " not in buf.split(b"\r\n", 1)[0]:
            raise OSError("handshake refused")
        self.rest = buf.split(b"\r\n\r\n", 1)[1]

    def send(self, obj):
        msg = json.dumps(obj).encode()
        mask, n = os.urandom(4), len(msg)
        head = bytes([0x81]) + (bytes([0x80 | n]) if n < 126 else
                                bytes([0x80 | 126]) + struct.pack(">H", n) if n < 65536 else
                                bytes([0x80 | 127]) + struct.pack(">Q", n))
        self.s.sendall(head + mask + bytes(b ^ mask[i % 4] for i, b in enumerate(msg)))

    def _read(self, k):
        while len(self.rest) < k:
            chunk = self.s.recv(65536)
            if not chunk:
                raise OSError("closed")
            self.rest += chunk
        out, self.rest = self.rest[:k], self.rest[k:]
        return out

    def recv(self):
        """The next text message, as JSON."""
        while True:
            b0, b1 = self._read(2)
            n = b1 & 0x7F
            if n == 126:
                n = struct.unpack(">H", self._read(2))[0]
            elif n == 127:
                n = struct.unpack(">Q", self._read(8))[0]
            data = self._read(n)
            if b0 & 0x0F == 8:
                raise OSError("closed by the page")
            if b0 & 0x0F == 1:
                return json.loads(data)

    def close(self):
        self.s.close()


def ws_eval(url, expr, timeout=5):
    ws = Ws(url, timeout)
    try:
        ws.send({"id": 1, "method": "Runtime.evaluate", "params": {"expression": expr, "returnByValue": True}})
        while True:
            reply = ws.recv()
            if reply.get("id") == 1:
                r = reply.get("result", {}).get("result", {})
                return r.get("value", r.get("description"))
    finally:
        ws.close()


# ---- page -> here: actions a patch asks for (window.<binding>(json)) ------------------------
# Scrolling. Injected pointer input (XTest, ydotool's wheel, the desktop's Remote Desktop portal)
# never visibly scrolls anything in the Steam session: in VR gamescope takes pointer input only from
# the lasers. Keyboard input does arrive, so the strips scroll with the arrow keys (PgUp/PgDn for a
# page), through ydotool, in the window you last clicked (the desktop included).
YDOTOOL = Path.home() / ".local/bin/ydotool"
YDOTOOL_SOCKET = f"/run/user/{os.getuid()}/.ydotool_socket"
KEY = {"up": 103, "down": 108, "left": 105, "right": 106, "pgup": 104, "pgdn": 109}


def keys(name, n):
    if n > 0 and YDOTOOL.exists():
        seq = [f"{KEY[name]}:1", f"{KEY[name]}:0"] * n
        subprocess.run([str(YDOTOOL), "key", *seq], env=dict(os.environ, YDOTOOL_SOCKET=YDOTOOL_SOCKET),
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=5)


def scroll(req):
    """Scroll strips: {dx, dy (steps; + = right / up), page (+1 up / -1 down)}."""
    dx, dy = (max(-60, min(60, int(req.get(k, 0)))) for k in ("dx", "dy"))
    page = int(req.get("page", 0))
    if page:
        keys("pgup" if page > 0 else "pgdn", 1)
    if dy:
        keys("up" if dy > 0 else "down", abs(dy))
    if dx:
        keys("right" if dx > 0 else "left", abs(dx))


def save_config(req):
    """A patch's settings panel in VR saves its settings: {name, config}. Only for an installed
    patch, a small JSON object; the next look (within 5 s) hands it to the page as usual."""
    name, conf = req.get("name"), req.get("config")
    if not (isinstance(name, str) and name in patches() and isinstance(conf, dict)):
        return
    text = json.dumps(conf)
    if len(text) > 4096:
        return
    CONF_DIR.mkdir(parents=True, exist_ok=True)
    tmp = CONF_DIR / f".{name}.json.tmp"
    tmp.write_text(text + "\n")
    tmp.replace(CONF_DIR / f"{name}.json")


BINDINGS = {"sftScroll": scroll, "sftConfig": save_config}


def listen():
    """Keep a connection to the dashboard page with the bindings registered, and act on calls."""
    while True:
        try:
            url = systemui()
            if not url:
                time.sleep(3)
                continue
            ws = Ws(url, timeout=None)
            # Runtime.enable first: the page gets the bindings again after a reload only while the
            # Runtime domain is on (without it the strips came back but didn't scroll, and the gear
            # said to reinstall). (It also sends the page's console messages: read and dropped below.)
            ws.send({"id": 1, "method": "Runtime.enable"})
            for i, name in enumerate(BINDINGS):
                ws.send({"id": i + 2, "method": "Runtime.addBinding", "params": {"name": name}})
            while True:
                m = ws.recv()
                if m.get("method") == "Runtime.bindingCalled" and m["params"].get("name") in BINDINGS:
                    # (Any script on the page can call these, with anything: a bad call is logged
                    # and skipped. Nothing restarts this thread, so nothing may end it.)
                    try:
                        BINDINGS[m["params"]["name"]](json.loads(m["params"].get("payload") or "{}"))
                    except Exception as e:
                        print(f"binding {m['params'].get('name')}: {e!r}", flush=True)
        except Exception:                 # (SteamVR or its page went away, or a bad message)
            time.sleep(3)


def systemui():
    with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/json/list", timeout=3) as r:
        targets = json.load(r)
    return next((t["webSocketDebuggerUrl"] for t in targets
                 if t.get("title") == "systemui" and t.get("webSocketDebuggerUrl")), None)


def patches():
    return {p.stem: p for p in sorted(PATCHES.glob("*.js"))}


# The patches' own text: T('...') and N('...') in their JS (dev/i18n.py collects the same).
JS_TEXT = re.compile(r"""\b[TN]\(\s*(?:'((?:[^'\\]|\\.)*)'|"((?:[^"\\]|\\.)*)")\s*\)""")
_js_texts = {}


def js_string(body):
    """A JS string literal's text (its quotes taken off)."""
    try:
        return json.loads('"' + body.replace("\\'", "'").replace('"', '\\"').replace('\\\\"', '\\"') + '"')
    except ValueError:
        return body


def i18n_expr(found):
    """window.__sftI18n = {lang, t: {English: translation}} for the patches' text, in the
    language the keyboard asks for now (lib/i18n.py looks no more than every 5 s)."""
    table, lang = {}, None
    if i18n is not None:
        i18n.refresh()
        lang = i18n.CATALOGUE.language
        for p in found.values():
            mtime = p.stat().st_mtime
            if _js_texts.get(p, (None,))[0] != mtime:
                _js_texts[p] = (mtime, {js_string(a or b) for a, b in JS_TEXT.findall(p.read_text(errors="replace"))})
            for text in _js_texts[p][1]:
                if _(text) != text:
                    table[text] = _(text)
    return f"window.__sftI18n = {json.dumps({'lang': lang, 't': table})}; 'i18n'"


def requires(path):
    req = dict(BASE)
    for m in re.finditer(r"^// requires: (.+?) \| (.+)$", path.read_text(errors="replace"), re.M):
        req[m[1]] = m[2]
    return req


def dashboard_code():
    """The code of the dashboard page the patches go into: the scripts systemui.html loads (not
    every bundle in the folder, some of which other pages load: a marker found only there isn't in
    the page). None without it (no files at all, the folder gone or moved)."""
    try:
        page = (DASHBOARD / "systemui.html").read_text(errors="replace")
        root = DASHBOARD.resolve()
        files = []
        for src in re.findall(r"""<script\b[^>]*?\bsrc\s*=\s*["']([^"'?#]+)""", page, re.I):
            f = (DASHBOARD / src).resolve()
            if "://" not in src and root in f.parents and f.is_file():
                files.append(f)
        return "".join(f.read_text(errors="replace") for f in files) or None
    except OSError:
        return None


def missing(path, code):
    """Pieces of SteamVR's dashboard code the patch needs but can't find ([] = compatible)."""
    if code is None:
        return ["SteamVR's dashboard files"]
    return [name for name, marker in requires(path).items() if marker not in code]


def stop_expr(name):
    n = json.dumps(name)          # (unregistered even if its stop() throws: ticked again, it's injected anew)
    return f"try {{ window.__sftPatches?.[{n}]?.stop?.(); }} catch (e) {{}} delete window.__sftPatches?.[{n}]; 'stopped'"


def config_expr(name, path):
    """window.__sftPatchConfig[NAME] = patches/NAME.json, and tell the patch if it changed."""
    try:
        conf = json.loads(path.read_text())
    except (OSError, ValueError):
        return None
    n, c = json.dumps(name), json.dumps(conf)
    return (f"(() => {{ const all = (window.__sftPatchConfig ??= {{}}); const c = {c};"
            f" if (JSON.stringify(all[{n}]) !== JSON.stringify(c)) {{ all[{n}] = c;"
            f" window.__sftPatches?.[{n}]?.configChanged?.(); }} }})()")


def watch():
    code, last, injected = dashboard_code(), None, None      # (injected: None until the page is seen)
    threading.Thread(target=listen, daemon=True).start()
    while True:
        found = patches()
        lib_bad = missing(LIB, code) if LIB.exists() else ["lib.js"]
        bad = {n: lib_bad + missing(p, code) for n, p in found.items()}
        CHECK_FILE.write_text("".join(f"{n}: SteamVR changed (missing: {', '.join(m)}); off\n"
                                      for n, m in bad.items() if m))
        try:
            url = systemui()
            if not url:
                result = "waiting for the dashboard"
            else:
                results = []
                if injected is None:
                    # Just started: the patches an earlier run left on the page count too (a tweak
                    # removed moments before this service restarted, as an install right after does).
                    names = ws_eval(url, "Object.keys(window.__sftPatches ?? {})")
                    injected = set(names) if isinstance(names, list) else set()
                for n in sorted(injected - set(found)):          # its tweak was removed
                    ws_eval(url, stop_expr(n))
                    injected.discard(n)
                if found and not lib_bad:
                    lib = LIB.read_text(errors="replace")  # (it skips itself if this very file is loaded)
                    wanted = hashlib.sha256(lib.encode()).hexdigest()[:16]
                    ws_eval(url, f"window.__sftLibWanted = {json.dumps(wanted)};\n" + lib)
                    ws_eval(url, i18n_expr(found))         # (the patches redraw with it)
                for n, p in found.items():
                    if bad[n]:
                        results.append(f"{n}: off (SteamVR changed)")
                        continue
                    conf = config_expr(n, CONF_DIR / f"{n}.json")
                    if conf:
                        ws_eval(url, conf)
                    results.append(f"{n}: {ws_eval(url, p.read_text(errors='replace'))}")
                    injected.add(n)
                result = "; ".join(results) or "no patches installed"
        except Exception as e:           # (OSError, and a reply cut short as SteamVR starts or stops)
            result = f"waiting for SteamVR's DevTools port ({e.__class__.__name__})"
        if result != last:
            print(result, flush=True)
            last = result
        time.sleep(5)


def status(name=None):
    if name and name not in patches():
        print(_("{name}: not installed").format(name=name))
        return 1
    for line in (CHECK_FILE.read_text().splitlines() if CHECK_FILE.exists() else []):
        if name is None or line.startswith(f"{name}:"):
            print(line)
            return 2
    try:
        url = systemui()
    except OSError:
        print("port closed (SteamVR not running, or not restarted since install)")
        return 2
    if not url:
        print(_("dashboard page not found"))
        return 2
    names = [name] if name else list(patches())
    ok = True
    for n in names:
        on = ws_eval(url, f"!!window.__sftPatches?.[{json.dumps(n)}]") is True
        # (": patched" and "port closed" stay English: lib/common.sh and status.sh match them)
        print(f"{n}: patched" if on else _("{name}: port open, not patched yet").format(name=n))
        ok = ok and on
    return 0 if ok else 2


def check():
    code = dashboard_code()
    bad = {n: missing(p, code) for n, p in patches().items()}
    bad["lib.js (shared helpers)"] = missing(LIB, code) if LIB.exists() else ["the file itself"]
    bad = {n: m for n, m in bad.items() if m}
    print("compatible" if not bad else "; ".join(f"{n} missing: {', '.join(m)}" for n, m in bad.items()))
    return bool(bad)


if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else ""
    fn = {"arm": arm, "restore": restore, "watch": watch, "check": check,
          "status": lambda: status(sys.argv[2] if len(sys.argv) > 2 else None)}.get(cmd)
    if fn is None:
        sys.exit(__doc__)
    sys.exit(fn() or 0)
