#!/usr/bin/env python3
"""Tests for the Python side: relayctl's WebSocket codec, and that its key table
has not drifted from the one in the controller page.

    python3 tests/test_relayctl.py

No third-party dependencies, so it runs anywhere the tools do.
"""
from __future__ import annotations

import inspect
import pathlib
import re
import struct
import sys

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent / "tools"))

import relayctl  # noqa: E402

failures = 0
checks = 0


def check(condition, label):
    global checks, failures
    checks += 1
    if not condition:
        failures += 1
        print(f"  FAIL  {label}")


def check_eq(actual, expected, label):
    global checks, failures
    checks += 1
    if actual != expected:
        failures += 1
        print(f"  FAIL  {label}: got {actual!r}, want {expected!r}")


print("handshake")
# RFC 6455 1.3's worked example -- the same vector the C++ suite asserts, so the
# two implementations are pinned to the same truth.
check_eq(relayctl.expect_accept("dGhlIHNhbXBsZSBub25jZQ=="),
         "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=", "Sec-WebSocket-Accept of the RFC example")
check(relayctl.expect_accept("differentKey==") != relayctl.expect_accept("dGhlIQ=="),
      "different keys give different accepts")

print("frame encoding")
frame = relayctl.encode_frame(relayctl.OP_TEXT, b"A1")
check_eq(len(frame), 2 + 4 + 2, "short masked frame length")
check_eq(frame[0], 0x81, "FIN + text opcode")
check_eq(frame[1] & 0x80, 0x80, "mask bit set (clients must mask)")
check_eq(frame[1] & 0x7F, 2, "payload length in the 7-bit field")
mask = frame[2:6]
check_eq(bytes(b ^ mask[i & 3] for i, b in enumerate(frame[6:])), b"A1",
         "payload round-trips through the mask")

long_payload = b"x" * 300
frame = relayctl.encode_frame(relayctl.OP_TEXT, long_payload)
check_eq(frame[1] & 0x7F, 126, "126 signals the 16-bit length")
check_eq(struct.unpack(">H", frame[2:4])[0], 300, "16-bit length value")
mask = frame[4:8]
check_eq(relayctl.mask_payload(frame[8:], mask), long_payload,
         "mask_payload is its own inverse")
check_eq(len(relayctl.encode_frame(relayctl.OP_TEXT, b"")), 6,
         "empty frame still carries a mask key")

print("frame decoding")
receiver = relayctl.Receiver()


def server_frame(opcode, payload):
    """The unmasked frames the board sends."""
    out = bytearray([0x80 | opcode, len(payload)])
    out += payload
    return bytes(out)


status = server_frame(relayctl.OP_TEXT, b'S{"armed":1,"usb":1}')
# Deliver it one byte at a time: the decoder must not invent a frame early.
emitted = []
for i in range(len(status)):
    emitted += list(receiver.feed(status[i:i + 1]))
check_eq(len(emitted), 1, "exactly one frame from byte-at-a-time delivery")
check_eq(emitted[0][0], relayctl.OP_TEXT, "decoded opcode")
check_eq(emitted[0][1], b'S{"armed":1,"usb":1}',
         "decoded payload, including the board's 'S' status prefix")

# Two frames in one read, plus a split across the boundary.
two = server_frame(relayctl.OP_TEXT, b"OK") + server_frame(relayctl.OP_PING, b"hi")
receiver = relayctl.Receiver()
got = list(receiver.feed(two))
check_eq(len(got), 2, "both frames decoded from one read")
check_eq(got[0][1], b"OK", "first frame payload")
check_eq(got[1][0], relayctl.OP_PING, "second frame opcode")

receiver = relayctl.Receiver()
got = list(receiver.feed(two[:5])) + list(receiver.feed(two[5:]))
check_eq(len(got), 2, "a frame split mid-payload is still decoded once")

# 16-bit length path, and the protocol violations the client refuses.
receiver = relayctl.Receiver()
big = bytearray([0x81, 126]) + struct.pack(">H", 130) + b"y" * 130
check_eq(list(receiver.feed(bytes(big)))[0][1], b"y" * 130, "16-bit length decoded")

receiver = relayctl.Receiver()
try:
    list(receiver.feed(bytes([0x81, 0x82]) + b"\x00" * 6))
    check(False, "a masked server frame must be rejected")
