"""End-to-end driver for the mplex client/server: runs the same scenarios against tai.py --mplex
(the oracle) and against build/mplex_cli --mplex, through a real pty, a real detached server and
real bash panes, then compares what the attached terminal shows and what the command line reports.

usage: python3 -I mplex_e2e.py <repo root> [--only NAME] [--cpp PATH] [--python-only]
"""
import fcntl
import importlib.util
import json
import os
import pty
import select
import shutil
import signal
import struct
import subprocess
import sys
import tempfile
import termios
import time

ROOT = sys.argv[1]
ARGS = sys.argv[2:]
ONLY = ARGS[ARGS.index("--only") + 1] if "--only" in ARGS else None
CPP = ARGS[ARGS.index("--cpp") + 1] if "--cpp" in ARGS else os.path.join(ROOT, "taicpp", "build", "mplex_cli")
PYTHON_ONLY = "--python-only" in ARGS

spec = importlib.util.spec_from_file_location("tai_mod", ROOT + "/tai.py")
tai = importlib.util.module_from_spec(spec)
sys.modules["tai_mod"] = tai
spec.loader.exec_module(tai)

PREFIX = b"\x02"
UP, DOWN, RIGHT, LEFT = b"\x1b[A", b"\x1b[B", b"\x1b[C", b"\x1b[D"
ROWS, COLS = 24, 80


class Harness:
    """One implementation (argv prefix) in a throwaway TAI_HOME with a throwaway bash home."""

    def __init__(self, label, argv, config=None):
        self.label = label
        self.argv = argv
        self.tmp = tempfile.mkdtemp(prefix="tai-mplex-e2e-")
        self.home = os.path.join(self.tmp, "tai")
        os.makedirs(self.home)
        self.bash_home = os.path.join(self.tmp, "home")
        os.makedirs(self.bash_home)
        with open(os.path.join(self.bash_home, ".bash_profile"), "w") as handle:
            handle.write("PS1='$ '\nunset PROMPT_COMMAND\nHISTFILE=/dev/null\n")
        base = 'status-right = "R"\nautomatic-rename = false\n' if config is None else config
        with open(os.path.join(self.home, "mplex.toml"), "w") as handle:
            handle.write("[options]\n" + base if "[options]" not in base else base)
        self.env = dict(os.environ)
        for name in ("TAI_PLEXER_ACTIVE", "TAI_PLEXER_SOCKET", "TAI_PLEXER_PANE", "TMUX", "VISUAL", "EDITOR"):
            self.env.pop(name, None)
        self.env.update({"TAI_HOME": self.home, "HOME": self.bash_home, "TERM": "xterm-256color", "SHELL": "/bin/bash",
                         "TAI_PASSPHRASE": ""})
        self.clients = []

    def cli(self, *words):
        done = subprocess.run(self.argv + list(words), env=self.env, stdin=subprocess.DEVNULL, capture_output=True,
                              text=True, timeout=30)
        return done.returncode, done.stdout, done.stderr

    def attach(self, *words):
        client = Client(self, list(words))
        self.clients.append(client)
        return client

    def close(self):
        for client in self.clients:
            client.kill()
        try:
            self.cli("kill-server")
        except Exception:  # noqa: BLE001
            pass
        time.sleep(0.3)
        shutil.rmtree(self.tmp, ignore_errors=True)


