#!/usr/bin/env python3
"""shell_screen_check.py - 验证 /shell 记录里对全屏程序的处理

交互式 shell 的原始输出里混着转义序列、重绘空格和全屏程序的整屏画面. 
不清洗的话模型看到的就是几百个空格加一串光标定位; 清过头又会把命令和
状态行一起吃掉. 这里用 printf 造出确定的输入来钉住行为：

  1. 备用屏幕 (vim/less/top 走这条路)——画面必须省略, 只留首尾标记; 
  2. 成片空格 (nano/top 靠空格重绘)——必须压掉, 但短缩进要原样保留; 
  3. 大输出里夹着的备用屏幕——跨 read 也要能认出来. 

用法：
    shell_screen_check.py <ligent 二进制> <mock base-url>
"""
import glob
import json
import os
import pty
import re
import select
import shutil
import sys
import time

ALT_TEXT = "备用屏幕里的机密内容"
NORMAL_TEXT = "正常输出"


def drive(binary, base, commands, wait=1.6, term="xterm-256color"):
    """在 pty 里跑 ligent, 进 /shell 执行 commands, 导出后返回记录文本. 

    term=None 表示不设 TERM, 用来观察退化行为. """
    saved = os.environ.copy()
    os.environ["LIGENT_API_KEY"] = "sk-test"
    # 没有 TERM 的话 less/vim 会退回非全屏模式, 测不到备用屏幕那条路. 
    if term is None:
        os.environ.pop("TERM", None)
    else:
        os.environ["TERM"] = term

    # 每轮都清空归档：归档名是随机的, 按文件名排序取"最后一个"会拿到上一轮. 
    bindir = os.path.dirname(os.path.abspath(binary))
    shutil.rmtree(os.path.join(bindir, "ligent.archive"), ignore_errors=True)

    pid, fd = pty.fork()
    if pid == 0:
        os.execv(binary, [binary, "--base-url", base, "--no-confirm", "--auto-save"])
        os._exit(127)

    # 行编辑器每次回到提示符都会打开括号粘贴 —— 用它当"轮到输入了"的信号. 
    # 只按"输出安静"等是不行的：等待动画每 250ms 画一帧, pty 永远不会安静
    # 超过 idle, 于是每次都要白等满整个额度. 
    PROMPT_READY = b"\x1b[?2004h"
    ready_seen = 0
    tail = b""

    def note(chunk):
        nonlocal tail, ready_seen
        buf = tail + chunk
        ready_seen += buf.count(PROMPT_READY)
        tail = buf[-(len(PROMPT_READY) - 1):]

    def drain(seconds, idle=0.4, until_ready=False):
        want = ready_seen + 1 if until_ready else None
        end = time.time() + seconds
        last = time.time()
        while time.time() < end:
            r, _, _ = select.select([fd], [], [], 0.02)
            if r:
                try:
                    chunk = os.read(fd, 65536)
                except OSError:
                    return
                if not chunk:
                    return
                note(chunk)
                last = time.time()
                if want is not None and ready_seen >= want:
                    return
            elif time.time() - last > idle:
                return

    try:
        drain(1.5, until_ready=True)   # 等横幅画完、回到提示符
        for line in ["/shell"] + commands + ["exit", "/export", "/exit"]:
            os.write(fd, line.encode("utf-8") + b"\n")
            drain(wait, until_ready=True)
    finally:
        for fn in (lambda: os.close(fd), lambda: os.waitpid(pid, 0)):
            try:
                fn()
            except OSError:
                pass
        os.environ.clear()
        os.environ.update(saved)

    bindir = os.path.dirname(os.path.abspath(binary))
    files = sorted(glob.glob(os.path.join(bindir, "ligent.archive", "*.json")))
    if not files:
        return None
    with open(files[0], encoding="utf-8") as fh:
        data = json.load(fh)
    users = [m["content"] for m in data["messages"] if m["role"] == "user"]
    return users[-1] if users else ""


def between_markers(rec, start="[全屏程序运行中, 画面内容已省略]", end="[全屏程序结束]"):
    """取出"全屏程序"首尾标记之间的内容. 

    命令回显必然把输入里的字面文本也录进去, 所以不能直接断言"机密文本不在
    记录里"——只能断言它没出现在真正被省略的那段区间内. """
    parts = rec.split(start)
    if len(parts) < 2:
        return None
    return parts[1].split(end)[0]


