#!/usr/bin/env python3
"""Read a full-screen TUI over a pty, and type at it.

`ssh_pty.py` captures the bytes a TUI draws. This reads the SCREEN, and the
difference is the whole reason a scripted wizard walk is possible at all.

ratatui redraws only the cells that changed. "step 1 of 6" becoming
"step 2 of 6" puts a single `2` on the wire behind a cursor move, so a script
watching the byte stream for the new title sees it exactly once — on the first
frame, where everything is written — and then silently never again. Every
subsequent screen is a diff against a screen the script never reconstructed.
Grepping the stream is not a weaker version of reading the screen; it is a
different thing that happens to work on frame one.

So this keeps a grid and applies the escapes to it. `Screen.text()` is what is
on the screen right now, which is a question with an answer, and
`Session.wait_for` is "wait until the screen says this".

The escape subset is crossterm's output and nothing more: absolute and relative
cursor moves, the two erases, and printable text. SGR, the alternate-screen
toggles, OSC and DCS are parsed only far enough to be skipped, because they
carry nothing a script reads. A sequence this does not know is skipped rather
than printed, which is the failure that would otherwise fill the grid with
`[38;5;` and look like the application misbehaving.

Not a terminal emulator. No scrollback, no autowrap, no character sets, no
line-drawing translation. A full-screen TUI positions every run it draws, so
none of that is reachable from here; if something starts depending on it, the
honest fix is to say so rather than to grow this file toward xterm.

Use:

    from tui import Session
    s = Session(["ssh", ...])
    s.wait_for(r"step 1 of 6", 60)
    s.send("alloytest\r")
    print(s.screen.text())
"""

import codecs
import fcntl
import os
import pty
import re
import select
import signal
import socket
import struct
import sys
import termios
import time

ROWS, COLS = 40, 120


class Timeout(RuntimeError):
    """What the screen said when the wait ran out, so a failure is readable."""

    def __init__(self, pattern, seconds, screen):
        super().__init__(
            "waited %gs for %r; the screen said:\n\n%s\n" % (seconds, pattern, screen)
        )
        self.screen = screen