except relayctl.RelayError:
    check(True, "masked server frame rejected")

receiver = relayctl.Receiver()
try:
    list(receiver.feed(bytes([0x01, 0x01, b"a"[0]])))
    check(False, "a fragmented frame must be rejected")
except relayctl.RelayError:
    check(True, "fragmented server frame rejected")

print("key table")
check_eq(relayctl.resolve_key("Enter"), 0x28, "Enter is usage 0x28")
check_eq(relayctl.resolve_key("0x28"), 0x28, "numeric usage accepted")
check_eq(relayctl.resolve_key("Delete"), 0x4C, "Delete is Delete-Forward 0x4C")
for bad, why in [("0x00", "usage 0 is not a key"), ("0x1ff", "usage must fit a byte"),
                 ("NotAKey", "unknown name")]:
    try:
        relayctl.resolve_key(bad)
        check(False, why)
    except (relayctl.RelayError, ValueError):
        check(True, why)

# The page and this tool must agree, or `relayctl key Enter` and pressing Enter in
# the browser would be different keys. Compare the JS object in web_page.h.
page = (pathlib.Path(__file__).resolve().parent.parent / "src" / "web_page.h").read_text()
match = re.search(r"const CODE_TO_USAGE = \{(.*?)\};", page, re.S)
check(match is not None, "CODE_TO_USAGE found in src/web_page.h")
if match:
    page_map = {}
    for name, usage in re.findall(r"([A-Za-z][A-Za-z0-9]*)\s*:\s*(0x[0-9a-fA-F]+)",
                                  match.group(1)):
        page_map[name] = int(usage, 16)
    check(len(page_map) > 80, f"page table parsed ({len(page_map)} entries)")
    only_page = sorted(set(page_map) - set(relayctl.CODE_TO_USAGE))
    only_cli = sorted(set(relayctl.CODE_TO_USAGE) - set(page_map))
    check_eq(only_page, [], "keys the page knows but relayctl does not")
    check_eq(only_cli, [], "keys relayctl knows but the page does not")
    mismatched = sorted(k for k in set(page_map) & set(relayctl.CODE_TO_USAGE)
                        if page_map[k] != relayctl.CODE_TO_USAGE[k])
    check_eq(mismatched, [], "same name, different usage id")
    check_eq(page_map.get("Enter"), 0x28, "Enter agrees on 0x28")
    check_eq(page_map.get("KeyA"), 0x04, "KeyA agrees on 0x04")

print("pasted text")
root = pathlib.Path(__file__).resolve().parent.parent
relay_src = (root / "tools" / "relayctl.py").read_text()
proto_src = (root / "src" / "proto.h").read_text()
text_src = (root / "src" / "text.h").read_text()

check_eq(relayctl.normalize_paste("a\r\nb"), "a\nb", "CRLF becomes one Enter")
check_eq(relayctl.normalize_paste("a\rb"), "a\nb", "a lone CR becomes Enter")
check_eq(relayctl.normalize_paste("a\nb"), "a\nb", "LF is left alone")
check_eq(relayctl.text_frame(0, "hi"), "P=0:6869", "the frame relayctl sends")
check_eq(relayctl.text_frame(250, "hi"), "P=250:6869", "the pace travels in decimal")
check_eq(relayctl.text_frame(0, "a"), "P=0:61", "one character")

# What the board's character map can type, read straight out of src/text.h: the
# switch spells out the punctuation and the ranges spell out the alphanumerics.
firmware = set()
for literal in re.findall(r"case '(\\.|[^'\\])':", text_src):
    firmware.add(ord(literal.encode("latin-1").decode("unicode_escape")))
check("ch >= 'a' && ch <= 'z'" in text_src, "the lowercase range is still a range")
check("ch >= 'A' && ch <= 'Z'" in text_src, "the uppercase range is still a range")
check("ch >= '0' && ch <= '9'" in text_src, "the digit range is still a range")
firmware |= set(range(ord("a"), ord("z") + 1))
firmware |= set(range(ord("A"), ord("Z") + 1))
firmware |= set(range(ord("0"), ord("9") + 1))

