"""Drive the cst-agent TUI in a pty and read the screen back (pyte).

Run from the repo root after `make`:
  nix-shell -p python3Packages.pyte --run "python3 scripts/tui_keys/keys.py"
  nix-shell -p python3Packages.pyte --run "python3 scripts/tui_keys/esc.py"
(or `pip install pyte`). A local fake OpenAI-compatible server answers the
model calls, so no API key is needed. Exit status is the failure count.
"""
import os, pty, sys, time, select, fcntl, termios, struct, signal, shutil, tempfile
import pyte

BIN = os.environ.get("CST_BIN") or os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "..", "..", "cst-agent")


class Term:
    def __init__(self, cols=120, rows=36, kitty=False, env=None, args=(), home=None,
                 ws=None, term="xterm-256color"):
        self.cols, self.rows, self.kitty = cols, rows, kitty
        self.screen = pyte.Screen(cols, rows)
        self.stream = pyte.ByteStream(self.screen)
        self.home = home or tempfile.mkdtemp(prefix="cstk_home_")
        self.ws = ws or tempfile.mkdtemp(prefix="cstk_ws_")
        e = dict(os.environ)
        e.update({"TERM": term, "HOME": self.home,
                  "XDG_CONFIG_HOME": os.path.join(self.home, ".config"),
                  "LANG": "C.UTF-8", "LC_ALL": "C.UTF-8",
                  "CST_PROVIDER": "deepseek", "DEEPSEEK_API_KEY": "sk-dummy-test",
                  "CST_METRICS": "0"})
        e.pop("NO_COLOR", None)
        if env:
            e.update(env)
        self.raw = b""
        pid, fd = pty.fork()
        if pid == 0:
            os.chdir(self.ws)
            os.execve(BIN, [BIN, "-N", *args], e)
        self.pid, self.fd = pid, fd
        fcntl.ioctl(fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
        self.pump(1.5)

    def pump(self, t=0.25):
        end = time.time() + t
        while True:
            left = end - time.time()
            if left <= 0:
                break
            r, _, _ = select.select([self.fd], [], [], left)
            if not r:
                continue
            try:
                d = os.read(self.fd, 65536)
            except OSError:
                break
            if not d:
                break
            self.raw += d
            if self.kitty and b"\x1b[?u" in d:
                os.write(self.fd, b"\x1b[?0u")
            self.stream.feed(d)

    def send(self, b, t=0.15):
        if isinstance(b, str):
            b = b.encode()
        os.write(self.fd, b)
        self.pump(t)

    def appcursor(self):
        return (1 << 5) in self.screen.mode  # pyte DECCKM

    def arrow(self, d, form=None, mod=None, t=0.15):
        c = {"up": "A", "down": "B", "right": "C", "left": "D"}[d]
        if mod:
            seq = "\x1b[1;%d%s" % (mod, c)
        else:
            f = form or ("ss3" if self.appcursor() else "csi")
            seq = ("\x1bO" if f == "ss3" else "\x1b[") + c
        self.send(seq, t)

    def lines(self):
        return [l.rstrip() for l in self.screen.display]

    def find(self, s):
        return [(i, l) for i, l in enumerate(self.lines()) if s in l]

    def cursor(self):
        return self.screen.cursor.y, self.screen.cursor.x

    def dump(self, title=""):
        print("=" * 20, title, "cursor", self.cursor())
        for i, l in enumerate(self.lines()):
            print("%2d|%s" % (i, l))

    def close(self):
        try:
            os.kill(self.pid, signal.SIGTERM)
            os.waitpid(self.pid, 0)
        except Exception:
            pass