#!/usr/bin/env python3
# Copyright (c) 2026 Storm Dragon
# Authors: Storm Dragon, Angus Kola
# SPDX-License-Identifier: GPL-3.0-or-later
"""Private Xorg Remote feasibility experiments, not a production input helper.

Always run this file without arguments. The launcher replaces DISPLAY with Xvfb
before creating any X connection. All keys in these experiments are synthetic;
they cannot establish physical/injected collision or hardware recovery safety.
"""
import ctypes as C
import os
from pathlib import Path
import select
import signal
import subprocess
import sys
import tempfile
import time
import unittest

from native_test_session import run_private_session


class KeyEvent(C.Structure):
    _fields_ = [("type", C.c_int), ("serial", C.c_ulong), ("send_event", C.c_int),
                ("display", C.c_void_p), ("window", C.c_ulong), ("root", C.c_ulong),
                ("subwindow", C.c_ulong), ("time", C.c_ulong), ("x", C.c_int),
                ("y", C.c_int), ("x_root", C.c_int), ("y_root", C.c_int),
                ("state", C.c_uint), ("keycode", C.c_uint), ("same_screen", C.c_int)]


class Event(C.Union):
    _fields_ = [("type", C.c_int), ("key", KeyEvent), ("pad", C.c_long * 24)]


