#!/usr/bin/env python3
"""phone-touch: in VR, laser clicks on the phone mirror land where you aim.

The Steam session's compositor (gamescope) maps the laser wrongly on a window narrower than its
1920x1080 screen: it adds the gap it would leave beside the window if it were centred on that
screen, so on the upright phone (485 wide) every click landed at the left edge. phone-fit makes
the window that narrow, so instead of the window taking the laser, a see-through layer over it in
SteamVR's dashboard does (the dashboard patch phone-touch.js): it reports where you press, drag and
let go, and this sends those to the phone as touches, the way scrcpy itself does: through a second
scrcpy server on the phone with only its control channel (no video), so nothing passes through the
Steam session's input at all.

  phone-touch DEVICE      DEVICE: the phone's adb address (ip:port); runs until scrcpy exits
  phone-touch --check     exit 1 (and say why, once per scrcpy version) if the installed scrcpy
                          isn't one this was tested with; `phone` then shows a plain window

The layer appears only while this runs: it calls in through the dashboard's DevTools port
(127.0.0.1:8087, opened by the VR environment helper) and says it's alive every 2 s. While the laser
rests on the mirror, the joystick scrolls the phone (read from the controller model, see Sticks):
SteamVR doesn't pass the joystick's scrolling to the layer.
"""
import importlib.machinery
import importlib.util
import json
import math
import os
import random
import re
import select
import shutil
import signal
import socket
import struct
import subprocess
import sys
import threading
import time
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 ftconf  # noqa: E402  (settings files: lib/ftconf.py, installed beside i18n.py)

DASHBOARD = Path.home() / ".local/share/frame-tweaks/vr-dashboard/sft-dashboard"
BINDING = "sftPhoneTouch"
JAR = "/data/local/tmp/sft-phone-touch.jar"     # scrcpy deletes its own copy once it has started
PIDFILE = Path(f"/run/user/{os.getuid()}/sft-phone-touch.pid")
DOWN, UP, MOVE = 0, 1, 2
FINGER = -2                                     # scrcpy's "generic finger" pointer id
# scrcpy versions (major) whose control protocol this was tested with: the touch and scroll
# messages below and the server's options. scrcpy has changed them between major versions
# before, and a mismatch would make taps go nowhere without an error, so with any other version
# this stays off and the mirror is a plain window. PHONE_TOUCH_ANY_SCRCPY=1 tries anyway.
# (lib/tested-versions records the exact version; tests check the two agree.)
SCRCPY_TESTED = ("5",)
STATE = Path.home() / ".local/state/frame-tweaks"


def log(*a):
    print("phone-touch:", *a, file=sys.stderr, flush=True)


def adb(dev, *args, **kw):
    # A phone that doesn't answer in time is an OSError, like one that refuses: the callers try
    # again later. (TimeoutExpired isn't one, and it ended phone-touch, or its joystick thread.)
    try:
        return subprocess.run(["adb", "-s", dev, *args], capture_output=True, text=True, timeout=20, **kw)
    except subprocess.TimeoutExpired:
        raise OSError(f"adb {args[0]}: no answer within 20 s") from None


def scrcpy_running():
    return subprocess.run(["pgrep", "-u", str(os.getuid()), "-x", "scrcpy"], stdout=subprocess.DEVNULL).returncode == 0


def scrcpy_server():
    """(version, path of the server file) of the installed scrcpy: the two must match."""
    exe = shutil.which("scrcpy") or "/home/linuxbrew/.linuxbrew/bin/scrcpy"
    out = subprocess.run([exe, "--version"], capture_output=True, text=True, timeout=10).stdout
    m = re.match(r"scrcpy (\S+)", out)
    jar = os.environ.get("SCRCPY_SERVER_PATH") or Path(os.path.realpath(exe)).parent.parent / "share/scrcpy/scrcpy-server"
    if not m or not Path(jar).is_file():
        sys.exit(_("phone-touch: can't find scrcpy's server"))
    return m.group(1), str(jar)


def tested(version):
    return version.split(".")[0] in SCRCPY_TESTED or os.environ.get("PHONE_TOUCH_ANY_SCRCPY") == "1"