expected = {9, 10, 13} | set(range(32, 127))
check_eq(sorted(firmware), sorted(expected),
         "the board types exactly tab, newline, CR and printable ASCII")

# ... and the tool's validator must agree with it on every ASCII character, once the
# line endings are normalised. A character the tool waves through that the board then
# refuses would fail the whole paste at the board instead of at the terminal.
for code in range(128):
    ch = chr(code)
    tool_says_ok = not relayctl.untypeable(relayctl.normalize_paste(ch))
    if tool_says_ok != (code in firmware):
        check(False, f"code {code:#04x}: tool={tool_says_ok} board={code in firmware}")
check(True, "tool and board agree on all 128 ASCII codes")

for bad in ["\u2019", "\u2014", "\u00e9", "\u00a0", "\x00", "\x07", "\x7f"]:
    check_eq(relayctl.untypeable("ok" + bad), [bad], f"refuses {bad!r}")
check_eq(relayctl.untypeable("line1\nline2\tcol\r"), ["\r"],
         "CR is flagged before normalisation, and never sent")

# The three numbers have to be the board's numbers, or a paste dies on E8 halfway
# through because the tool split the text somewhere the firmware will not accept.
config_src = (root / "src" / "config.h").read_text()


def config_value(name):
    # Decimal or hex: kMouseDeltaMax is 127 and kMouseButtonMask is 0x1F, and both are
    # numbers a client has to agree with the firmware about.
    found = re.search(rf"\b{name}\s*=\s*(0[xX][0-9a-fA-F]+|\d+)\s*;", config_src)
    return int(found.group(1), 0) if found else None


check_eq(relayctl.TEXT_FRAME_MAX, config_value("kTextFrameMax"),
         "relayctl TEXT_FRAME_MAX == kTextFrameMax")
check_eq(relayctl.MAX_PASTE_CHARS, config_value("kTextJobCap"),
         "relayctl MAX_PASTE_CHARS == kTextJobCap")
check_eq(relayctl.DEFAULT_TEXT_PACE_MS, config_value("kTextPeriodMsDefault"),
         "relayctl DEFAULT_TEXT_PACE_MS == kTextPeriodMsDefault")
page_src = (root / "src" / "web_page.h").read_text()
check_eq(int(re.search(r"const TEXT_FRAME_MAX = (\d+);", page_src).group(1)),
         config_value("kTextFrameMax"), "page TEXT_FRAME_MAX == kTextFrameMax")
check_eq(int(re.search(r"const TEXT_JOB_MAX = (\d+);", page_src).group(1)),
         config_value("kTextJobCap"), "page TEXT_JOB_MAX == kTextJobCap")
check(config_value("kTextFrameMax") * 2 + 7 <= config_value("kWsMessageCap"),
      "the largest text frame fits one WebSocket message")

# _drain answered the board's ping with two methods that do not exist, which only
# surfaced on a session longer than the 15 s ping interval. A paste is such a session.
check("_send_frame" not in relay_src, "_drain no longer calls the missing _send_frame")
check("_check_error" not in relay_src, "_drain no longer calls the missing _check_error")
drained = relayctl.Relay._drain
src_lines = inspect.getsource(drained)
check("OP_PONG" in src_lines, "_drain answers a ping")
check("raise RelayError" in src_lines, "_drain still surfaces an error frame")

print("mouse")
check_eq(relayctl.mouse_frame(0x01, 30, -4), "M=01:30:-4:0:0",
         "a mouse frame is the button byte then four movements")
check_eq(relayctl.mouse_frame(0), "M=00:0:0:0:0",
         "and all five fields are always written -- the firmware parses exactly five")
check_eq(relayctl.mouse_frame(0x1f, -127, 127, -127, 127), "M=1f:-127:127:-127:127",
         "the full signed range")

for bad, why in (((0x20, 0, 0, 0, 0), "a sixth button bit has nowhere to go"),
                 ((0, 128, 0, 0, 0), "128 is more than one report"),
                 ((0, -128, 0, 0, 0), "and so is -128"),
                 ((0, 0, 0, 128, 0), "the wheel is one byte too")):
    try:
        relayctl.mouse_frame(*bad)
        check(False, f"mouse_frame should refuse: {why}")
    except relayctl.RelayError:
        check(True, f"mouse_frame refuses: {why}")