class Client:
    def __init__(self, harness, words):
        self.harness = harness
        master, slave = pty.openpty()
        fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", ROWS, COLS, 0, 0))
        self.master = master
        self.proc = subprocess.Popen(harness.argv + words, stdin=slave, stdout=slave, stderr=slave, env=harness.env,
                                     start_new_session=True, close_fds=True)
        os.close(slave)
        self.screen = tai.TuiVtScreen(ROWS, COLS)
        self.raw = b""

    def pump(self, seconds=0.2):
        end = time.monotonic() + seconds
        while time.monotonic() < end:
            ready, _w, _x = select.select([self.master], [], [], 0.05)
            if ready:
                try:
                    data = os.read(self.master, 65536)
                except OSError:
                    return
                if not data:
                    return
                self.raw += data
                self.screen.feed(data)

    def send(self, data, settle=0.25):
        os.write(self.master, data)
        self.pump(settle)

    def text(self):
        return self.screen.text()

    def wait_for(self, label, predicate, timeout=10.0):
        end = time.monotonic() + timeout
        while time.monotonic() < end:
            self.pump(0.1)
            if predicate(self.text()):
                return
        raise AssertionError("%s: timed out waiting for %s; screen:\n%s" % (self.harness.label, label, self.text()))

    def alive(self):
        return self.proc.poll() is None

    def wait_exit(self, timeout=10.0):
        end = time.monotonic() + timeout
        while time.monotonic() < end:
            self.pump(0.1)
            if self.proc.poll() is not None:
                self.pump(0.2)
                return self.proc.returncode
        raise AssertionError("%s: client did not exit" % self.harness.label)

    def kill(self):
        if self.proc.poll() is None:
            try:
                os.killpg(self.proc.pid, signal.SIGKILL)
            except OSError:
                pass
        try:
            os.close(self.master)
        except OSError:
            pass


def screen_body(client):
    """The screen minus the status row, trailing blanks stripped, for comparison."""
    lines = client.text().split("\n")
    return "\n".join(line.rstrip() for line in lines)


# ---- scenarios: each returns a list of (label, observation) ----------------------------------

def sc_lifecycle(h):
    seen = []
    client = h.attach()
    client.wait_for("shell prompt and window list", lambda t: "0:bash*" in t and "$" in t)
    seen.append(("fresh", screen_body(client)))
    client.send(b"echo hello-$((40+2))\r", 0.6)
    seen.append(("echo", screen_body(client)))
    rc, out, err = h.cli("list-sessions", "-F", "#{session_name}:#{session_windows}")
    seen.append(("ls", (rc, out, err)))
    client.send(PREFIX + b"d", 0.5)
    code = client.wait_exit()
    seen.append(("detach exit", code))
    rc, out, err = h.cli("list-sessions", "-F", "#{session_name}:#{session_attached}")
    seen.append(("ls after detach", (rc, out, err)))
    again = h.attach("attach")
    again.wait_for("reattach", lambda t: "hello-42" in t)
    seen.append(("reattached", screen_body(again)))
    again.send(b"exit\r", 0.8)
    code = again.wait_exit()
    seen.append(("exit code", code))
    time.sleep(0.5)
    seen.append(("server gone", h.cli("list-sessions")[0]))
    return seen


def sc_panes_windows(h):
    seen = []
    client = h.attach()
    client.wait_for("start", lambda t: "0:bash*" in t)
    client.send(PREFIX + b"%", 0.6)
    client.send(b"echo right-pane\r", 0.5)
    seen.append(("split h", screen_body(client)))
    client.send(PREFIX + b'"', 0.6)
    seen.append(("split v", screen_body(client)))
    client.send(PREFIX + b"o", 0.3)
    client.send(PREFIX + b"c", 0.7)
    client.wait_for("second window", lambda t: "1:bash*" in t)
    seen.append(("new window", screen_body(client)))
    client.send(PREFIX + b"p", 0.4)
    seen.append(("previous", screen_body(client)))
    client.send(PREFIX + b"z", 0.4)
    seen.append(("zoom", screen_body(client)))
    client.send(PREFIX + b"z", 0.4)
    client.send(PREFIX + b"," + b"\x15" + b"renamed\r", 0.5)
    seen.append(("rename", screen_body(client)))
    for cmd in (["list-windows"], ["list-panes"], ["list-panes", "-a", "-F", "#{pane_index} #{pane_width}x#{pane_height}"],
                ["display-message", "-p", "#{window_name} #{window_panes}"], ["select-layout", "even-vertical"],
                ["list-panes", "-F", "#{pane_index} #{pane_width}x#{pane_height} #{pane_top} #{pane_left}"]):
        seen.append((" ".join(cmd), tuple(x.replace(h.home, "HOME") if isinstance(x, str) else x for x in h.cli(*cmd))))
    client.send(PREFIX + b"x", 0.4)
    client.send(b"y", 0.6)
    seen.append(("kill pane", screen_body(client)))
    return seen


