#!/usr/bin/env python3
"""Kitbly Remote help – official desktop helper (session-only).

Lets a trusted viewer see this computer and, after you click Allow, move the
mouse and type. Nothing is recorded. Kitbly staff will never ask you to run
this program. Press Ctrl+C or q to stop instantly.

Usage:
  python3 kitbly-remote-helper.py
  python3 kitbly-remote-helper.py --allow --code-out /tmp/code --display :99   # unattended / tests

Optional packages (pip install):
  websockets aiortc av numpy mss pynput

Without aiortc the helper still works: JPEG frames go over a WebRTC data
channel (via aiortc if present, otherwise a documented error asking to pip
install). Mouse/keyboard injection uses pynput, or xdotool if pynput is missing.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import ssl
import sys
import time
import urllib.parse

HOST_DEFAULT = "kitbly.com"

# --------------------------------------------------------------------------- CLI
def parse_args():
    p = argparse.ArgumentParser(description="Kitbly Remote help helper")
    p.add_argument("--host", default=os.environ.get("KITBLY_REMOTE_HOST", HOST_DEFAULT))
    p.add_argument("--name", default=os.environ.get("KITBLY_REMOTE_NAME", "My computer"))
    p.add_argument("--allow", action="store_true", help="Auto-accept the next viewer (tests / unattended)")
    p.add_argument("--code-out", default="", help="Write the 6-digit code to this file")
    p.add_argument("--display", default=os.environ.get("DISPLAY", ""), help="X11 display, e.g. :99")
    p.add_argument("--fps", type=int, default=8)
    p.add_argument("--max-width", type=int, default=1280)
    p.add_argument("--insecure", action="store_true", help="Skip TLS verify (local tests only)")
    return p.parse_args()


def log(msg):
    print(msg, flush=True)


# --------------------------------------------------------------------------- screen
def grabber(display, max_width):
    if display:
        os.environ["DISPLAY"] = display
    try:
        import mss  # type: ignore
        sct = mss.mss()
        mon = sct.monitors[1] if len(sct.monitors) > 1 else sct.monitors[0]

        def grab():
            shot = sct.grab(mon)
            return shot.width, shot.height, bytes(shot.bgra), "bgra"
        return grab, mon["width"], mon["height"]
    except Exception:
        pass
    # fallback: ImageMagick / ffmpeg import
    def grab():
        import subprocess, tempfile
        from pathlib import Path
        fd, path = tempfile.mkstemp(suffix=".png")
        os.close(fd)
        try:
            env = os.environ.copy()
            if display:
                env["DISPLAY"] = display
            r = subprocess.run(["import", "-window", "root", path], env=env, capture_output=True, timeout=5)
            if r.returncode != 0:
                subprocess.run(["ffmpeg", "-y", "-f", "x11grab", "-video_size", "800x600",
                                "-i", display or ":0", "-frames", "1", path], env=env, capture_output=True, timeout=8)
            from struct import unpack
            # decode via Pillow if present else raw ffmpeg ppm
            try:
                from PIL import Image
                im = Image.open(path).convert("RGB")
                return im.size[0], im.size[1], im.tobytes(), "rgb"
            except Exception:
                return 0, 0, b"", "rgb"
        finally:
            try:
                os.remove(path)
            except Exception:
                pass
    return grab, 0, 0


def scale_wh(w, h, max_width):
    if not w or w <= max_width:
        return w, h
    return max_width, max(1, int(h * max_width / w))


def to_jpeg(w, h, data, fmt, max_width, quality=55):
    tw, th = scale_wh(w, h, max_width)
    try:
        from PIL import Image
        mode = "RGB"
        raw = data
        if fmt == "bgra":
            # BGRA -> RGB
            b = bytearray(w * h * 3)
            for i in range(w * h):
                o = i * 4
                d = i * 3
                b[d] = data[o + 2]; b[d + 1] = data[o + 1]; b[d + 2] = data[o]
            raw, mode = bytes(b), "RGB"
        im = Image.frombytes(mode, (w, h), raw)
        if (tw, th) != (w, h):
            im = im.resize((tw, th), Image.BILINEAR)
        import io
        buf = io.BytesIO()
        im.save(buf, format="JPEG", quality=quality, optimize=False)
        return buf.getvalue(), tw, th
    except Exception:
        return None, w, h


# --------------------------------------------------------------------------- input
class Injector:
    def __init__(self, display=""):
        if display:
            os.environ["DISPLAY"] = display
        self.w = 1
        self.h = 1
        self.ok = False
        self.backend = "none"
        try:
            from pynput.mouse import Button, Controller as MC  # type: ignore
            from pynput.keyboard import Key, Controller as KC  # type: ignore
            self.mc, self.kc, self.Button, self.Key = MC(), KC(), Button, Key
            self.backend = "pynput"
            self.ok = True
        except Exception:
            self.backend = "xdotool" if _which("xdotool") else "none"
            self.ok = self.backend == "xdotool"

    def size(self, w, h):
        self.w, self.h = max(1, w), max(1, h)

    def _xy(self, x, y):
        return int(float(x) * self.w), int(float(y) * self.h)

    def mouse(self, a, x, y, b=0, dy=0):
        X, Y = self._xy(x, y)
        if self.backend == "pynput":
            self.mc.position = (X, Y)
            btn = self.Button.right if b == 2 else self.Button.middle if b == 1 else self.Button.left
            if a == "down":
                self.mc.press(btn)
            elif a == "up":
                self.mc.release(btn)
            elif a == "wheel":
                self.mc.scroll(0, -1 if dy > 0 else 1)
        elif self.backend == "xdotool":
            import subprocess
            env = os.environ.copy()
            cmd = ["xdotool", "mousemove", "--sync", str(X), str(Y)]
            if a == "down":
                cmd = ["xdotool", "mousemove", "--sync", str(X), str(Y), "mousedown", str((b or 0) + 1)]
            elif a == "up":
                cmd = ["xdotool", "mousemove", "--sync", str(X), str(Y), "mouseup", str((b or 0) + 1)]
            elif a == "wheel":
                cmd = ["xdotool", "mousemove", "--sync", str(X), str(Y), "click", "5" if dy > 0 else "4"]
            subprocess.Popen(cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    def key(self, a, key, code="", mods=None):
        if a != "down":
            # pynput: also release
            if self.backend == "pynput":
                k = self._pkey(key, code)
                if k is not None:
                    self.kc.release(k)
            return
        if self.backend == "pynput":
            k = self._pkey(key, code)
            if k is not None:
                self.kc.press(k)
                if key and len(key) == 1:
                    self.kc.release(k)
        elif self.backend == "xdotool":
            import subprocess
            name = key if key and len(key) == 1 else (code or key or "")
            if name:
                subprocess.Popen(["xdotool", "key", name], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    def _pkey(self, key, code):
        special = {
            "Enter": self.Key.enter, "Tab": self.Key.tab, "Backspace": self.Key.backspace,
            "Escape": self.Key.esc, "ArrowLeft": self.Key.left, "ArrowRight": self.Key.right,
            "ArrowUp": self.Key.up, "ArrowDown": self.Key.down, "Delete": self.Key.delete,
            "Home": self.Key.home, "End": self.Key.end, "PageUp": self.Key.page_up,
            "PageDown": self.Key.page_down, "Shift": self.Key.shift, "Control": self.Key.ctrl,
            "Alt": self.Key.alt, "Meta": self.Key.cmd, " ": self.Key.space,
        }
        if key in special:
            return special[key]
        if key and len(key) == 1:
            return key
        return None


def _which(name):
    from shutil import which
    return which(name)


# --------------------------------------------------------------------------- signaling + webrtc (aiortc)
async def run(args):
    if args.display:
        os.environ["DISPLAY"] = args.display
    try:
        import websockets
    except ImportError:
        log("Please install dependencies:\n  pip install websockets aiortc av numpy mss pynput pillow")
        return 2
    try:
        from aiortc import RTCPeerConnection, RTCSessionDescription, RTCIceCandidate, RTCConfiguration, RTCIceServer
        from aiortc import RTCDataChannel
        from aiortc.contrib.media import MediaPlayer
        have_rtc = True
    except ImportError:
        have_rtc = False
        log("aiortc is required. pip install aiortc av numpy")
        return 2

    grab, sw, sh = grabber(args.display, args.max_width)
    inj = Injector(args.display)
    # probe one frame for size
    try:
        w, h, _, _ = grab()
        if w:
            inj.size(w, h)
            sw, sh = w, h
    except Exception as e:
        log("screen grab failed: %s" % e)

    scheme = "wss"
    ws_url = f"{scheme}://{args.host}/api/remote/ws"
    sslopt = None
    if args.insecure:
        sslopt = ssl.create_default_context()
        sslopt.check_hostname = False
        sslopt.verify_mode = ssl.CERT_NONE

    log("Kitbly Remote help helper")
    log("Only run this if YOU started it. Kitbly staff will never ask you to.")
    log("Connecting to %s …" % args.host)
    async with websockets.connect(ws_url, ssl=sslopt, max_size=2**20, ping_interval=20) as ws:
        await ws.send(json.dumps({"t": "hello", "name": args.name, "role": "helper"}))
        ice_servers = []
        code = ""
        pc = None
        channel = None
        allowed = False
        stop = asyncio.Event()

        async def send(m):
            await ws.send(json.dumps(m))

        def ice_cfg(lst):
            servers = []
            for s in lst or []:
                urls = s.get("urls") if isinstance(s, dict) else None
                if not urls:
                    continue
                if isinstance(urls, str):
                    urls = [urls]
                servers.append(RTCIceServer(urls=urls, username=s.get("username"), credential=s.get("credential")))
            return RTCConfiguration(iceServers=servers or [RTCIceServer(urls=["stun:stun.l.google.com:19302"])])

        async def attach_track(peer):
            """Prefer a real video track; also start a JPEG pusher on the data channel as backup."""
            try:
                from av import VideoFrame
                import fractions
                from aiortc import VideoStreamTrack

                class Screen(VideoStreamTrack):
                    kind = "video"

                    def __init__(self):
                        super().__init__()
                        self._last = 0

                    async def recv(self):
                        pts, tb = await self.next_timestamp()
                        delay = 1 / max(3, args.fps)
                        now = time.monotonic()
                        if now - self._last < delay:
                            await asyncio.sleep(delay - (now - self._last))
                        self._last = time.monotonic()
                        try:
                            w, h, data, fmt = grab()
                            inj.size(w, h)
                            # convert bgra/rgb to ndarray rgb
                            import numpy as np
                            if fmt == "bgra":
                                arr = np.frombuffer(data, dtype=np.uint8).reshape((h, w, 4))
                                rgb = arr[:, :, [2, 1, 0]].copy()
                            else:
                                rgb = np.frombuffer(data, dtype=np.uint8).reshape((h, w, 3)).copy()
                            tw, th = scale_wh(w, h, args.max_width)
                            if (tw, th) != (w, h):
                                # cheap box scale
                                ys = (np.linspace(0, h - 1, th)).astype(int)
                                xs = (np.linspace(0, w - 1, tw)).astype(int)
                                rgb = rgb[ys][:, xs]
                            frame = VideoFrame.from_ndarray(rgb, format="rgb24")
                            frame.pts = pts
                            frame.time_base = tb
                            return frame
                        except Exception:
                            frame = VideoFrame.from_ndarray(
                                __import__("numpy").zeros((180, 320, 3), dtype="uint8"), format="rgb24")
                            frame.pts = pts
                            frame.time_base = tb
                            return frame

                peer.addTrack(Screen())
            except Exception as e:
                log("video track not available (%s); will use JPEG frames" % e)

        async def jpeg_loop():
            while not stop.is_set():
                await asyncio.sleep(1 / max(3, args.fps))
                ch = channel
                if not ch or ch.readyState != "open" or not allowed:
                    continue
                try:
                    w, h, data, fmt = grab()
                    if not w:
                        continue
                    inj.size(w, h)
                    jpg, _, _ = to_jpeg(w, h, data, fmt, args.max_width)
                    if jpg:
                        ch.send(bytes([1]) + jpg)
                except Exception:
                    pass

        async def start_pc(servers, as_offerer=True):
            nonlocal pc, channel
            if pc:
                await pc.close()
            pc = RTCPeerConnection(ice_cfg(servers))

            @pc.on("icecandidate")
            async def on_ice(c):
                if c:
                    await send({"t": "sig", "d": {"ice": {
                        "candidate": c.candidate, "sdpMid": c.sdpMid,
                        "sdpMLineIndex": c.sdpMLineIndex
                    }}})

            @pc.on("datachannel")
            def on_dc(ch):
                nonlocal channel
                channel = ch
                hook_dc(ch)

            await attach_track(pc)
            channel = pc.createDataChannel("help", ordered=True)
            hook_dc(channel)
            if as_offerer:
                offer = await pc.createOffer()
                await pc.setLocalDescription(offer)
                await send({"t": "sig", "d": {"sdp": {"type": pc.localDescription.type, "sdp": pc.localDescription.sdp}}})
            asyncio.create_task(jpeg_loop())

        def hook_dc(ch):
            @ch.on("message")
            def on_message(data):
                if not allowed:
                    return
                if isinstance(data, (bytes, bytearray)):
                    return
                try:
                    m = json.loads(data)
                except Exception:
                    return
                if m.get("t") == "m":
                    inj.mouse(m.get("a"), m.get("x", 0), m.get("y", 0), m.get("b", 0), m.get("dy", 0))
                elif m.get("t") == "k":
                    inj.key(m.get("a"), m.get("key") or "", m.get("code") or "", m.get("mods"))

        log("BEING SHARED will show once someone is accepted. Press Ctrl+C to stop.")
        async for raw in ws:
            try:
                m = json.loads(raw)
            except ValueError:
                continue
            t = m.get("t")
            if t == "hi":
                await send({"t": "create", "mode": "control"})
            elif t == "room":
                code = m.get("code") or ""
                log("Your code: %s %s" % (code[:3], code[3:]))
                log("Link: %s" % m.get("url"))
                log("Waiting. A viewer must be Accepted before they can control this PC.")
                if args.code_out:
                    open(args.code_out, "w").write(code)
                ice_servers = m.get("ice") or ice_servers
            elif t == "ask":
                who, name = m.get("from"), m.get("name") or "Guest"
                log("VIEWER ASKS: %s  (id %s)" % (name, who))
                if args.allow:
                    log("Auto-allow (--allow).")
                    await send({"t": "accept", "who": who})
                else:
                    log("Type 'allow' + Enter to accept, or 'deny' to refuse.")
                    # non-blocking stdin
                    loop = asyncio.get_event_loop()
                    def read_yes():
                        try:
                            return sys.stdin.readline().strip().lower()
                        except Exception:
                            return "deny"
                    ans = await loop.run_in_executor(None, read_yes)
                    if ans in ("allow", "a", "y", "yes", "ok"):
                        await send({"t": "accept", "who": who})
                    else:
                        await send({"t": "decline", "who": who})
            elif t == "accepted":
                allowed = True
                ice_servers = m.get("ice") or ice_servers
                log("***** BEING SHARED / CONTROL ENABLED *****")
                log("Visible session. Press Ctrl+C or close this window to revoke instantly.")
                await start_pc(ice_servers, True)
            elif t == "sig":
                d = m.get("d") or {}
                if not pc:
                    continue
                if d.get("sdp"):
                    sdp = d["sdp"]
                    await pc.setRemoteDescription(RTCSessionDescription(sdp=sdp["sdp"] if isinstance(sdp, dict) else sdp, type=sdp.get("type", "answer") if isinstance(sdp, dict) else "answer"))
                elif d.get("ice"):
                    c = d["ice"]
                    try:
                        cand = c.get("candidate") if isinstance(c, dict) else None
                        if cand:
                            await pc.addIceCandidate(RTCIceCandidate(
                                sdpMid=c.get("sdpMid"), sdpMLineIndex=c.get("sdpMLineIndex"),
                                candidate=cand if isinstance(cand, str) else cand.get("candidate", "")
                            ))
                    except Exception:
                        pass
            elif t in ("ended", "expired"):
                log("Session ended (%s)." % m.get("why", t))
                allowed = False
                stop.set()
                if pc:
                    await pc.close()
                break
            elif t == "err":
                log("Error: %s" % m.get("e"))
                if m.get("e") in ("rate", "slow"):
                    break
        stop.set()
        if pc:
            await pc.close()
    return 0


def main():
    args = parse_args()
    try:
        return asyncio.run(run(args))
    except KeyboardInterrupt:
        log("Stopped. Control revoked.")
        return 0


if __name__ == "__main__":
    sys.exit(main() or 0)