# A movement bigger than one report is split here, exactly as the page splits it, and the
# split is exact -- clamping it would move the cursor half as far as was asked and say
# nothing about it.
for dx, dy in ((300, 0), (0, -260), (300, -260), (127, 127), (1, -1), (-500, 500)):
    parts = relayctl.split_moves(dx, dy)
    check_eq(sum(p[0] for p in parts), dx, f"the X split of {dx},{dy} adds up exactly")
    check_eq(sum(p[1] for p in parts), dy, f"and the Y of {dx},{dy}")
    check(all(abs(p[0]) <= relayctl.MOUSE_DELTA_MAX and abs(p[1]) <= relayctl.MOUSE_DELTA_MAX
              for p in parts), f"and no part of {dx},{dy} exceeds one report")

check_eq(relayctl.mouse_frames("click", [], "left"),
         [(0x01, 0, 0, 0, 0), (0, 0, 0, 0, 0)],
         "a click is the button restated down then restated up")
check_eq(relayctl.mouse_frames("release", []), [(0, 0, 0, 0, 0)],
         "release restates the whole button byte as nothing")
drag = relayctl.mouse_frames("drag", ["300", "0"], "right")
check_eq(drag[0], (0x02, 0, 0, 0, 0), "a drag presses first")
check(drag[-1] == (0, 0, 0, 0, 0), "and releases last")
check(all(f[0] == 0x02 for f in drag[1:-1]),
      "and every movement in between carries the button, so it is a drag and not a move")
check_eq(sum(f[1] for f in drag), 300, "and the drag travels the whole distance")
check_eq(relayctl.mouse_frames("wheel", ["-3"]), [(0, 0, 0, -3, 0)],
         "wheel notches go where the report's wheel byte is")
for bad_args, why in ((("move", ["1"], "left"), "move needs dx and dy"),
                      (("wheel", [], "left"), "wheel needs a count"),
                      (("move", ["1", "2"], "thumb"), "an unknown button is refused"),
                      (("wiggle", [], "left"), "an unknown action is refused")):
    try:
        relayctl.mouse_frames(*bad_args)
        check(False, f"mouse_frames should refuse: {why}")
    except (relayctl.RelayError, ValueError):
        check(True, f"mouse_frames refuses: {why}")

# Every frame this tool can build must be one the firmware parses. The shape is asserted
# rather than a reimplementation of the parser, because the parser itself is tested
# against the same frames in tests/test_proto.cpp.
mframe = re.compile(r"^M=[0-9a-f]{2}:-?\d{1,3}:-?\d{1,3}:-?\d{1,3}:-?\d{1,3}$")
built = [relayctl.mouse_frame(*f)
         for f in relayctl.mouse_frames("drag", ["500", "-500", "4", "-4"], "middle")]
check(all(mframe.match(f) for f in built), f"every built frame parses ({len(built)})")
check(all(len(f) <= config_value("kWsMessageCap") for f in built),
      "and every one fits in one WebSocket message")

# The numbers are the board's. If they drift, the board starts answering E6 to frames
# this tool believed were legal.
check_eq(relayctl.MOUSE_DELTA_MAX, config_value("kMouseDeltaMax"),
         "relayctl MOUSE_DELTA_MAX == kMouseDeltaMax")
check_eq(relayctl.MOUSE_BUTTON_MASK, config_value("kMouseButtonMask"),
         "relayctl MOUSE_BUTTON_MASK == kMouseButtonMask")
check_eq(int(re.search(r"const MOUSE_DELTA_MAX = (\d+);", page_src).group(1)),
         config_value("kMouseDeltaMax"), "page MOUSE_DELTA_MAX == kMouseDeltaMax")
check_eq(int(re.search(r"const MOUSE_BUTTON_MASK = (0x[0-9a-f]+);", page_src).group(1), 16),
         config_value("kMouseButtonMask"), "page MOUSE_BUTTON_MASK == kMouseButtonMask")
check_eq(int(re.search(r"const MOUSE_FRAME_MS = (\d+);", page_src).group(1)),
         relayctl.MOUSE_FRAME_MS, "the page and the CLI pace mouse frames identically")