def sc_commands(h):
    seen = []
    assert h.cli("new-session", "-d", "-s", "work")[0] == 0
    time.sleep(0.5)
    for cmd in (["list-sessions", "-F", "#{session_name}"], ["has-session", "-t", "work"], ["has-session", "-t", "nope"],
                ["show-options", "-g"], ["show-options", "-g", "status"], ["set-option", "-g", "status-right", "X"],
                ["show-options", "-gv", "status-right"], ["set-option", "-g", "nosuch", "1"], ["bogus-command"],
                ["list-commands"], ["list-keys", "-T", "root"], ["list-keys", "-N"], ["list-keys", "-T", "copy-mode-vi"],
                ["bind-key", "-n", "M-x", "display-message", "hi"], ["list-keys", "-T", "root"],
                ["send-keys", "-t", "work", "echo marker-$((1+2))", "Enter"], ["new-window", "-t", "work", "-n", "two"],
                ["list-windows", "-t", "work", "-F", "#{window_index}:#{window_name}:#{window_flags}"],
                ["rename-session", "-t", "work", "renamed"], ["list-sessions", "-F", "#{session_name}"],
                ["split-window", "-t", "renamed:two", "-h"], ["list-panes", "-t", "renamed:two", "-F", "#{pane_index}:#{pane_width}"],
                ["resize-pane", "-t", "renamed:two.0", "-x", "20"], ["list-panes", "-t", "renamed:two", "-F", "#{pane_index}:#{pane_width}"],
                ["set-buffer", "hello"], ["show-buffer"], ["list-buffers", "-F", "#{buffer_name}:#{buffer_size}"],
                ["capture-pane", "-p", "-t", "renamed:0"], ["kill-window", "-t", "renamed:two"],
                ["list-windows", "-t", "renamed"], ["display-message", "-p", "-t", "renamed", "#S"],
                ["new-session", "-d", "-s", "work2", "-n", "w", "sleep 30"], ["list-windows", "-t", "work2", "-F", "#{window_name}"],
                ["kill-session", "-t", "work2"], ["list-sessions", "-F", "#{session_name}"]):
        time.sleep(0.15)
        rc, out, err = h.cli(*cmd)
        seen.append((" ".join(cmd), (rc, out.replace(h.home, "HOME"), err.replace(h.home, "HOME"))))
    return seen


def sc_copy_mode(h):
    seen = []
    client = h.attach()
    client.wait_for("start", lambda t: "0:bash*" in t)
    client.send(b"seq 1 60\r", 0.8)
    client.send(PREFIX + b"[", 0.4)
    seen.append(("enter", screen_body(client)))
    client.send(b"\x1b[5~", 0.4)
    seen.append(("page up", screen_body(client)))
    client.send(b"\x02", 0.1)
    client.send(b"q", 0.4)
    client.send(PREFIX + b"[", 0.3)
    client.send(b"\x00", 0.2)
    client.send(UP + UP + UP, 0.3)
    client.send(b"\x1bw", 0.4)
    seen.append(("after copy", screen_body(client)))
    rc, out, err = h.cli("show-buffer")
    seen.append(("buffer", (rc, out, err)))
    client.send(b"echo $((6*7))\r", 0.6)
    client.send(PREFIX + b"]", 0.5)
    seen.append(("paste", screen_body(client)))
    return seen


def sc_config(h):
    seen = []
    client = h.attach()
    client.wait_for("start", lambda t: "1:bash*" in t)
    seen.append(("status", client.text().split("\n")[-1].rstrip()))
    client.send(b"\x01c", 0.7)
    client.wait_for("new window by C-a c", lambda t: "2:bash*" in t)
    seen.append(("prefix C-a c", client.text().split("\n")[-1].rstrip()))
    client.send(b"\x1b[1;3D", 0.4)
    seen.append(("M-Left", client.text().split("\n")[-1].rstrip()))
    import re
    seen.append(("show-messages", tuple(re.sub(r"\d\d:\d\d:\d\d ", "", x) if isinstance(x, str) else x for x in h.cli("show-messages"))))
    return seen


CONFIG_TOML = '''status-right = "R"
automatic-rename = false
prefix = "C-a"
base-index = 1
status-position = "top"
status-left = "<#S> "

[keys.root]
"M-Left" = "previous-window"

[keys.prefix]
"b" = false
"nosuchkey!!" = "new-window"
"x" = "nosuchcommand"
'''


