#!/usr/bin/env python3
#
# SPDX-License-Identifier: GPL-3.0 and BSD-3-clause
#
# Copyright 2024 Johannes Schauer Marin Rodrigues <josch@mister-muffin.de>
#
# inspired by https://github.com/barneygale/pytest-vnc (GPL-3)
# Copyright 2022 - 2023, Barney Gale <barney.gale@gmail.com>
#
# image search by https://github.com/asweigart/pyscreeze (BSD-3-clause)
# Copyright Al Sweigart <al@inventwithpython.com>

import socket
from collections import namedtuple
from dataclasses import field
from PIL import Image
import numpy
import time
import sys
import re

Point = namedtuple("Point", "x y")
Rect = namedtuple("Rect", "x y width height")


# Knuth-Morris-Pratt search algorithm implementation (to be used by screen capture)
def _kmp(needle, haystack):
    # build table of shift amounts
    shifts = [1] * (len(needle) + 1)
    shift = 1
    for pos in range(len(needle)):
        while shift <= pos and needle[pos] != needle[pos - shift]:
            shift += shifts[pos - shift]
        shifts[pos + 1] = shift

    # do the actual search
    startPos = 0
    matchLen = 0
    for c in haystack:
        while matchLen == len(needle) or matchLen >= 0 and needle[matchLen] != c:
            startPos += shifts[matchLen]
            matchLen -= shifts[matchLen]
        matchLen += 1
        if matchLen == len(needle):
            yield startPos


def find_needle_in_haystack(needleImage, haystackImage):
    if isinstance(needleImage, str):
        with open(needleImage, "rb") as needleFileObj:
            needleImage = Image.open(needleFileObj)
            assert needleImage.mode == "RGB"
            needleImageData = tuple(needleImage.getdata())
    else:
        needleImageData = tuple(needleImage.getdata())

    if isinstance(haystackImage, str):
        with open(haystackImage, "rb") as haystackFileObj:
            haystackImage = Image.open(haystackFileObj)
            if haystackImage.mode == "RGBA":
                haystackImage = haystackImage.convert("RGB")
            assert haystackImage.mode == "RGB"
            haystackImageData = tuple(haystackImage.getdata())
    else:
        haystackImageData = tuple(haystackImage.getdata())

    # search all rows from the topmost row down to the last row that the needle
    # can possibly be found in
    for y in range(haystackImage.height - needleImage.height + 1):
        for matchx in _kmp(
            needleImageData[0 : needleImage.width],  # first row of needle
            haystackImageData[y * haystackImage.width : (y + 1) * haystackImage.width],
        ):
            # the first row of the needle was found in the current row of the
            # haystack, so check if the rest of the needle can be found at this
            # position
            foundMatch = True
            for searchy in range(1, needleImage.height):
                haystackStart = (searchy + y) * haystackImage.width + matchx
                needleStart = searchy * needleImage.width
                if (
                    needleImageData[needleStart : needleStart + needleImage.width]
                    != haystackImageData[
                        haystackStart : haystackStart + needleImage.width
                    ]
                ):
                    foundMatch = False
                    break
            if foundMatch:
                # Match found, report the x, y, width, height of where the matching region is in haystack.
                return Rect(matchx, y, needleImage.width, needleImage.height)
    return None


# parsing /usr/include/X11/keysymdef.h seems to be intended, see
# src/util/makekeys.c in src:libx11
keysymdef = []
keysymre = re.compile(r"^#define (?:XF86)?XK_(\w+)\s+0x(\w+)(?:\s+/\*\s+U\+(\w+))?")
for fname in ["/usr/include/X11/keysymdef.h", "/usr/include/X11/XF86keysym.h"]:
    with open(fname) as f:
        for line in f:
            m = re.match(keysymre, line)
            if m is None:
                continue
            name, sym, uni = m.groups()
            sym = int(sym, 16)
            uni = int(uni, 16) if uni else None
            keysymdef.append((name, sym, uni))