def main():
    if len(sys.argv) < 3:
        print("用法: shell_screen_check.py <二进制> <base-url>")
        return 2
    binary, base = sys.argv[1], sys.argv[2]
    bindir = os.path.dirname(os.path.abspath(binary))
    archive = os.path.join(bindir, "ligent.archive")

    fails = 0
    total = 0

    def check(name, ok, detail=""):
        nonlocal fails, total
        total += 1
        if ok:
            print("  ✓ %s" % name)
        else:
            fails += 1
            print("  ✗ %s%s" % (name, ("　" + detail) if detail else ""))

    # ---- 1. 备用屏幕：画面省略, 只留首尾标记 ----
    print("备用屏幕 (vim/less 路径)")
    rec = drive(binary, base, [
        r"printf '\033[?1049h" + ALT_TEXT + r"\033[?1049l'",
        "echo " + NORMAL_TEXT,
    ])
    if rec is None:
        check("拿到记录", False, "没有生成归档")
        return 1
    check("拿到记录", True)
    body = between_markers(rec)
    check("首尾标记成对出现", body is not None, "没有找到完整的一对标记")
    check("全屏区间内没有画面内容", body is not None and ALT_TEXT not in body,
          "机密内容漏出来了")
    check("留下开始标记", "全屏程序运行中" in rec)
    check("留下结束标记", "全屏程序结束" in rec)
    check("标记之外的输出还在", NORMAL_TEXT in rec)
    check("记录里没有裸转义符", "\x1b" not in rec)

    # ---- 2. 成片空格：压掉, 但缩进保留 ----
    print("成片空格 (nano/top 重绘路径)")
    rec = drive(binary, base, [
        r"printf 'A'; printf '%0.s ' $(seq 1 200); printf 'B\n'",
        r"printf '    缩进四个空格\n'",
    ])
    if rec is None:
        check("拿到记录", False, "没有生成归档")
        return 1
    check("拿到记录", True)
    check("200 个空格被压缩", not re.search(r" {9,}", rec), "还留着成片空白")
    check("压缩后仍保留 A 和 B", "A" in rec and "B" in rec)
    check("短缩进原样保留", "    缩进四个空格" in rec, "缩进被误伤")
    blank = sum(1 for ch in rec if ch in " \n\t") / max(1, len(rec))
    check("空白占比低于 40%", blank < 0.40, "实际 %.0f%%" % (100 * blank))

    # ---- 3. 大输出里夹着的备用屏幕 ----
    print("大输出 (跨 read)里的备用屏幕")
    rec = drive(binary, base, [
        "head -c 8000 /dev/zero | tr '\\0' 'x'; printf '\\033[?1049h"
        + ALT_TEXT + "\\033[?1049l'",
    ])
    if rec is None:
        check("拿到记录", False, "没有生成归档")
        return 1
    check("拿到记录", True)
    body = between_markers(rec)
    check("跨块后仍认得出备用屏幕", body is not None and ALT_TEXT not in body,
          "切断后漏了画面")
    check("跨块标记成对出现", body is not None)

    # ---- 4. 同一个会话里进两次全屏, 标记必须成对 ----
    print("连续两次进入全屏")
    rec = drive(binary, base, [
        r"printf '\033[?1049hAAA\033[?1049l'; printf '\033[?1049hBBB\033[?1049l'",
    ])
    if rec is None:
        check("拿到记录", False, "没有生成归档")
        return 1
    check("拿到记录", True)
    check("两次开始标记", rec.count("[全屏程序运行中, 画面内容已省略]") == 2,
          "实际 %d 个" % rec.count("[全屏程序运行中, 画面内容已省略]"))
    check("两次结束标记", rec.count("[全屏程序结束]") == 2,
          "实际 %d 个" % rec.count("[全屏程序结束]"))

    # ---- 5. TERM 退化时补一个像样的值 ----
    # TERM=dumb 时 nano 会放弃备用屏幕、改成逐行重绘, 整段忽略就失效了. 
    # 而 ligent 自己的 REPL 就在用光标定位, 说明终端并非真的 dumb. 
    for bad in ("dumb", None):
        label = "TERM=%s" % (bad if bad else "(未设置)")
        print("退化的 %s" % label)
        rec = drive(binary, base, ["echo TERM_IS=$TERM"], term=bad)
        if rec is None:
            check("拿到记录", False, "没有生成归档")
            continue
        check("%s 时补上可用 TERM" % label, "TERM_IS=xterm-256color" in rec,
              "实际记录：%s" % rec.strip().replace("\n", " | ")[-80:])

    shutil.rmtree(archive, ignore_errors=True)
    print("\n%d 项检查, %d 项失败" % (total, fails))
    return 1 if fails else 0


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