def sc_errors(h):
    seen = []
    real_cli = h.cli
    h.cli = lambda *w: tuple(x.replace(h.tmp, "TMP") if isinstance(x, str) else x for x in real_cli(*w))
    seen.append(("unknown", h.cli("nosuchcommand")))
    seen.append(("ambiguous", h.cli("l")))
    seen.append(("no server", h.cli("list-sessions")))
    seen.append(("flags", h.cli("new-session", "-Z")))
    seen.append(("noterm", h.cli("new-session")))
    seen.append(("attach noterm", h.cli("attach-session")))
    return seen


def sc_native_flow(h):
    """The control socket: both controllers (tai.py's and the C++ one) against this server."""
    seen = []
    client = h.attach()
    client.wait_for("start", lambda t: "0:bash*" in t)
    cpp_test = os.path.join(ROOT, "taicpp", "build", "mplex_flow_test")
    helpers = [("py-controller", [sys.executable, os.path.join(ROOT, "taicpp", "tests", "mplex_native_oracle.py"), ROOT]),
               ("cpp-controller", [cpp_test, "--native"])]
    for label, argv in helpers:
        for answer in ("go", "cancel"):
            env = dict(h.env)
            env["TAI_PLEXER_SOCKET"] = os.path.join(h.home, "plexer", "server.sock")
            env["TAI_PLEXER_PANE"] = "0"
            proc = subprocess.Popen(argv, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
            client.wait_for("%s first dialog" % label, lambda t: "Environment" in t and "staging" in t)
            seen.append((label + " " + answer + " dialog 1", screen_body(client)))
            if answer == "cancel":
                client.send(b"\x1b", 0.5)
            else:
                client.send(DOWN, 0.3)
                seen.append((label + " " + answer + " moved", screen_body(client)))
                client.send(b"\r", 0.6)
                client.wait_for("text prompt", lambda t: "Tag" in t)
                seen.append((label + " " + answer + " prompt", client.text().split("\n")[-1].rstrip()))
                client.send(b"\x15x1\r", 0.6)
                client.wait_for("actions", lambda t: "Deploy" in t)
                seen.append((label + " " + answer + " actions", screen_body(client)))
                client.send(b"\r", 0.6)
            out, _ = proc.communicate(timeout=30)
            seen.append((label + " " + answer + " output", out))
            client.send(b"\x15", 0.3)
    return seen


SCENARIOS = [("lifecycle", sc_lifecycle, None), ("panes_windows", sc_panes_windows, None), ("commands", sc_commands, None),
             ("copy_mode", sc_copy_mode, None), ("config", sc_config, CONFIG_TOML), ("native_flow", sc_native_flow, None), ("errors", sc_errors, None)]


def run_all(label, argv):
    results = {}
    for name, fn, config in SCENARIOS:
        if ONLY and name != ONLY:
            continue
        harness = Harness(label, argv, config)
        try:
            results[name] = fn(harness)
        except Exception as exc:  # noqa: BLE001
            results[name] = [("EXCEPTION", repr(exc))]
        finally:
            harness.close()
    return results


def main():
    python_argv = [sys.executable, ROOT + "/tai.py", "--mplex"]
    oracle = run_all("python", python_argv)
    if PYTHON_ONLY:
        for name, rows in oracle.items():
            print("==", name)
            for label, value in rows:
                print("--", label, "\n", value)
        return 0
    cpp = run_all("cpp", [CPP, "--mplex"])
    bad = 0
    for name in oracle:
        a, b = oracle[name], cpp.get(name, [])
        if a == b:
            print("PASS", name, "(%d observations)" % len(a))
            continue
        bad += 1
        print("FAIL", name)
        for (la, va), (lb, vb) in zip(a, b):
            if (la, va) != (lb, vb):
                print("  step:", la, "/", lb)
                print("  python:", json.dumps(va, default=str)[:1500])
                print("  cpp   :", json.dumps(vb, default=str)[:1500])
        if len(a) != len(b):
            print("  length differs: python %d, cpp %d" % (len(a), len(b)))
    print("%d of %d scenarios identical" % (len(oracle) - bad, len(oracle)))
    return 1 if bad else 0


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