#!/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):
    found = re.search(rf"\b{name}\s*=\s*(\d+)\s*;", config_src)
    return int(found.group(1)) 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("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'"):
    check(command in proto_src, f"proto.h parses {command}")
check("if (c == 'Y' && n == 0 && out.mods == 0)" in proto_src,
      "proto.h still requires a chord to press something")

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