# Button bits, agreed between the page and this tool even though the page keys them by
# DOM number and this keys them by name. Getting one wrong is a clean frame that is the
# wrong button -- the Alt+Tab failure again, one layer down.
check_eq({name: bits for name, bits in relayctl.MOUSE_BUTTONS.items()},
         {"left": 0x01, "right": 0x02, "middle": 0x04, "back": 0x08, "forward": 0x10},
         "the button bits are the report descriptor's order")
check("0: 0x01, 2: 0x02, 1: 0x04, 3: 0x08, 4: 0x10" in page_src,
      "the page maps DOM button numbers to those same bits (DOM 1 is middle, 2 is right)")

print("protocol strings")

check_eq("S=%02x:%s" % (0x0A, "04,28"), "S=0a:04,28", "state frame format")
check_eq("Y=%02x:%s" % (0x05, "4c"), "Y=05:4c", "chord frame format")
check_eq("T" + "hunter2", "Thunter2", "auth frame format")

# The single-character commands relayctl emits must be ones proto.h accepts, and
# proto.h's own documented list must still mention them: a change on either side
# has to be a deliberate edit in two places.
for command in ("A1", "A0", "R", "H", "K"):
    check(f'"{command}"' in relay_src, f'relayctl sends "{command}"')
for command in ("'A'", "'R'", "'H'", "'K'", "'T'", "'S'", "'Y'", "'C'", "'P'", "'M'"):
    check(command in proto_src, f"proto.h parses {command}")
check("case 'M':" in proto_src, "proto.h routes M= to the mouse parser")

# The two mouse error codes have to be in the README's list. A client that cannot look up
# what E11 means cannot tell the operator why the pointer is not moving, and this suite is
# the only place that notices when the code and the document stop agreeing.
readme_src = (root / "README.md").read_text()
for code, where in (("10 the mouse interface is off", "proto.h"),
                    ("11 the target asked for boot protocol", "proto.h")):
    check(code in proto_src, f"proto.h's error list documents {code}")
    check(code.split()[0] in readme_src,
          f"README's error list has code {code.split()[0]}")
check("mouse interface is off" in readme_src and "boot protocol" in readme_src,
      "README explains what both mouse refusals mean, not just their numbers")
check("`M=<btn>`" in readme_src or "M=<btn>" in readme_src or "`M=`" in readme_src,
      "README's wire protocol table lists the M= command")
check("if (c == 'Y' && n == 0 && out.mods == 0)" in proto_src,
      "proto.h still requires a chord to press something")

print("the fail-safe wiring")
# Source pins, and deliberately honest about being source pins: control_server.cpp and
# main.cpp need a socket and a board, so nothing here runs them. What they were caught
# doing is invisible to a compiler and to every behavioural test on the host, and each of
# these three is a sentence that a later edit can quietly undo.
main_src = (root / "src" / "main.cpp").read_text()
server_src = (root / "src" / "control_server.cpp").read_text()

# The backstop must not ask whether a controller is linked: a controller that has gone
# away is exactly the case where linked() is false, which made the silence timer
# uncallable in the one situation it exists for -- the board sat armed with a key held
# forever, and read-only /status traffic kept it that way.
check("if (armed && control.controllerSeen() &&" in main_src,
      "loop()'s silence disarm runs off controllerSeen(), not linked()")
check("if (armed && control.linked()" not in main_src,
      "nothing gates an armed-and-silent disarm on the link still being there")

# Accepting a TCP connection is not the operator being alive.
accept = server_src[server_src.index("void ControlServer::acceptNew"):
                    server_src.index("size_t ControlServer::readSome")]
# Code lines only -- the function carries a comment explaining why it does NOT stamp the
# clock, and a check that flagged the comment would be worse than no check at all.
accept_code = "\n".join(l for l in accept.splitlines()
                        if not l.strip().startswith(("//", "*", "/*")))
check("_lastInboundMs" not in accept_code,
      "acceptNew() does not refresh the controller-silence clock")

# And closing a read-only /status request must not refresh it either -- the first fix
# stamped it on every close, which held a dead controller's armed state open exactly as
# effectively as the original bug did.
close_fn = server_src[server_src.index("void ControlServer::closeLink"):
                      server_src.index("}", server_src.index("_onLinkChange(_ctx, false)"))]
check("if (wasLinked) _lastInboundMs = millis();" in close_fn,
      "closeLink() stamps the silence clock only when a controller link is what closed")
