#!/usr/bin/env python3
"""repl_drive.py - 用 pty 驱动 ligent 的 REPL, 跑多轮对话并打印全部输出. 

交互式行为 (多轮上下文、/命令、行编辑)没法用管道测 —— 管道模式会把
整段 stdin 当成一个任务. 所以这里开一个伪终端, 逐条发送提示词. 

用法: 
    repl_drive.py --bin ./ligent --base-url URL --prompt "第一轮" --prompt "第二轮"
"""
import argparse
import os
import pty
import re
import select
import sys
import time


def drive(binary, args, prompts, wait_per_prompt=2.0, quiet_after=1.5,
          raw_keys=None):
    pid, fd = pty.fork()
    if pid == 0:
        os.execv(binary, [binary] + args)
        os._exit(127)

    out = bytearray()
    # 行编辑器每次回到提示符都会打开括号粘贴, 所以这个序列就是"轮到用户输入了"
    # 的信号. 用它做同步点, 比"等输出安静下来"可靠也快得多. 
    PROMPT_READY = b"\x1b[?2004h"
    ready_seen = 0
    tail = b""

    def note(chunk):
        # 标记可能被 read 切断, 所以要连着上一块的尾巴一起找
        nonlocal tail, ready_seen
        buf = tail + chunk
        ready_seen += buf.count(PROMPT_READY)
        tail = buf[-(len(PROMPT_READY) - 1):]

    def drain(timeout, idle=0.5, until_ready=False):
        """等到"回到提示符" (until_ready), 或者输出彻底安静下来. 

        必须按标记等, 不能只按安静等: 等待动画每 250ms 画一帧, pty 永远
        不会安静超过 idle, 于是每次都要白等满 timeout —— 而 mock 其实几
        毫秒就答复了. 这是整个 e2e 从 126 秒掉下来的关键. 
        """
        want = ready_seen + 1 if until_ready else None
        deadline = time.time() + timeout
        last = time.time()
        while time.time() < deadline:
            r, _, _ = select.select([fd], [], [], 0.02)
            if r:
                try:
                    chunk = os.read(fd, 65536)
                except OSError:
                    return False
                if not chunk:
                    return False
                out.extend(chunk)
                note(chunk)
                last = time.time()
                if want is not None and ready_seen >= want:
                    return True
            elif time.time() - last > idle:
                return True
        return True

    drain(quiet_after, until_ready=True)    # 等横幅画完、第一次回到提示符
    if raw_keys is not None:
        # 逐字节发送原始按键序列 (用于测方向键等编辑行为)
        for chunk, pause in raw_keys:
            os.write(fd, chunk if isinstance(chunk, bytes) else chunk.encode("utf-8"))
            drain(pause)
        os.write(fd, b"/exit\n")
        drain(0.8)
        try:
            os.close(fd)
        except OSError:
            pass
        try:
            os.waitpid(pid, 0)
        except ChildProcessError:
            pass
        return re.sub(r"\x1b\[[0-9;]*m", "", out.decode("utf-8", "replace"))

    for p in prompts:
        os.write(fd, (p + "\n").encode("utf-8"))
        if not drain(wait_per_prompt, until_ready=True):
            break
    os.write(fd, b"/exit\n")
    drain(0.8)

    try:
        os.close(fd)
    except OSError:
        pass
    try:
        os.waitpid(pid, 0)
    except ChildProcessError:
        pass

    text = out.decode("utf-8", "replace")
    return re.sub(r"\x1b\[[0-9;]*m", "", text)      # 去掉 ANSI 样式


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--bin", default="./ligent")
    ap.add_argument("--base-url", required=True)
    ap.add_argument("--api-key", default="test")
    ap.add_argument("--prompt", action="append", default=[])
    ap.add_argument("--extra", action="append", default=[])
    ap.add_argument("--wait", type=float, default=2.0)
    ap.add_argument("--keys", action="append", default=[],
                    help="原始按键序列 (支持 \\x1b 转义), 按顺序发送")
    a = ap.parse_args()

    # --no-history: 被驱动的实例不该往用户真实的 ligent.history 里写测试垃圾. 
    # 需要历史的调用方用 --extra=--history PATH (排在后面, 会覆盖这一条). 
    args = ["--api-key", a.api_key, "--base-url", a.base_url, "--no-color",
            "--no-confirm", "--no-auto-save", "--no-history"] + a.extra
    if a.keys:
        raw = [(k.encode("utf-8").decode("unicode_escape").encode("latin-1"), 0.25)
               for k in a.keys]
        sys.stdout.write(drive(a.bin, args, [], a.wait, raw_keys=raw))
    else:
        sys.stdout.write(drive(a.bin, args, a.prompt, a.wait))
    return 0


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