class X11:
    def __init__(self):
        self.x = C.CDLL("libX11.so.6")
        self.xtst = C.CDLL("libXtst.so.6")
        signatures = {
            "XOpenDisplay": ([C.c_char_p], C.c_void_p),
            "XCloseDisplay": ([C.c_void_p], C.c_int),
            "XDefaultRootWindow": ([C.c_void_p], C.c_ulong),
            "XStringToKeysym": ([C.c_char_p], C.c_ulong),
            "XKeysymToKeycode": ([C.c_void_p, C.c_ulong], C.c_ubyte),
            "XSync": ([C.c_void_p, C.c_int], C.c_int),
            "XQueryKeymap": ([C.c_void_p, C.POINTER(C.c_ubyte)], C.c_int),
            "XCreateSimpleWindow": ([C.c_void_p, C.c_ulong, C.c_int, C.c_int,
                                     C.c_uint, C.c_uint, C.c_uint, C.c_ulong,
                                     C.c_ulong], C.c_ulong),
            "XSelectInput": ([C.c_void_p, C.c_ulong, C.c_long], C.c_int),
            "XMapWindow": ([C.c_void_p, C.c_ulong], C.c_int),
            "XSetInputFocus": ([C.c_void_p, C.c_ulong, C.c_int, C.c_ulong], C.c_int),
            "XGrabKeyboard": ([C.c_void_p, C.c_ulong, C.c_int, C.c_int,
                               C.c_int, C.c_ulong], C.c_int),
            "XUngrabKeyboard": ([C.c_void_p, C.c_ulong], C.c_int),
            "XPending": ([C.c_void_p], C.c_int),
            "XNextEvent": ([C.c_void_p, C.POINTER(Event)], C.c_int),
        }
        for name, (arguments, result) in signatures.items():
            function = getattr(self.x, name)
            function.argtypes, function.restype = arguments, result
        self.xtst.XTestFakeKeyEvent.argtypes = [C.c_void_p, C.c_uint, C.c_int, C.c_ulong]
        self.xtst.XTestFakeKeyEvent.restype = C.c_int
        self.display = self.x.XOpenDisplay(None)
        if not self.display:
            raise RuntimeError("private X display unavailable")
        self.root = self.x.XDefaultRootWindow(self.display)

    def close(self):
        if self.display:
            self.x.XCloseDisplay(self.display)
            self.display = None

    def sync(self):
        self.x.XSync(self.display, False)

    def code(self, name):
        code = self.x.XKeysymToKeycode(self.display, self.x.XStringToKeysym(name.encode()))
        if not code:
            raise ValueError("fixture keysym unavailable")
        return code

    def key(self, code, down):
        if not self.xtst.XTestFakeKeyEvent(self.display, code, down, 0):
            raise RuntimeError("XTEST failed")
        self.sync()

    def held(self, code):
        keys = (C.c_ubyte * 32)()
        self.x.XQueryKeymap(self.display, keys)
        return bool(keys[code // 8] & (1 << (code % 8)))

    def focus_window(self):
        window = self.x.XCreateSimpleWindow(self.display, self.root, 0, 0, 200, 100, 0, 0, 0)
        self.x.XSelectInput(self.display, window, 3)  # KeyPressMask | KeyReleaseMask
        self.x.XMapWindow(self.display, window)
        self.x.XSetInputFocus(self.display, window, 1, 0)
        self.sync()

    def grab(self):
        return self.x.XGrabKeyboard(self.display, self.root, False, 1, 1, 0)

    def events(self):
        self.sync()
        result = []
        while self.x.XPending(self.display):
            event = Event()
            self.x.XNextEvent(self.display, C.byref(event))
            if event.type in (2, 3):
                result.append((event.type, event.key.keycode, bool(event.key.send_event)))
        return result


def helper():
    """One private descriptor, fixed two-byte records, bounded heartbeat lifetime.

    This deliberately has no network, lease API, production entry point or
    parent-death signal. EOF must allow the independent owner to release keys.
    Stdout acknowledgements are test synchronization, not a production protocol.
    """
    x = X11()
    held = set()
    pending = b""
    deadline = time.monotonic() + 3
    reason = "eof"
    print("ready", flush=True)
    try:
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0 or not select.select([0], [], [], max(0, remaining))[0]:
                reason = "expired"
                break
            data = os.read(0, 2 - len(pending))
            if not data:
                break
            pending += data
            if len(pending) != 2:
                continue
            if time.monotonic() >= deadline:
                reason = "expired"
                break
            operation, code = pending
            pending = b""
            # Only an explicit heartbeat renews the watchdog, not key traffic.
            if operation == ord("H") and code == 0:
                deadline = time.monotonic() + 3
            elif operation in (ord("D"), ord("U")) and 8 <= code <= 255:
                if operation == ord("D"):
                    if code not in held:
                        if x.held(code):
                            reason = "collision"
                            break
                        held.add(code)
                        x.key(code, True)
                elif code in held:
                    x.key(code, False)
                    held.remove(code)
            else:
                reason = "invalid"
                break
            print("ok", flush=True)
    finally:
        for code in sorted(held):
            x.key(code, False)
        x.close()
    print(reason, flush=True)


def owner(write_fd):
    # Simulates the reader's exclusive write end. Killing/stopping this process
    # must not kill the helper before it has performed cleanup.
    try:
        while data := os.read(0, 256):
            os.write(write_fd, data)  # <= PIPE_BUF; blocking pipe, one writer.
    finally:
        os.close(write_fd)


class RemoteInputTests(unittest.TestCase):
    def setUp(self):
        self.x = X11()
        self.addCleanup(self.x.close)
        self.x.focus_window()
        self.shift = self.x.code("Shift_L")
        self.control = self.x.code("Control_L")
        self.a = self.x.code("a")
        # The fixture alone owns these synthetic keys, even after failed tests.
        self.addCleanup(self.release_fixture_keys)

    def release_fixture_keys(self):
        for code in (self.shift, self.control, self.a):
            self.x.key(code, False)

    def spawn(self, *arguments, **kwargs):
        process = subprocess.Popen([sys.executable, "-B", __file__, *arguments], **kwargs)
        def stop():
            if process.poll() is None:
                process.kill()
            process.wait(timeout=5)
            for stream in (process.stdin, process.stdout, process.stderr):
                if stream:
                    stream.close()
        self.addCleanup(stop)
        return process

    def line(self, process):
        self.assertTrue(select.select([process.stdout], [], [], 5)[0], "helper response timed out")
        line = process.stdout.readline()
        self.assertTrue(line, "helper exited without response")
        return line.rstrip().decode()

    def start_helper(self):
        read_fd, write_fd = os.pipe()
        try:
            # The test owns neither pipe end after launch. The two children are
            # siblings so the harness can reap both after killing the owner.
            helper_process = self.spawn("--helper", stdin=read_fd, stdout=subprocess.PIPE, bufsize=0)
            owner_process = self.spawn("--owner", str(write_fd), pass_fds=(write_fd,),
                                       stdin=subprocess.PIPE, bufsize=0)
        finally:
            os.close(read_fd)
            os.close(write_fd)
        self.assertEqual(self.line(helper_process), "ready")
        return owner_process, helper_process

    def send(self, owner_process, helper_process, operation, code=0):
        owner_process.stdin.write(bytes((ord(operation), code)))
        self.assertEqual(self.line(helper_process), "ok")

    def assert_recovered(self, helper_process, reason, started):
        self.assertEqual(self.line(helper_process), reason)
        self.assertEqual(helper_process.wait(timeout=5), 0)
        elapsed = time.monotonic() - started
        self.assertLess(elapsed, 4, "cleanup exceeded watchdog plus scheduling allowance")
        self.assertFalse(self.x.held(self.shift))
        self.assertFalse(self.x.held(self.control))
        self.x.events()
        self.x.key(self.a, True)
        self.x.key(self.a, False)
        self.assertEqual(self.x.events(), [(2, self.a, False), (3, self.a, False)])
        print(f"{reason} cleanup observed within {elapsed:.3f}s", flush=True)

    def test_reader_sigkill_releases_owned_modifiers(self):
        owner_process, helper_process = self.start_helper()
        self.send(owner_process, helper_process, "D", self.shift)
        self.send(owner_process, helper_process, "D", self.control)
        self.assertTrue(self.x.held(self.shift))
        self.assertTrue(self.x.held(self.control))
        started = time.monotonic()
        owner_process.kill()
        owner_process.wait(timeout=5)
        self.assert_recovered(helper_process, "eof", started)

    def test_stopped_reader_expires_without_eof(self):
        owner_process, helper_process = self.start_helper()
        self.send(owner_process, helper_process, "D", self.shift)
        self.send(owner_process, helper_process, "H")
        started = time.monotonic()
        owner_process.send_signal(signal.SIGSTOP)
        self.assert_recovered(helper_process, "expired", started)

    def test_heartbeat_and_balanced_owned_keys(self):
        owner_process, helper_process = self.start_helper()
        self.send(owner_process, helper_process, "D", self.shift)
        self.send(owner_process, helper_process, "D", self.shift)  # Duplicate down.
        self.send(owner_process, helper_process, "U", self.control)  # Not owned.
        self.send(owner_process, helper_process, "H")
        self.assertTrue(self.x.held(self.shift))
        self.send(owner_process, helper_process, "U", self.shift)
        self.send(owner_process, helper_process, "U", self.shift)
        self.assertFalse(self.x.held(self.shift))
        owner_process.stdin.close()
        self.assert_recovered(helper_process, "eof", time.monotonic())

    def test_heartbeat_extends_expiry(self):
        owner_process, helper_process = self.start_helper()
        self.send(owner_process, helper_process, "D", self.shift)
        # This is a real watchdog timing experiment. Cross the original expiry
        # while keeping the refreshed deadline comfortably in the future.
        time.sleep(2)
        self.send(owner_process, helper_process, "H")
        time.sleep(1.2)
        self.assertIsNone(helper_process.poll())
        self.assertTrue(self.x.held(self.shift))
        self.send(owner_process, helper_process, "U", self.shift)
        owner_process.stdin.close()
        self.assert_recovered(helper_process, "eof", time.monotonic())

    def test_existing_external_key_is_not_released(self):
        owner_process, helper_process = self.start_helper()
        self.x.key(self.shift, True)
        self.send(owner_process, helper_process, "U", self.shift)
        self.assertTrue(self.x.held(self.shift))
        self.send(owner_process, helper_process, "D", self.control)
        owner_process.stdin.write(bytes((ord("D"), self.shift)))
        self.assertEqual(self.line(helper_process), "collision")
        self.assertEqual(helper_process.wait(timeout=5), 0)
        self.assertTrue(self.x.held(self.shift), "cleanup released an unowned external key")
        self.assertFalse(self.x.held(self.control), "collision failed to clean owned keys")

    def test_later_same_code_xtest_collision_remains_unresolved(self):
        owner_process, helper_process = self.start_helper()
        self.send(owner_process, helper_process, "D", self.shift)
        # Both clients use the server's XTEST keyboard. The global bitmap cannot
        # represent separate ownership when a second source presses this code.
        self.x.key(self.shift, True)
        owner_process.stdin.close()
        self.assertEqual(self.line(helper_process), "eof")
        self.assertEqual(helper_process.wait(timeout=5), 0)
        self.assertFalse(self.x.held(self.shift))
        print("later same-code XTEST press is not preserved by ledger cleanup; "
              "physical-device collision safety remains unproven", flush=True)

    def test_fragmented_record_eof_cleans_owned_keys(self):
        owner_process, helper_process = self.start_helper()
        self.send(owner_process, helper_process, "D", self.shift)
        owner_process.stdin.write(b"D")
        owner_process.stdin.close()
        self.assert_recovered(helper_process, "eof", time.monotonic())

    def test_invalid_record_cleans_owned_keys(self):
        owner_process, helper_process = self.start_helper()
        self.send(owner_process, helper_process, "D", self.shift)
        owner_process.stdin.write(b"?\x00")
        self.assert_recovered(helper_process, "invalid", time.monotonic())

    def test_exclusive_capture_and_grab_owner_sigkill(self):
        capture = self.spawn("--capture", stdin=subprocess.PIPE, stdout=subprocess.PIPE, bufsize=0)
        self.assertEqual(self.line(capture), "ready")
        self.x.key(self.a, True)
        self.x.key(self.a, False)
        capture.stdin.write(b"e")
        self.assertEqual(self.line(capture), f"[(2, {self.a}, False), (3, {self.a}, False)]")
        self.assertEqual(self.x.events(), [], "captured keys reached the focused application")
        capture.kill()
        capture.wait(timeout=5)
        # Synchronize with server teardown by acquiring the released grab.
        deadline = time.monotonic() + 3
        while self.x.grab() != 0:
            self.assertLess(time.monotonic(), deadline, "dead owner's grab survived")
            time.sleep(0.005)
        self.x.x.XUngrabKeyboard(self.x.display, 0)
        self.x.sync()
        self.x.key(self.a, True)
        self.x.key(self.a, False)
        self.assertEqual(self.x.events(), [(2, self.a, False), (3, self.a, False)])

    def test_xtest_connection_close_does_not_release_key(self):
        source = X11()
        try:
            source.key(self.shift, True)
        finally:
            source.close()
        self.assertTrue(self.x.held(self.shift), "X server behavior changed; revisit recovery design")
        # XTEST's send_event is false, so that bit cannot authorize local consent.
        self.assertIn((2, self.shift, False), self.x.events())


def capture():
    x = X11()
    try:
        if x.grab() != 0:
            raise RuntimeError("private keyboard grab failed")
        print("ready", flush=True)
        while os.read(0, 1):
            print(x.events(), flush=True)
    finally:
        x.close()


def main():
    if len(sys.argv) == 1:
        with tempfile.TemporaryDirectory(prefix="cthulhu-remote-input-") as temporary:
            environment = dict(os.environ, CTHULHU_REMOTE_INPUT_PRIVATE=temporary,
                               XDG_CONFIG_HOME=temporary, GOCACHE="/tmp/cthulhu-cache",
                               GOPROXY="off")
            return run_private_session(
                ["xvfb-run", "-a", "-s", "-screen 0 1024x768x24", "dbus-run-session", "--",
                 sys.executable, "-B", __file__, "--inside"], env=environment, check=True).returncode
    if not os.environ.get("CTHULHU_REMOTE_INPUT_PRIVATE"):
        sys.exit("Use native_remote_input.py without arguments to create a private display")
    if sys.argv[1] == "--helper":
        helper()
    elif sys.argv[1] == "--owner":
        owner(int(sys.argv[2]))
    elif sys.argv[1] == "--capture":
        capture()
    elif sys.argv[1] == "--inside":
        suite = unittest.defaultTestLoader.loadTestsFromTestCase(RemoteInputTests)
        if not unittest.TextTestRunner(verbosity=2).run(suite).wasSuccessful():
            return 1
        # Check the real command owner's existing exclusive-capture and restore
        # paths as well as the standalone recovery experiment.
        subprocess.run([
            "go", "test", "-race", "./internal/platform/xorg", "-count=1", "-timeout=60s",
            "-run", "^TestNative(InternalUICapturesAndReleases|LearnModeKeepsKeysFromFocusedApplication|RemoteInjectedKeyObservation|RemoteRepeatNormalizesHeldKey|RemoteCaptureForwardsBypassesAndEscapes|LocalUIRejectsInjectedApproval|SuspendedQueuedKey|XTermHandoffKeepsKeypadReleased)$",
            "-v"], cwd=Path(__file__).resolve().parents[1], check=True,
            env=dict(os.environ, CTHULHU_NATIVE_X11="1", CTHULHU_NATIVE_REMOTE_INPUT="1"))
    else:
        sys.exit("unknown private fixture mode")
    return 0


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