check("const bool wasLinked = _phase == Phase::WebSocket;" in close_fn,
      "closeLink() treats anything that reached the WebSocket phase as able to leave "
      "keys down, not only what is still marked authenticated")

print("the target-side verifier")
# Twice now the verifier itself was the bug -- mislabelled LEDs, unnamed key codes that
# looked like lost keystrokes -- so the thing that decides what "it arrived" means gets
# tested too, and the mouse half of it especially: it is new, and a watcher that misread
# REL_X would either cry wolf or miss a mouse that worked.
import watch_target as wt                                    # noqa: E402

# EV_REL is what makes a node a pointing device. Nothing here trusts a device name.
check_eq(wt.classify(1 << wt.EV_REL | 1 << wt.EV_KEY), "mouse", "EV_REL means mouse")
check_eq(wt.classify(1 << wt.EV_KEY | 1 << wt.EV_LED), "keyboard",
         "EV_KEY without EV_REL means keyboard")
check_eq(wt.classify(1 << wt.EV_ABS), None,
         "an absolute-only device is neither, and is not watched as a mouse")
check_eq(wt.capability_bits("/nonexistent/does/not/exist"), 0,
         "an unreadable sysfs file is 0 rather than an exception")

m = wt.MouseWatch(False)
check(m.event(0, 0, 20) is None and m.event(0, 1, -6) is None,
      "a raw axis event is not a line -- one report is one line")
check(m.event(0, 0, 15) is None, "and more movement in the same report accumulates")
line = m.flush(0.1)
check(line is not None and "+35" in line and "-6" in line,
      f"the report comes out as the sum of its axes ({line!r})")
check(m.flush(0.2) is None, "a report with no movement prints nothing")
check_eq(m.x, 35, "net X is the exact sum of what the target received")
check_eq(m.y, -6, "and so is net Y")
check_eq(m.travel, 41, "and the path length counts both, not the diagonal")

# REL_WHEEL_HI_RES is 1/240 of a notch. Counting it as a notch would report 240x the
# scroll that happened, which is the same class of error as the mislabelled LEDs.
# The axis numbers, pinned to the kernel's, not to the HID report's byte order. Getting
# this wrong made a working wheel look dead -- see the comment on REL_NAMES.
check_eq((wt.REL_X, wt.REL_Y, wt.REL_WHEEL, wt.REL_HWHEEL), (0, 1, 8, 6),
         "axis codes are the kernel's: REL_WHEEL is 8, not 2")
check_eq((wt.REL_WHEEL_HI_RES, wt.REL_HWHEEL_HI_RES), (11, 12),
         "and the hi-res wheel codes are 11/12, not 6/7")

m2 = wt.MouseWatch(False)
check(m2.event(0, wt.REL_WHEEL, -3) is not None,
      "a wheel notch is printed as it happens, not coalesced with movement")
m2.event(0, wt.REL_WHEEL_HI_RES, 720)   # the same notches again, finer
check_eq(m2.wheel, -3, "hi-res wheel events are not counted as extra notches")
check("REL_WHEEL" in "\n".join(m2.report()),
      "and the summary names the axis it saw, so a wheel that arrived is visible")
m3 = wt.MouseWatch(False)
m3.event(0, 40, 5)                       # an axis this tool does not know
check(any("rel40" in l for l in m3.report()),
      "an unnamed axis code is counted and shown, never dropped in silence")
m2.event(0, 272, 1)
check("still held" in "\n".join(m2.report()),
      "a button the target still believes is down is called out, in those words")
m2.event(0, 272, 0)
check("still held" not in "\n".join(m2.report()),
      "and the warning goes away when it is released")

k = wt.KeyboardWatch(False)
k.event(0, 30, 1)
k.event(0, 30, 0)
check("a" in "".join(k.typed), "the keyboard watcher still reassembles typed text")
kc = wt.KeyboardWatch(False)
for code in (29, 56, 111):
    kc.event(0, code, 1)
check(any("Ctrl+Alt+Del" in l for l in kc.report()),
      "and still names Ctrl+Alt+Del as one combination, not three keys")

print(f"relayctl              {checks} checks, {failures} failures")
sys.exit(1 if failures else 0)