#!/usr/bin/env python3
"""
trackscreen - use a laptop touchscreen as a multitouch touchpad.

How it works
------------
* Touchscreen mode (default): the touchscreen is only *read* passively; the
  compositor keeps receiving every touch exactly as before (full multitouch).
* Trackpad mode: the touchscreen is grabbed (EVIOCGRAB) so the compositor stops
  seeing it, and its touches are re-emitted through a virtual multitouch
  *touchpad* (uinput). libinput then treats it like any other touchpad, so you
  get relative pointer motion, acceleration, tap-to-click, two-finger
  scrolling, pinch-zoom, 3/4-finger swipes, ... all configurable in your
  desktop's normal touchpad settings.  The grab is released automatically by
  the kernel if this program exits or crashes.

Toggle between the modes with a five-finger tap (configurable), by sending
SIGUSR1 to the running instance, or by running `trackscreen.py --toggle`
(handy for a keyboard shortcut).
"""

import argparse
import array
import fcntl
import math
import os
import select
import shutil
import signal
import socket
import subprocess
import sys
import time

try:
    import evdev
    from evdev import AbsInfo, UInput, ecodes as e
except ImportError:  # pragma: no cover
    sys.exit("python-evdev is required (pip install evdev / pacman -S python-evdev)")

VIRTUAL_NAME = "Trackscreen Virtual Touchpad"
UNITS_PER_MM = 100  # resolution of the virtual touchpad's coordinate space
TOUCHSCREEN, TRACKPAD = "touchscreen", "trackpad"
TOOL_BUTTONS = [e.BTN_TOOL_FINGER, e.BTN_TOOL_DOUBLETAP, e.BTN_TOOL_TRIPLETAP,
                e.BTN_TOOL_QUADTAP, e.BTN_TOOL_QUINTTAP]

DEBUG = False


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


def dbg(*a):
    if DEBUG:
        log(*a)


# --------------------------------------------------------------------------
# Device discovery
# --------------------------------------------------------------------------

def udev_props(path):
    """Read udev properties of a device node straight from the udev database."""
    props = {}
    try:
        rdev = os.stat(path).st_rdev
        with open(f"/run/udev/data/c{os.major(rdev)}:{os.minor(rdev)}") as f:
            for line in f:
                if line.startswith("E:"):
                    k, _, v = line[2:].strip().partition("=")
                    props[k] = v
    except OSError:
        pass
    return props


def is_touchscreen(dev):
    caps = dev.capabilities(absinfo=False)
    absc = set(caps.get(e.EV_ABS, []))
    keys = set(caps.get(e.EV_KEY, []))
    if not {e.ABS_MT_SLOT, e.ABS_MT_POSITION_X, e.ABS_MT_POSITION_Y,
            e.ABS_MT_TRACKING_ID} <= absc:
        return False  # we need multitouch protocol B
    if keys & {e.BTN_TOOL_PEN, e.BTN_STYLUS, e.BTN_TOOL_RUBBER}:
        return False  # a pen/stylus interface
    return e.INPUT_PROP_DIRECT in dev.input_props()


def candidate_devices():
    """Yield (path, InputDevice|None, udev_props); device is None if not readable."""
    # (evdev.list_devices() silently skips unreadable nodes; we want to report those)
    paths = [os.path.join("/dev/input", n) for n in os.listdir("/dev/input")
             if n.startswith("event") and n[5:].isdigit()]
    for path in sorted(paths, key=lambda p: int(p.rsplit("event", 1)[-1])):
        props = udev_props(path)
        try:
            dev = evdev.InputDevice(path)
        except PermissionError:
            yield path, None, props
            continue
        except OSError:
            continue
        yield path, dev, props