# Colour channel orders
pixel_formats = {
    "bgra": b"\x20\x18\x00\x01\x00\xff\x00\xff\x00\xff\x10\x08\x00\x00\x00\x00",
    "rgba": b"\x20\x18\x00\x01\x00\xff\x00\xff\x00\xff\x00\x08\x10\x00\x00\x00",
    "argb": b"\x20\x18\x01\x01\x00\xff\x00\xff\x00\xff\x10\x08\x00\x00\x00\x00",
    "abgr": b"\x20\x18\x01\x01\x00\xff\x00\xff\x00\xff\x00\x08\x10\x00\x00\x00",
}

encodings = {
    6,  # zlib
}

key_codes = {}
key_codes.update((name, code) for name, code, char in keysymdef)
key_codes.update((chr(char), code) for name, code, char in keysymdef if char)
key_codes["Del"] = key_codes["Delete"]
key_codes["Esc"] = key_codes["Escape"]
key_codes["Cmd"] = key_codes["Super_L"]
key_codes["Alt"] = key_codes["Alt_L"]
key_codes["Ctrl"] = key_codes["Control_L"]
key_codes["Super"] = key_codes["Super_L"]
key_codes["Shift"] = key_codes["Shift_L"]
key_codes["Backspace"] = key_codes["BackSpace"]
key_codes["Space"] = key_codes["space"]


def write_mouse(sock, buttons, position):
    sock.sendall(
        b"\x05"
        + buttons.to_bytes(1, "big")
        + position.x.to_bytes(2, "big")
        + position.y.to_bytes(2, "big")
    )
    time.sleep(0.05)


def slice_rect(rect, *channels):
    return (
        slice(rect.y, rect.y + rect.height),
        slice(rect.x, rect.x + rect.width),
    ) + channels


def capture(sock, rect: Rect):
    sock.sendall(
        b"\x03\x00"
        + rect.x.to_bytes(2, "big")
        + rect.y.to_bytes(2, "big")
        + rect.width.to_bytes(2, "big")
        + rect.height.to_bytes(2, "big")
    )
    pixels = numpy.zeros((rect.height, rect.width, 4), "B")
    while True:
        update_type = read_int(sock, 1)
        assert update_type == 0  # video
        read(sock, 1)  # padding
        for _ in range(read_int(sock, 2)):
            area_rect = Rect(
                read_int(sock, 2),
                read_int(sock, 2),
                read_int(sock, 2),
                read_int(sock, 2),
            )
            area_encoding = read_int(sock, 4)
            assert area_encoding == 0  # raw not compressed
            area = read(sock, area_rect.height * area_rect.width * 4)
            with open("screenshots/screenshot.raw", "wb") as f:
                f.write(area)
            area = numpy.ndarray((area_rect.height, area_rect.width, 4), "B", area)
            pixels[slice_rect(area_rect)] = area
            pixels[slice_rect(area_rect, 3)] = 255
        if pixels[slice_rect(rect, 3)].all():
            return pixels[slice_rect(rect)]


def read(sock: socket, length: int) -> bytes:
    data = b""
    while len(data) < length:
        data += sock.recv(length - len(data))
    return data


def read_int(sock: socket, length: int) -> int:
    return int.from_bytes(read(sock, length), "big")


def move(sock, point):
    write_mouse(sock, 0, point)
    time.sleep(1)


def click(sock, point):
    write_mouse(sock, 1, point)
    time.sleep(1)
    write_mouse(sock, 0, point)
    time.sleep(1)


def keypress(sock, key):
    data = key_codes[key].to_bytes(4, "big")
    sock.sendall(b"\x04\x01\x00\x00" + data)
    time.sleep(0.1)
    sock.sendall(b"\x04\x00\x00\x00" + data)
    time.sleep(1)


def get_center_match(sock, rect, needle):
    pixels = capture(sock, rect)
    haystack = Image.fromarray(pixels)
    haystack.save("screenshot.png", mode="RGB")
    rect = find_needle_in_haystack(needle, "screenshot.png")
    if rect is None:
        return None
    center = Point(int(rect.x + rect.width / 2), int(rect.y + rect.height / 2))
    print(f"found match for {needle} at {center}", file=sys.stderr)
    return center


def screenshot(sock, rect, filename):
    pixels = capture(sock, rect)
    image = Image.fromarray(pixels)
    image.save(filename, mode="RGB")