class Screen:
    """A character grid, and the escapes needed to keep it current."""

    def __init__(self, rows=ROWS, cols=COLS):
        self.rows, self.cols = rows, cols
        self.grid = [[" "] * cols for _ in range(rows)]
        self.row = self.col = 0
        # Incremental, because a read can split a multi-byte character and a
        # per-chunk decode would put a replacement character in the grid.
        self._decode = codecs.getincrementaldecoder("utf-8")("replace")
        # Holds a partial escape across reads for the same reason.
        self._buf = ""

    def text(self):
        """The screen as lines, right-trimmed. Rows are padded, and trailing
        blanks on every line would make every pattern need to allow for them."""
        return "\n".join("".join(row).rstrip() for row in self.grid)

    def feed(self, data):
        self._buf += self._decode.decode(data)
        s, i, n = self._buf, 0, len(self._buf)
        while i < n:
            ch = s[i]
            if ch == "\x1b":
                nxt = self._escape(s, i)
                if nxt is None:
                    break  # incomplete; the rest of it is in the next read
                i = nxt
                continue
            i += 1
            if ch == "\r":
                self.col = 0
            elif ch == "\n":
                self._newline()
            elif ch == "\b":
                self.col = max(0, self.col - 1)
            elif ch == "\t":
                self.col = min(self.cols - 1, (self.col // 8 + 1) * 8)
            elif ch >= " ":
                self._put(ch)
            # Everything else is a control character a TUI does not use.
        self._buf = s[i:]

    # ---- the grid ----

    def _put(self, ch):
        self.grid[self.row][self.col] = ch
        # Clamped rather than wrapped. A TUI positions every run it draws, so
        # the right margin is never reached by accident, and wrapping would
        # scroll the screen out from under a script on a full-width line.
        self.col = min(self.col + 1, self.cols - 1)

    def _newline(self):
        if self.row + 1 < self.rows:
            self.row += 1

    def _clamp(self, row, col):
        self.row = max(0, min(row, self.rows - 1))
        self.col = max(0, min(col, self.cols - 1))

    def _blank(self, row, start, end):
        for c in range(max(0, start), min(end, self.cols)):
            self.grid[row][c] = " "

    # ---- escapes ----

    def _escape(self, s, i):
        """Index just past the sequence starting at `i`, or None if it is not
        all here yet."""
        if i + 1 >= len(s):
            return None
        kind = s[i + 1]
        if kind == "[":
            j = i + 2
            while j < len(s) and "\x30" <= s[j] <= "\x3f":
                j += 1
            while j < len(s) and "\x20" <= s[j] <= "\x2f":
                j += 1
            if j >= len(s):
                return None
            self._csi(s[i + 2 : j], s[j])
            return j + 1
        if kind in "]P^_":
            # OSC, DCS, PM, APC: a string terminated by BEL or by ST.
            j = i + 2
            while j < len(s):
                if s[j] == "\x07":
                    return j + 1
                if s[j] == "\x1b":
                    if j + 1 >= len(s):
                        return None
                    if s[j + 1] == "\\":
                        return j + 2
                j += 1
            return None
        if "\x20" <= kind <= "\x2f":
            # A two-character intermediate, such as a character-set select.
            return None if i + 2 >= len(s) else i + 3
        return i + 2

    def _csi(self, params, final):
        # `?` and friends mark private modes: cursor visibility, the alternate
        # screen, bracketed paste. None of them move the cursor or erase, so
        # skipping them is not an approximation.
        if params[:1] in ("?", "<", ">", "="):
            return
        nums = [int(p) if p.isdigit() else 0 for p in params.split(";")] if params else []

        def arg(k, default=1):
            # Zero means "the default" in every one of these, per ECMA-48.
            return nums[k] if k < len(nums) and nums[k] else default

        if final in "Hf":
            self._clamp(arg(0) - 1, arg(1) - 1)
        elif final == "A":
            self._clamp(self.row - arg(0), self.col)
        elif final == "B":
            self._clamp(self.row + arg(0), self.col)
        elif final == "C":
            self._clamp(self.row, self.col + arg(0))
        elif final == "D":
            self._clamp(self.row, self.col - arg(0))
        elif final == "G":
            self._clamp(self.row, arg(0) - 1)
        elif final == "d":
            self._clamp(arg(0) - 1, self.col)
        elif final == "J":
            # Mode 0 is the default here, so `arg` is wrong for J and K.
            mode = nums[0] if nums else 0
            if mode == 0:
                self._blank(self.row, self.col, self.cols)
                for r in range(self.row + 1, self.rows):
                    self._blank(r, 0, self.cols)
            elif mode == 1:
                for r in range(0, self.row):
                    self._blank(r, 0, self.cols)
                self._blank(self.row, 0, self.col + 1)
            else:
                for r in range(self.rows):
                    self._blank(r, 0, self.cols)
        elif final == "K":
            mode = nums[0] if nums else 0
            if mode == 0:
                self._blank(self.row, self.col, self.cols)
            elif mode == 1:
                self._blank(self.row, 0, self.col + 1)
            else:
                self._blank(self.row, 0, self.cols)
        # SGR and everything else changes how the screen looks, not what it says.


class Session:
    """A child under a pty of a known size, and the screen it is drawing."""

    def __init__(self, argv, rows=ROWS, cols=COLS, log_path=None):
        self.screen = Screen(rows, cols)
        self.argv = argv
        self.log = open(log_path, "wb") if log_path else None
        self.pid, self.fd = pty.fork()
        if self.pid == 0:
            os.execvp(argv[0], argv)
            os._exit(127)
        fcntl.ioctl(self.fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
        self.eof = False

    def pump(self, seconds=0.2):
        """Read for up to `seconds`, applying whatever arrives to the screen."""
        end = time.time() + seconds
        while time.time() < end:
            r, _, _ = select.select([self.fd], [], [], max(0.0, end - time.time()))
            if not r:
                continue
            try:
                chunk = os.read(self.fd, 65536)
            except OSError:
                self.eof = True
                return
            if not chunk:
                self.eof = True
                return
            self.screen.feed(chunk)
            if self.log:
                self.log.write(chunk)
                self.log.flush()

    def wait_for(self, pattern, seconds, poll=0.2):
        """Block until the screen matches, and return the match.

        A regex over the whole screen rather than over the stream: see the
        module docstring for why the stream cannot answer this."""
        end = time.time() + seconds
        rx = re.compile(pattern)
        while True:
            hit = rx.search(self.screen.text())
            if hit:
                return hit
            if time.time() >= end:
                raise Timeout(pattern, seconds, self.screen.text())
            if self.eof:
                # One more look: the last frame before the child left may
                # carry the thing being waited for.
                if rx.search(self.screen.text()):
                    return rx.search(self.screen.text())
                raise Timeout(pattern, seconds, self.screen.text() + "\n\n(the session ended)")
            self.pump(poll)

    def send(self, text, settle=0.4):
        """Type, then let the frame that answers arrive before anyone reads."""
        os.write(self.fd, text.encode())
        self.pump(settle)

    def close(self):
        try:
            os.kill(self.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
        try:
            os.waitpid(self.pid, 0)
        except ChildProcessError:
            pass
        if self.log:
            self.log.close()


class SocketSession(Session):
    """The same screen, over a unix socket instead of a pty.

    The guest's serial console is a chardev socket, and on the installer medium
    GRUB draws its menu there and `alloy-debug-shell@ttyS0` puts a root bash
    there. Both are things a scenario has to read and answer, and neither is
    reachable through a pty.

    Subclasses Session for its `wait_for` and `screen`, and replaces only the
    transport: there is no child process here, so nothing to fork, size or
    reap. The terminal size is the guest's `stty`, not an ioctl from this end.
    """

    def __init__(self, path, rows=ROWS, cols=COLS, log_path=None, connect_timeout=60.0):
        self.screen = Screen(rows, cols)
        self.rows, self.cols = rows, cols
        self.log = open(log_path, "wb") if log_path else None
        self.eof = False
        end = time.time() + connect_timeout
        while True:
            try:
                self.sock = socket.socket(socket.AF_UNIX)
                self.sock.connect(path)
                break
            except (FileNotFoundError, ConnectionRefusedError):
                if time.time() >= end:
                    raise
                time.sleep(0.2)
        self.fd = self.sock.fileno()

    def send(self, text, settle=0.4):
        self.sock.sendall(text.encode())
        self.pump(settle)

    def close(self):
        try:
            self.sock.close()
        except OSError:
            pass
        if self.log:
            self.log.close()


def self_test():
    """Check the parser, because a wrong parse still produces a screen.

    Same reason check-installed.sh and check-rust-stage.sh carry one: the part
    that fails invisibly is the part that turns input into a verdict, and a
    verdict is what this file produces. Every case below is a shape crossterm
    actually emits.
    """
    fails = []

    def check(label, got, want):
        if got != want:
            fails.append("%s\n  got:  %r\n  want: %r" % (label, got, want))

    # The case this file exists for, and the one a stream search cannot answer:
    # a redraw that rewrites one cell. Search the stream for "step 2 of 6" and
    # it is not there; read the screen and it is.
    s = Screen(3, 20)
    s.feed(b"\x1b[2J\x1b[1;1Hstep 1 of 6\x1b[1;6H2")
    check("a one-cell redraw", s.text().splitlines()[0], "step 2 of 6")

    # SGR carries how the screen looks, not what it says, and must leave
    # nothing behind. This is the failure that fills a grid with `[38;5;`.
    s = Screen(1, 30)
    s.feed(b"\x1b[1;38;5;203mred\x1b[0m ok")
    check("SGR is skipped", s.text(), "red ok")

    # Erase to end of line, which is how a footer's hints are replaced when a
    # step changes the shorter set for the longer one.
    s = Screen(1, 20)
    s.feed(b"\x1b[1;1Hesc back\x1b[1;1Hr reboot\x1b[K")
    check("erase to end of line", s.text(), "r reboot")

    # And erase-all, which is the first thing a full repaint does.
    s = Screen(2, 10)
    s.feed(b"\x1b[1;1Hgone\x1b[2;1Halso\x1b[2J\x1b[1;1Hnew")
    check("erase display", s.text(), "new\n")

    # An escape split across two reads. A parser that consumed the partial
    # sequence would print `[2;` into the grid and lose the move.
    s = Screen(2, 10)
    s.feed(b"a\x1b[2")
    s.feed(b";1Hb")
    check("a split escape", s.text(), "a\nb")

    # A UTF-8 character split across two reads. The list marker is three bytes,
    # and a per-chunk decode would put a replacement character in the gutter
    # and make every row read as unselected.
    s = Screen(1, 10)
    s.feed(b"\xe2\x96")
    s.feed(b"\xb6 vda")
    check("a split character", s.text(), "\u25b6 vda")

    # OSC, which crossterm emits to set the window title. Terminated by BEL
    # here and by ST in the next case; a parser that knew only one would print
    # the other's payload.
    s = Screen(1, 20)
    s.feed(b"\x1b]0;alloy\x07ok")
    check("OSC ended by BEL", s.text(), "ok")
    s = Screen(1, 20)
    s.feed(b"\x1b]0;alloy\x1b\\ok")
    check("OSC ended by ST", s.text(), "ok")

    # Relative moves and the column/row absolutes.
    s = Screen(2, 10)
    s.feed(b"\x1b[1;1Habc\x1b[1D\x1b[1BX\x1b[1;1H\x1b[5GY")
    # `abc` leaves the cursor at column 3; one left is 2, one down writes the X
    # there, and `5G` is column 4 counting from one.
    check("relative moves", s.text(), "abc Y\n  X")

    # A pane the way one really arrives: a border, a gutter marker, and rows.
    # The assertion is the question install_drive.py asks of the disk step.
    s = Screen(4, 40)
    s.feed(
        b"\x1b[2J"
        b"\x1b[1;1H\xe2\x94\x8c\x1b[1;2H select a disk"
        b"\x1b[2;1H\xe2\x94\x82\x1b[2;3H  vda        40.0 GiB"
        b"\x1b[3;1H\xe2\x94\x82\x1b[3;3H  sdb       931.5 GiB"
        # ...and then the selection moves down, which rewrites two gutters.
        b"\x1b[3;3H\xe2\x96\xb6"
    )
    row = re.search("\u25b6" + r"\s+(\S+)", s.text())
    check("the selected row is readable", row and row.group(1), "sdb")

    if fails:
        for f in fails:
            print("self-test: " + f, file=sys.stderr)
        print("self-test: %d failed" % len(fails), file=sys.stderr)
        return 1
    print("self-test: the parser is right about all ten shapes")
    return 0


def main():
    """Ad-hoc use: run a command, wait for a pattern, print the screen."""
    if len(sys.argv) > 1 and sys.argv[1] == "--self-test":
        raise SystemExit(self_test())
    if len(sys.argv) < 3:
        raise SystemExit("usage: tui.py <seconds> <pattern> <command> [args...]\n"
                         "       an empty pattern just reads for <seconds>")
    seconds, pattern, argv = float(sys.argv[1]), sys.argv[2], sys.argv[3:]
    s = Session(argv)
    try:
        if pattern:
            s.wait_for(pattern, seconds)
        else:
            s.pump(seconds)
        print(s.screen.text())
    except Timeout as e:
        print(e, file=sys.stderr)
        raise SystemExit(1)
    finally:
        s.close()


if __name__ == "__main__":
    main()