def check():
    """0 if the installed scrcpy is a tested one; else say so (a notification once per version)."""
    version, _jar = scrcpy_server()
    if tested(version):
        return 0
    why = _("scrcpy {version} isn't a version Frame Tweaks has been tested with "
            "({tested}), so the mirror's VR extras are off: aimed laser "
            "clicks and joystick scrolling. It works as a plain window until a Frame Tweaks update.").format(
                version=version, tested=", ".join(v + ".x" for v in SCRCPY_TESTED))
    log(why)
    seen = STATE / "phone-touch-untested"
    try:
        if seen.read_text().strip() == version:
            return 1
    except OSError:
        pass
    try:
        STATE.mkdir(parents=True, exist_ok=True)
        seen.write_text(version + "\n")
        subprocess.run(["notify-send", "-a", _("Android Mirror"), _("VR extras are off"), why], timeout=10)
    except (OSError, subprocess.SubprocessError):
        pass
    return 1


def screen_size(dev):
    """The phone screen's size in pixels, upright (an "Override size" from `wm size` wins)."""
    sizes = re.findall(r"(\d+)x(\d+)", adb(dev, "shell", "wm", "size").stdout)
    if not sizes:
        raise OSError("no screen size")
    w, h = map(int, sizes[-1])
    return min(w, h), max(w, h)