def find_touchscreen(spec=None):
    """Return an InputDevice for the touchscreen, or raise SystemExit with advice."""
    denied = []
    for path, dev, props in candidate_devices():
        if dev is None:
            if props.get("ID_INPUT_TOUCHSCREEN") == "1":
                denied.append(path)
            continue
        if dev.name == VIRTUAL_NAME:
            dev.close()
            continue
        if spec:
            ok = (os.path.realpath(spec) == path or spec.lower() in dev.name.lower())
        else:
            ok = is_touchscreen(dev)
        if ok:
            return dev
        dev.close()
    if denied:
        raise LookupError(
            f"touchscreen {', '.join(denied)} exists but is not readable by this user.\n"
            "Install the udev rule shipped with trackscreen (70-trackscreen.rules), "
            "or add yourself to the 'input' group, then log out and in again.")
    raise LookupError("no multitouch touchscreen found" + (f" matching {spec!r}" if spec else ""))


# --------------------------------------------------------------------------
# Geometry: touchscreen coordinates -> virtual touchpad coordinates
# --------------------------------------------------------------------------

class Geometry:
    def __init__(self, dev, size_mm=None, rotate=0, matrix=None, speed=1.0):
        self.ax = dev.absinfo(e.ABS_MT_POSITION_X)
        self.ay = dev.absinfo(e.ABS_MT_POSITION_Y)
        rx = max(1, self.ax.max - self.ax.min)
        ry = max(1, self.ay.max - self.ay.min)
        props = udev_props(dev.path)

        # Physical size of the panel in its native orientation.
        if size_mm:
            w_mm, h_mm, src = size_mm[0], size_mm[1], "--size"
        elif self.ax.resolution > 0 and self.ay.resolution > 0:
            w_mm, h_mm, src = rx / self.ax.resolution, ry / self.ay.resolution, "kernel"
        elif props.get("ID_INPUT_WIDTH_MM") and props.get("ID_INPUT_HEIGHT_MM"):
            w_mm, h_mm, src = float(props["ID_INPUT_WIDTH_MM"]), float(props["ID_INPUT_HEIGHT_MM"]), "udev"
        else:
            w_mm, h_mm, src = 300.0, 300.0 * ry / rx, "guess"
            log("could not determine the touchscreen's physical size; guessing. "
                "Use --size WxH (in mm) for correct pointer speed.")
        self.mm_per_unit_x = w_mm / rx
        self.mm_per_unit_y = h_mm / ry

        # Optional libinput calibration matrix (a b c d e f), as used for
        # rotated panels; taken from udev unless given explicitly.
        if matrix is None and props.get("LIBINPUT_CALIBRATION_MATRIX"):
            try:
                matrix = [float(v) for v in props["LIBINPUT_CALIBRATION_MATRIX"].split()]
            except ValueError:
                matrix = None
        if matrix is not None and len(matrix) != 6:
            raise SystemExit("calibration matrix needs 6 numbers")
        self.m = matrix or [1, 0, 0, 0, 1, 0]
        self.rot = rotate % 360
        if self.rot not in (0, 90, 180, 270):
            raise SystemExit("--rotate must be 0, 90, 180 or 270")

        swap = abs(self.m[0]) < 0.5 and abs(self.m[4]) < 0.5  # matrix rotates by 90/270
        if self.rot in (90, 270):
            swap = not swap
        self.w_mm, self.h_mm = (h_mm, w_mm) if swap else (w_mm, h_mm)
        self.max_x = max(1, round(self.w_mm * UNITS_PER_MM))
        self.max_y = max(1, round(self.h_mm * UNITS_PER_MM))
        # libinput derives physical distances from the resolution; lying about
        # it scales pointer/gesture speed.
        self.res = max(1, round(UNITS_PER_MM / speed))
        log(f"panel {w_mm:.0f}x{h_mm:.0f} mm ({src}), touchpad {self.w_mm:.0f}x{self.h_mm:.0f} mm, "
            f"rotate {self.rot}, speed x{speed:g}")

    def map(self, x, y):
        u = (x - self.ax.min) / (self.ax.max - self.ax.min or 1)
        v = (y - self.ay.min) / (self.ay.max - self.ay.min or 1)
        m = self.m
        u, v = m[0] * u + m[1] * v + m[2], m[3] * u + m[4] * v + m[5]
        if self.rot == 90:
            u, v = 1 - v, u
        elif self.rot == 180:
            u, v = 1 - u, 1 - v
        elif self.rot == 270:
            u, v = v, 1 - u
        u = min(1.0, max(0.0, u))
        v = min(1.0, max(0.0, v))
        return round(u * self.max_x), round(v * self.max_y)

    def dist_mm(self, x0, y0, x1, y1):
        return math.hypot((x1 - x0) * self.mm_per_unit_x, (y1 - y0) * self.mm_per_unit_y)