def wait_visible(sock, rect, needle, timeout):
    starttime = time.time()
    while (time.time() - starttime) < timeout:
        pixels = capture(sock, rect)
        haystack = Image.fromarray(pixels)
        haystack.save("screenshots/screenshot.png", mode="RGB")
        if find_needle_in_haystack(needle, "screenshots/screenshot.png") is not None:
            return True
    # timeout reached
    return False


class RecvNullExc(Exception):
    pass


def connect(host, port, timeout):
    starttime = time.time()
    while (time.time() - starttime) < timeout:
        try:
            sock = socket.create_connection(("127.0.0.1", 5910), timeout=5)
            print("connected", file=sys.stderr)
            intro = sock.recv(12)
            if len(intro) != 12:
                raise RecvNullExc
            print("received intro", file=sys.stderr)
        except ConnectionRefusedError:
            print("trying to connect...", file=sys.stderr)
            time.sleep(2)
            pass
        except RecvNullExc:
            print("cannot read, retrying...", file=sys.stderr)
            time.sleep(2)
            pass
        else:
            print("leaving loop", file=sys.stderr)
            break
    assert intro[:4] == b"RFB "
    sock.sendall(b"RFB 003.008\n")

    num_auth_types = read_int(sock, 1)
    assert num_auth_types == 1, num_auth_types
    auth_type = int.from_bytes(read(sock, num_auth_types))
    assert auth_type == 1, auth_type

    sock.sendall(b"\x01")

    auth_result = read_int(sock, 4)
    assert auth_result == 0

    sock.sendall(b"\x01")
    width, height = read_int(sock, 2), read_int(sock, 2)
    assert (width, height) == (1920, 1080)
    rect = Rect(0, 0, width, height)
    read(sock, 16)
    read(sock, read_int(sock, 4))
    sock.sendall(
        b"\x00\x00\x00\x00"
        + pixel_formats["rgba"]
        + b"\x02\x00"
        + len(encodings).to_bytes(2, "big")
        + b"".join(encoding.to_bytes(4, "big") for encoding in encodings)
    )

    return sock, rect


def main():
    sock = None
    rect = None
    mousepos = Point(0, 0)
    for line in sys.stdin:
        if line.startswith("#"):
            continue
        line = line.strip()
        if not line:
            continue
        print(f"handling line: {line}", file=sys.stderr)
        tokens = line.split()
        if len(tokens) < 1:
            raise Exception("no tokens")
        match tokens[0]:
            case "connect":
                if len(tokens) != 4:
                    raise Exception(f"missing argument {line}")
                sock, rect = connect(tokens[1], int(tokens[2]), int(tokens[3]))
            case "sleep":
                if len(tokens) != 2:
                    raise Exception(f"missing argument {line}")
                time.sleep(float(tokens[1]))
            case "capture":
                if len(tokens) != 2:
                    raise Exception(f"missing argument {line}")
                if sock is None or rect is None:
                    raise Exception("must connect first")
                screenshot(sock, rect, tokens[1])
            case "key":
                if sock is None or rect is None:
                    raise Exception("must connect first")
                if len(tokens) != 2:
                    raise Exception(f"missing argument {line}")
                keypress(sock, tokens[1])
            case "mousemove":
                if sock is None or rect is None:
                    raise Exception("must connect first")
                if len(tokens) == 3:
                    mousepos = Point(int(tokens[1]), int(tokens[2]))
                elif len(tokens) == 2:
                    mousepos = get_center_match(sock, rect, tokens[1])
                    if mousepos is None:
                        raise Exception(f"no match for {tokens[1]} found")
                else:
                    raise Exception(f"missing argument {line}")
                move(sock, mousepos)
            case "click":
                if sock is None or rect is None:
                    raise Exception("must connect first")
                click(sock, mousepos)
            case "wait_visible":
                if sock is None or rect is None:
                    raise Exception("must connect first")
                if len(tokens) != 3:
                    raise Exception(f"missing argument {line}")
                if not wait_visible(sock, rect, tokens[1], int(tokens[2])):
                    raise Exception(
                        f"no match for {tokens[1]} found within {tokens[2]} s"
                    )
            case _:
                raise Exception(f"unknown command: {tokens[0]}")

    if sock is not None:
        sock.close()


if __name__ == "__main__":
    main()