class Phone:
    """A control-only scrcpy server on the phone and the connection to it."""

    def __init__(self, dev):
        self.dev, self.sock, self.server, self.port = dev, None, None, None
        self.version, self.jar = scrcpy_server()
        self.lock = threading.Lock()

    def connect(self):
        self.close()
        self.size = screen_size(self.dev)
        r = adb(self.dev, "push", self.jar, JAR)
        if r.returncode:
            raise OSError(f"push: {r.stderr.strip()}")
        scid = random.randrange(1, 0x7FFFFFFF)
        self.server = subprocess.Popen(
            ["adb", "-s", self.dev, "shell", f"CLASSPATH={JAR}", "app_process", "/", "com.genymobile.scrcpy.Server",
             self.version, f"scid={scid:08x}", "log_level=warn", "video=false", "audio=false", "control=true",
             "tunnel_forward=true", "clipboard_autosync=false", "cleanup=false", "power_on=false",
             "send_device_meta=false"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        # The tunnel ends in a socket file in your runtime folder (only you can reach it), not a
        # TCP port on 127.0.0.1, which any local program (sandboxed apps included) could connect to.
        path = f"/run/user/{os.getuid()}/sft-phone-touch-{scid:08x}.sock"
        r = adb(self.dev, "forward", f"localfilesystem:{path}", f"localabstract:scrcpy_{scid:08x}")
        if r.returncode:
            raise OSError(f"adb forward failed: {r.stderr.strip()}")
        self.port = path
        for _ in range(100):       # the tunnel accepts before the server listens: wait for its first byte
            s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            try:
                s.settimeout(3)
                s.connect(path)
                if s.recv(1) == b"\0":
                    self.sock = s
                    return
            except OSError:
                pass
            s.close()
            time.sleep(0.1)
        raise OSError("the phone's server didn't answer")

    def close(self):
        if self.sock:
            self.sock.close()
            self.sock = None
        if self.port:
            adb(self.dev, "forward", "--remove", f"localfilesystem:{self.port}")
            try:
                os.unlink(self.port)
            except OSError:
                pass
            self.port = None
        if self.server:
            self.server.terminate()
            self.server = None

    def point(self, u, v, sideways):
        """u, v: 0..1 across and down the picture as you see it; sideways: the phone is turned."""
        w, h = self.size[::-1] if sideways else self.size
        return round(min(1, max(0, u)) * (w - 1)), round(min(1, max(0, v)) * (h - 1)), w, h

    def touch(self, action, u, v, sideways):
        self.send(lambda: struct.pack(">BBqiiHHHII", 2, action, FINGER, *self.point(u, v, sideways),
                                      0 if action == UP else 0xFFFF, 0, 0), retry=action != MOVE)

    def scroll(self, u, v, sideways, dx, dy):
        """A mouse wheel at u, v: dx, dy in notches, + = right / down (as a page's wheel event)."""
        # scrcpy's fixed point: -1..1 carries -16..16 notches; Android's + scrolls up / right.
        fix = lambda n: max(-0x8000, min(0x7FFF, round(max(-16, min(16, n)) / 16 * 0x8000)))
        self.send(lambda: struct.pack(">BiiHHhhI", 3, *self.point(u, v, sideways), fix(dx), fix(-dy), 0), retry=False)

    def send(self, message, retry):
        """message(): the bytes, made once connected (they carry the screen size connect() reads,
        so a first connection that failed is made at the first touch, as main() expects)."""
        with self.lock:
            for attempt in (1, 2):
                try:
                    if not self.sock:
                        self.connect()
                    self.sock.sendall(message())
                    return
                except OSError as e:
                    log("reconnecting:", e)
                    self.close()
                    if attempt == 2 or not retry:
                        return


# The gear's settings panel (phone-touch.js): its choices, as lines of ~/.config/frame-tweaks/phone.conf (which
# `phone` reads). Only these keys are written; anything else in the file is kept as it is.
CONF = Path.home() / ".config/frame-tweaks/phone.conf"      # (before 0.6: ~/.config/phone.conf)
# Frame rates: the Steam Frame display's refresh rates, plus 30 and 60.
CHOICES = {"PHONE_MAX_SIZE": ("720", "800", "1024", "1280", "1600", "1920"),
           "PHONE_BIT_RATE": ("1M", "1500K", "2M", "3M", "4M", "6M", "8M"),
           "PHONE_MAX_FPS": ("30", "60", "72", "80", "90", "96", "108", "120", "144"),
           "PHONE_STICK_SPEED": ("8", "15", "25"), "PHONE_CLIPBOARD": ("both", "off")}
DEFAULTS = {"PHONE_MAX_SIZE": "1024", "PHONE_BIT_RATE": "2M", "PHONE_MAX_FPS": "60", "PHONE_STICK_SPEED": "15",
            "PHONE_CLIPBOARD": "both"}
PHONE_PIDFILE = Path(os.environ.get("XDG_RUNTIME_DIR") or f"/run/user/{os.getuid()}") / "sft-phone.pid"   # written by `phone` (same folder)


def read_conf(path=CONF):
    conf = dict(DEFAULTS)
    conf.update({k: v for k, v in ftconf.read_shell(path).items() if k in CHOICES and v in CHOICES[k]})
    return conf


def write_conf(changes, path=CONF):
    """Set these keys (allowed values only) in the file, keeping its other lines."""
    ftconf.set_shell(path, {k: str(v) for k, v in changes.items() if k in CHOICES and str(v) in CHOICES[k]})


def reset_conf(path=CONF):
    """Back to the defaults: remove our keys' lines, keeping the rest of the file."""
    if path.exists():
        ftconf.set_shell(path, drop=CHOICES)


def running(pidfile, name):
    """The pid in pidfile, if that process is still `name` (not a stale file whose pid was reused)."""
    try:
        pid = int(pidfile.read_text())
        cmd = Path(f"/proc/{pid}/cmdline").read_bytes().split(b"\0")
    except (OSError, ValueError):
        return None
    return pid if any(Path(a.decode(errors="replace")).name == name for a in cmd[:3]) else None


def restart_stream():
    """Ask `phone` to start scrcpy again (with the settings as they are now)."""
    pid = running(PHONE_PIDFILE, "phone")
    if pid:
        os.kill(pid, 10)                                              # SIGUSR1
    else:
        log("restart: phone isn't running")


class Sticks:
    """The controllers' thumbsticks, read the way Arrow Pad reads the grips: SteamVR doesn't hand the
    dashboard (or a background program) the joystick on the Frame, but the controller model's
    "thumbstick" part still tilts. Its rotation from rest: about 20 degrees at full travel, up/down
    about the part's x axis, left/right mostly about its z axis."""
    FULL_DEG = 20.0

    def __init__(self):
        import openvr                             # in the shared venv (see install.sh); None without
        self.vr = openvr
        openvr.init(openvr.VRApplication_Background)
        self.sys, self.rm, self.inp = openvr.VRSystem(), openvr.VRRenderModels(), openvr.VRInput()
        self.mode = openvr.RenderModel_ControllerMode_State_t()
        self.hands = [(role, self.inp.getInputSourceHandle(f"/user/hand/{hand}")) for hand, role in
                      (("left", openvr.TrackedControllerRole_LeftHand), ("right", openvr.TrackedControllerRole_RightHand))]

    def tilt(self):
        """(right, up) of the stick pushed furthest, each -1..1; (0, 0) when none is."""
        best = (0.0, 0.0)
        for role, path in self.hands:
            try:
                d = self.sys.getTrackedDeviceIndexForControllerRole(role)
                if d == self.vr.k_unTrackedDeviceIndexInvalid or not self.sys.isTrackedDeviceConnected(d):
                    continue
                model = self.sys.getStringTrackedDeviceProperty(d, self.vr.Prop_RenderModelName_String)
                ok1, cur = self.rm.getComponentStateForDevicePath(model, "thumbstick", path, self.mode)
                ok2, rest = self.rm.getComponentState(model, "thumbstick", self.vr.VRControllerState_t(), self.mode)
            except Exception:
                continue
            if not (ok1 and ok2):
                continue
            x, _, z = rotation_vector(cur.mTrackingToComponentRenderModel.m, rest.mTrackingToComponentRenderModel.m)
            t = (max(-1.0, min(1.0, -z / self.FULL_DEG * 1.2)), max(-1.0, min(1.0, -x / self.FULL_DEG)))
            if max(map(abs, t)) > max(map(abs, best)):
                best = t
        return best


def rotation_vector(a, b):
    """Rotation from b to a (3x4 matrices, rotation part), as axis x angle in degrees."""
    r = [[sum(a[i][k] * b[j][k] for k in range(3)) for j in range(3)] for i in range(3)]
    angle = math.acos(max(-1.0, min(1.0, (r[0][0] + r[1][1] + r[2][2] - 1) / 2)))
    if angle < 1e-6:
        return 0.0, 0.0, 0.0
    k = math.degrees(angle) / (2 * math.sin(angle))
    return (r[2][1] - r[1][2]) * k, (r[0][2] - r[2][0]) * k, (r[1][0] - r[0][1]) * k


def stick_rate(t, dead=0.2, top=15.0):
    """Wheel notches a second for a stick tilt t (-1..1): none in the dead zone, then faster
    the further it's pushed."""
    m = (abs(t) - dead) / (1 - dead)
    return 0.0 if m <= 0 else math.copysign(top * m ** 1.6, t)


def scroll_with_sticks(phone, state):
    """While the laser is on the mirror (and not pressing), the joystick scrolls the phone there:
    up/down or left/right, whichever it's pushed more."""
    try:
        sticks = Sticks()
    except Exception as e:                        # no openvr, or SteamVR isn't running
        log("no joystick scrolling:", e)
        return
    tick, checked = 1 / 30, 0.0
    while True:
        time.sleep(tick)
        # Whether scrcpy still runs: every 2 s, as serve() asks (a pgrep takes some 20 ms of
        # processor time; 30 a second took a good part of a core for as long as the mirror ran).
        if time.monotonic() - checked > 2:
            if not scrcpy_running():
                return
            checked = time.monotonic()
        at = state.get("hover")
        quiet = time.monotonic() - state.get("since", 0)
        # (A press with no drag or let-go for 3 s: its let-go went missing, don't wait for it.)
        if not at or (state.get("pressed") and quiet < 3) or quiet > 60:
            continue
        right, up = sticks.tilt()
        top = state.get("speed", 15.0)
        if abs(up) >= abs(right):
            dy = -stick_rate(up, top=top) * tick  # stick up: scroll up (the wheel's "up")
            dx = 0.0
        else:
            dx, dy = stick_rate(right, top=top) * tick, 0.0
        if dx or dy:
            phone.scroll(*at, dx, dy)


def dashboard():
    loader = importlib.machinery.SourceFileLoader("sft_dashboard", str(DASHBOARD))
    spec = importlib.util.spec_from_loader("sft_dashboard", loader)
    mod = importlib.util.module_from_spec(spec)
    loader.exec_module(mod)
    return mod


def publish(ws, state):
    """Hand the page the settings as they are in the file (and the joystick speed to the sticks)."""
    conf = read_conf()
    state["speed"] = float(conf["PHONE_STICK_SPEED"])
    ws.send({"id": 900000 + int(time.monotonic() * 10) % 99999, "method": "Runtime.evaluate",
             "params": {"expression": f"window.__sftPhoneTouchConf = {json.dumps(conf)}"}})


def serve(phone, sd, state):
    """Take the layer's calls from the dashboard page; tell it every 2 s that we're here."""
    while scrcpy_running():
        try:
            url = sd.systemui()
            if not url:
                time.sleep(3)
                continue
            ws = sd.Ws(url, timeout=2)
            # (Runtime.enable first: without it a reloaded page doesn't get the binding again)
            ws.send({"id": 1, "method": "Runtime.enable"})
            ws.send({"id": 2, "method": "Runtime.addBinding", "params": {"name": BINDING}})
            n, beat = 2, 0.0
            publish(ws, state)
            while True:
                now = time.monotonic()
                if now - beat > 2:
                    if not scrcpy_running():
                        break
                    n += 1
                    ws.send({"id": n, "method": "Runtime.evaluate",
                             "params": {"expression": "window.__sftPhoneTouchAlive = Date.now()"}})
                    beat = now
                # Wait for a message without a timeout cutting one in half.
                if not ws.rest and not select.select([ws.s], [], [], 0.5)[0]:
                    continue
                m = ws.recv()
                if m.get("method") == "Runtime.bindingCalled" and m["params"].get("name") == BINDING:
                    try:
                        req = json.loads(m["params"].get("payload") or "{}")
                        if req["a"] in ("set", "reset", "restart"):
                            if req["a"] == "set":
                                write_conf(req.get("conf") or {})
                            elif req["a"] == "reset":
                                reset_conf()
                            else:
                                restart_stream()
                            publish(ws, state)
                            n += 1
                            continue
                        if req["a"] == "leave":
                            state["hover"] = None
                            continue
                        at = (float(req["u"]), float(req["v"]), float(req["w"]) > float(req["h"]))
                        if req["a"] == "scroll":          # (from the scroll strips: not where the laser is)
                            phone.scroll(*at, float(req.get("dx", 0)), float(req.get("dy", 0)))
                            continue
                        state.update(hover=at, since=time.monotonic(), pressed=req["a"] in ("down", "move"))
                        if req["a"] != "hover":
                            phone.touch({"down": DOWN, "move": MOVE, "up": UP}[req["a"]], *at)
                    except (KeyError, ValueError, TypeError) as e:
                        log("bad call:", e)
            ws.send({"id": n + 1, "method": "Runtime.evaluate", "params": {"expression": "window.__sftPhoneTouchAlive = 0"}})
            ws.close()
            return
        except (OSError, ValueError) as e:
            log("dashboard:", e)
            time.sleep(3)


def main():
    if len(sys.argv) != 2:
        sys.exit(__doc__)
    if sys.argv[1] == "--check":
        sys.exit(check())
    if not DASHBOARD.is_file():
        return                                    # no VR environment helper: no layer to talk to
    # One at a time: a reconnect (new address) replaces the one before. The one replaced cleans
    # up as when scrcpy closes (Python's own SIGTERM ends it at once, leaving its tunnel in adb).
    signal.signal(signal.SIGTERM, lambda *a: sys.exit(0))
    old = running(PIDFILE, "phone-touch")
    if old and old != os.getpid():
        try:
            os.kill(old, 15)
        except OSError:
            pass
    PIDFILE.write_text(str(os.getpid()))
    phone = Phone(sys.argv[1])
    try:
        for _ in range(30):                       # scrcpy may still be starting
            if scrcpy_running():
                break
            time.sleep(1)
        try:
            phone.connect()                       # now, so the first tap isn't slow
        except OSError as e:
            log("phone:", e)                      # tried again at the first touch
        state = {}                                # where the laser is on the mirror (from the layer)
        threading.Thread(target=scroll_with_sticks, args=(phone, state), daemon=True).start()
        serve(phone, dashboard(), state)
    finally:
        phone.close()
        try:
            if PIDFILE.read_text() == str(os.getpid()):
                PIDFILE.unlink()
        except OSError:
            pass


if __name__ == "__main__":
    main()