def make_virtual_touchpad(geo, nslots):
    ax = AbsInfo(value=0, min=0, max=geo.max_x, fuzz=0, flat=0, resolution=geo.res)
    ay = AbsInfo(value=0, min=0, max=geo.max_y, fuzz=0, flat=0, resolution=geo.res)
    caps = {
        # No BTN_LEFT on purpose: libinput then enables tap-to-click by default.
        e.EV_KEY: [e.BTN_TOUCH] + TOOL_BUTTONS,
        e.EV_ABS: [
            (e.ABS_X, ax), (e.ABS_Y, ay),
            (e.ABS_MT_SLOT, AbsInfo(0, 0, nslots - 1, 0, 0, 0)),
            (e.ABS_MT_POSITION_X, ax), (e.ABS_MT_POSITION_Y, ay),
            (e.ABS_MT_TRACKING_ID, AbsInfo(0, 0, 65535, 0, 0, 0)),
        ],
    }
    return UInput(caps, name=VIRTUAL_NAME, vendor=0x7473, product=0x0001, version=1,
                  bustype=e.BUS_VIRTUAL, input_props=[e.INPUT_PROP_POINTER])


# --------------------------------------------------------------------------
# Multitouch state
# --------------------------------------------------------------------------

def _eviocgmtslots(size):
    return (2 << 30) | (size << 16) | (ord("E") << 8) | 0x0A  # _IOC(_IOC_READ,'E',0x0a,size)


def read_mt_slots(fd, code, nslots):
    buf = array.array("i", [code] + [0] * nslots)
    fcntl.ioctl(fd, _eviocgmtslots(buf.itemsize * len(buf)), buf, True)
    return list(buf[1:])


class Slot:
    __slots__ = ("tid", "x", "y")

    def __init__(self):
        self.tid, self.x, self.y = -1, 0, 0


class TapDetector:
    """Recognises an N-finger tap: N+ fingers, short, nearly motionless."""

    def __init__(self, geo, fingers, max_time, max_move_mm):
        self.geo, self.fingers, self.max_time, self.max_move = geo, fingers, max_time, max_move_mm
        self.reset()

    def reset(self):
        self.start, self.max_n, self.moved, self.origin = None, 0, False, {}

    def update(self, now, slots):
        n = 0
        for s in slots:
            if s.tid < 0:
                continue
            n += 1
            o = self.origin.setdefault(s.tid, (s.x, s.y))
            if not self.moved and self.geo.dist_mm(o[0], o[1], s.x, s.y) > self.max_move:
                self.moved = True
        if n and self.start is None:
            self.start = now
        self.max_n = max(self.max_n, n)

    def finish(self, now):
        ok = (self.start is not None and self.max_n >= self.fingers and not self.moved
              and now - self.start <= self.max_time)
        dbg(f"touch sequence ended: fingers={self.max_n} moved={self.moved} "
            f"time={now - (self.start or now):.2f}s -> {'TOGGLE' if ok else 'no toggle'}")
        self.reset()
        return ok


# --------------------------------------------------------------------------
# Main engine
# --------------------------------------------------------------------------

class Trackscreen:
    def __init__(self, args):
        self.args = args
        self.mode = args.start_mode
        self.pending_toggle = False
        self.dev = self.ui = None

    # -- device lifecycle ---------------------------------------------------

    def open(self):
        a = self.args
        self.dev = find_touchscreen(a.device)
        self.nslots = self.dev.absinfo(e.ABS_MT_SLOT).max + 1
        self.geo = Geometry(self.dev, a.size, a.rotate, a.matrix, a.speed)
        self.ui = make_virtual_touchpad(self.geo, self.nslots)
        self.slots = [Slot() for _ in range(self.nslots)]
        self.vtid = [-1] * self.nslots       # tracking ids as last sent to the virtual pad
        self.vtool = -1
        self.dirty = set()
        self.dropping = False
        fingers = min(a.toggle_fingers, self.nslots)
        if fingers != a.toggle_fingers:
            log(f"touchscreen only tracks {self.nslots} fingers; toggle tap uses {fingers}")
        if a.toggle_fingers and fingers < 4:
            log("warning: toggle tap with <4 fingers will also produce clicks/gestures")
        self.tap = TapDetector(self.geo, fingers, a.tap_time, a.tap_move) if a.toggle_fingers else None
        self.resync()
        log(f"using {self.dev.path}: {self.dev.name!r} ({self.nslots} slots), "
            f"virtual touchpad {self.ui.device.path if self.ui.device else VIRTUAL_NAME!r}")
        self.apply_mode(notify=False)

    def close(self):
        if self.ui:
            try:
                self.lift_all_virtual()
            except OSError:
                pass
            self.ui.close()
            self.ui = None
        if self.dev:
            try:
                self.dev.ungrab()
            except OSError:
                pass
            self.dev.close()
            self.dev = None

    def resync(self):
        """Re-read the full multitouch state from the kernel (start-up, SYN_DROPPED)."""
        fd = self.dev.fd
        tids = read_mt_slots(fd, e.ABS_MT_TRACKING_ID, self.nslots)
        xs = read_mt_slots(fd, e.ABS_MT_POSITION_X, self.nslots)
        ys = read_mt_slots(fd, e.ABS_MT_POSITION_Y, self.nslots)
        for i, s in enumerate(self.slots):
            if (s.tid, s.x, s.y) != (tids[i], xs[i], ys[i]):
                s.tid, s.x, s.y = tids[i], xs[i], ys[i]
                self.dirty.add(i)
        self.cur = min(max(self.dev.absinfo(e.ABS_MT_SLOT).value, 0), self.nslots - 1)

    # -- modes --------------------------------------------------------------

    def apply_mode(self, notify=True):
        if self.mode == TRACKPAD:
            try:
                self.dev.grab()
            except OSError as ex:
                log(f"cannot grab touchscreen ({ex}); staying in touchscreen mode")
                self.mode = TOUCHSCREEN
        else:
            try:
                self.dev.ungrab()
            except OSError:
                pass  # was not grabbed
        log(f"mode: {self.mode}")
        if notify:
            self.notify()

    def toggle(self):
        self.mode = TRACKPAD if self.mode == TOUCHSCREEN else TOUCHSCREEN
        self.apply_mode()

    def notify(self):
        if self.args.no_notify or not shutil.which("notify-send"):
            return
        trackpad = self.mode == TRACKPAD
        try:
            subprocess.Popen(
                ["notify-send", "-a", "Trackscreen", "-t", "1500", "-e",
                 "-i", "input-touchpad" if trackpad else "input-touchscreen",
                 "-h", "string:x-canonical-private-synchronous:trackscreen",
                 "Trackpad mode" if trackpad else "Touchscreen mode"],
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
        except OSError:
            pass

    # -- event handling -----------------------------------------------------

    def handle(self, ev):
        if ev.type == e.EV_SYN:
            if ev.code == e.SYN_DROPPED:
                self.dropping = True
            elif ev.code == e.SYN_REPORT:
                if self.dropping:
                    self.dropping = False
                    self.resync()
                self.frame()
            return
        if self.dropping or ev.type != e.EV_ABS:
            return
        if ev.code == e.ABS_MT_SLOT:
            self.cur = ev.value if 0 <= ev.value < self.nslots else self.cur
            return
        s = self.slots[self.cur]
        if ev.code == e.ABS_MT_TRACKING_ID:
            s.tid = ev.value
        elif ev.code == e.ABS_MT_POSITION_X:
            s.x = ev.value
        elif ev.code == e.ABS_MT_POSITION_Y:
            s.y = ev.value
        else:
            return
        self.dirty.add(self.cur)

    def frame(self):
        now = time.monotonic()
        if self.tap:
            self.tap.update(now, self.slots)
        if self.mode == TRACKPAD:
            self.forward()
        self.dirty.clear()
        if not any(s.tid >= 0 for s in self.slots):  # all fingers lifted
            tapped = self.tap.finish(now) if self.tap else False
            if tapped != self.pending_toggle:  # both cancel each other out
                self.toggle()
            self.pending_toggle = False

    def forward(self):
        """Re-emit the changed touch slots on the virtual touchpad."""
        ui, geo = self.ui, self.geo
        for i in sorted(self.dirty):
            s = self.slots[i]
            ui.write(e.EV_ABS, e.ABS_MT_SLOT, i)
            tid = s.tid & 0xFFFF if s.tid >= 0 else -1
            if tid != self.vtid[i]:
                ui.write(e.EV_ABS, e.ABS_MT_TRACKING_ID, tid)
                self.vtid[i] = tid
            if tid >= 0:
                x, y = geo.map(s.x, s.y)
                ui.write(e.EV_ABS, e.ABS_MT_POSITION_X, x)
                ui.write(e.EV_ABS, e.ABS_MT_POSITION_Y, y)
        active = [s for s in self.slots if s.tid >= 0]
        self.emit_common(active)

    def emit_common(self, active):
        ui, n = self.ui, len(active)
        tool = min(n, len(TOOL_BUTTONS)) - 1  # -1: no finger
        if tool != self.vtool:
            if self.vtool >= 0:
                ui.write(e.EV_KEY, TOOL_BUTTONS[self.vtool], 0)
            if tool >= 0:
                ui.write(e.EV_KEY, TOOL_BUTTONS[tool], 1)
            ui.write(e.EV_KEY, e.BTN_TOUCH, 1 if n else 0)
            self.vtool = tool
        if active:
            x, y = self.geo.map(active[0].x, active[0].y)
            ui.write(e.EV_ABS, e.ABS_X, x)
            ui.write(e.EV_ABS, e.ABS_Y, y)
        ui.syn()

    def lift_all_virtual(self):
        for i, tid in enumerate(self.vtid):
            if tid >= 0:
                self.ui.write(e.EV_ABS, e.ABS_MT_SLOT, i)
                self.ui.write(e.EV_ABS, e.ABS_MT_TRACKING_ID, -1)
                self.vtid[i] = -1
        self.emit_common([])

    # -- signals ------------------------------------------------------------

    def request_toggle(self):
        if any(s.tid >= 0 for s in self.slots):
            self.pending_toggle = not self.pending_toggle  # applied once fingers lift
        else:
            self.toggle()

    # -- main loop ----------------------------------------------------------

    def run(self):
        rsock, wsock = socket.socketpair()
        rsock.setblocking(False)
        wsock.setblocking(False)
        signal.set_wakeup_fd(wsock.fileno())
        flags = {"toggle": 0, "stop": False}

        def on_toggle(*_):
            flags["toggle"] += 1

        def on_stop(*_):
            flags["stop"] = True

        signal.signal(signal.SIGUSR1, on_toggle)
        signal.signal(signal.SIGTERM, on_stop)
        signal.signal(signal.SIGINT, on_stop)
        signal.signal(signal.SIGHUP, on_stop)

        waiting_logged = False
        try:
            while not flags["stop"]:
                if self.dev is None:
                    try:
                        self.open()
                        waiting_logged = False
                    except LookupError as ex:
                        if not self.args.wait:
                            raise SystemExit(f"trackscreen: {ex}")
                        if not waiting_logged:
                            log(f"{ex}; waiting for it to appear...")
                            waiting_logged = True
                        select.select([rsock], [], [], 2.0)
                        self._drain(rsock)
                        continue

                r, _, _ = select.select([self.dev.fd, rsock], [], [])
                if rsock in r:
                    self._drain(rsock)
                while flags["toggle"]:
                    flags["toggle"] -= 1
                    self.request_toggle()
                if self.dev.fd in r:
                    try:
                        for ev in self.dev.read():
                            self.handle(ev)
                    except BlockingIOError:
                        pass
                    except OSError as ex:  # device vanished (suspend, driver reload, ...)
                        log(f"lost touchscreen ({ex}); will reconnect")
                        self.close()
                        if not self.args.wait:
                            self.args.wait = True
        finally:
            self.close()

    @staticmethod
    def _drain(sock):
        try:
            while sock.recv(4096):
                pass
        except BlockingIOError:
            pass


# --------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------

def pidfile_path():
    return os.path.join(os.environ.get("XDG_RUNTIME_DIR") or f"/tmp/trackscreen-{os.getuid()}",
                        "trackscreen.pid")


def running_pid():
    try:
        with open(pidfile_path()) as f:
            pid = int(f.read().strip())
        with open(f"/proc/{pid}/cmdline", "rb") as f:
            if b"trackscreen" not in f.read():
                return None
        return pid
    except (OSError, ValueError):
        return None


def list_devices():
    for path, dev, props in candidate_devices():
        ts = props.get("ID_INPUT_TOUCHSCREEN") == "1"
        if dev is None:
            if ts:
                print(f"{path}  (no permission)  udev says: touchscreen")
            continue
        if is_touchscreen(dev) or ts:
            print(f"{path}  {dev.name!r}  {'<- auto-detected' if is_touchscreen(dev) else ''}")
        dev.close()


def parse_size(s):
    try:
        w, h = s.lower().split("x")
        return float(w), float(h)
    except ValueError:
        raise argparse.ArgumentTypeError("expected WxH in millimetres, e.g. 307x187")


def main():
    p = argparse.ArgumentParser(
        description="Use the touchscreen as a multitouch touchpad; toggle with a five-finger tap.")
    p.add_argument("--toggle", action="store_true",
                   help="toggle the mode of the running instance and exit")
    p.add_argument("--list-devices", action="store_true", help="list touchscreens and exit")
    p.add_argument("-d", "--device", help="touchscreen device path or (part of) its name "
                   "(default: auto-detect)")
    p.add_argument("--start-mode", choices=[TOUCHSCREEN, TRACKPAD], default=TOUCHSCREEN)
    p.add_argument("--toggle-fingers", type=int, default=5,
                   help="number of fingers for the toggle tap; 0 disables it (default 5)")
    p.add_argument("--tap-time", type=float, default=0.6,
                   help="max duration of the toggle tap in seconds (default 0.6)")
    p.add_argument("--tap-move", type=float, default=12.0,
                   help="max finger movement during the toggle tap in mm (default 12)")
    p.add_argument("--speed", type=float, default=1.0,
                   help="pointer/gesture speed multiplier on top of the desktop's touchpad "
                   "settings (default 1.0)")
    p.add_argument("--rotate", type=int, default=0,
                   help="rotate touch coordinates clockwise by 0/90/180/270 degrees, e.g. if "
                   "your built-in display is rotated")
    p.add_argument("--matrix", type=float, nargs=6, metavar="N",
                   help="calibration matrix 'a b c d e f' (default: LIBINPUT_CALIBRATION_MATRIX "
                   "from udev, if set)")
    p.add_argument("--size", type=parse_size, metavar="WxH",
                   help="physical touchscreen size in mm (default: from kernel/udev)")
    p.add_argument("--no-wait", dest="wait", action="store_false",
                   help="exit instead of waiting if no touchscreen is available")
    p.add_argument("--no-notify", action="store_true", help="no desktop notification on toggle")
    p.add_argument("--debug", action="store_true")
    args = p.parse_args()

    global DEBUG
    DEBUG = args.debug

    if args.list_devices:
        list_devices()
        return
    if args.toggle:
        pid = running_pid()
        if not pid:
            sys.exit("trackscreen is not running")
        os.kill(pid, signal.SIGUSR1)
        return
    if args.speed <= 0:
        p.error("--speed must be positive")

    pid = running_pid()
    if pid and pid != os.getpid():
        sys.exit(f"trackscreen is already running (pid {pid}); use --toggle")
    pf = pidfile_path()
    os.makedirs(os.path.dirname(pf), exist_ok=True)
    with open(pf, "w") as f:
        f.write(str(os.getpid()))
    try:
        Trackscreen(args).run()
    finally:
        if running_pid() == os.getpid():
            os.unlink(pf)


if __name__ == "__main__":
    main()