#!/usr/bin/env python3
"""confirm_check.py - 确认提示上的每一种回答都必须被正确处理

重点是 Ctrl-C。

`confirm_tool` 用 fgets 读输入，而 SIGINT 会让 fgets 以 EINTR 失败、返回
NULL；代码把 NULL 当成"没人能回答，按默认值放行"。于是**在确认提示
上按 Ctrl-C 等于批准这次操作**。run_command 侥幸自己查了中断标志才没执行，
write_file / edit_file 没查，文件照写。

这个脚本把三种回答都跑一遍，直接看磁盘上有没有留下东西：
    Ctrl-C  -> 必须拒绝
    n       -> 必须拒绝
    y       -> 必须执行（否则就是"修过头"了）

用法：
    confirm_check.py <ligent 二进制> <mock base-url>
"""
import os
import pty
import select
import sys
import time

# 探针文件放在 build/ 下，别弄脏仓库根目录
# WRITE_TEST 写这一个；WRITE2_TEST 写 PROBE1 / PROBE2（一轮里两次）
PROBE = "build/tmp/confirm_probe.txt"
PROBE1 = "build/tmp/confirm_probe1.txt"
PROBE2 = "build/tmp/confirm_probe2.txt"
# 只认稳定前缀：选项列表是会变的（现在是 Y/N/Always/neVer），
# 写死整串会在改提示文案时把测试弄坏。
MARK = "继续？"
PROMPT_READY = b"\x1b[?2004h"


def run_case(binary, base, answer, wait_confirm=6.0):
    """跑一轮：触发 write_file → 等确认提示出现 → 回答 → 看文件在不在。"""
    try:
        os.unlink(PROBE)
    except OSError:
        pass

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

    out = bytearray()

    def pump(seconds, idle=0.3):
        """读到"安静"或超时为止"""
        end = time.time() + seconds
        last = time.time()
        while time.time() < end:
            r, _, _ = select.select([fd], [], [], 0.02)
            if r:
                try:
                    d = os.read(fd, 65536)
                except OSError:
                    return False
                if not d:
                    return False
                out.extend(d)
                last = time.time()
            elif time.time() - last > idle:
                return True
        return True

    def wait_for(text, seconds):
        """一直读到 text 出现（或超时）。按标记等，不靠 sleep 猜。"""
        end = time.time() + seconds
        while time.time() < end:
            if text.encode("utf-8") in bytes(out):
                return True
            r, _, _ = select.select([fd], [], [], 0.05)
            if r:
                try:
                    d = os.read(fd, 65536)
                except OSError:
                    return False
                if not d:
                    return False
                out.extend(d)
        return text.encode("utf-8") in bytes(out)

    pump(1.2)                                  # 横幅 + 第一个提示符
    os.write(fd, "WRITE_TEST 写个文件\n".encode())
    got_prompt = wait_for(MARK, wait_confirm)
    if got_prompt:
        os.write(fd, answer)
    pump(1.5)
    os.write(fd, b"\x03")
    pump(0.3)
    try:
        os.close(fd)
    except OSError:
        pass
    try:
        os.waitpid(pid, 0)
    except OSError:
        pass
    return got_prompt, os.path.exists(PROBE)


def run_turn_scope(binary, base):
    """Always / neVer 的作用域必须是"这一轮"，不是整个会话。

    一轮里让模型连写两个文件：第一次问、答 Always，第二次不该再问。
    然后开新的一轮（下一条用户消息），必须**重新问** —— 会话级的开关是
    /yolo，不是这个。

    全程按标记同步（"继续？"出现 / 回到提示符），不靠 sleep 猜时间。
    """
    for p in (PROBE1, PROBE2):
        try:
            os.unlink(p)
        except OSError:
            pass

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

    out = bytearray()
    state = {"ready": 0, "tail": b""}
    READY = b"\x1b[?2004h"     # 行编辑器每次回到提示符都会发它

    def note(chunk):
        buf = state["tail"] + chunk
        state["ready"] += buf.count(READY)
        state["tail"] = buf[-(len(READY) - 1):]

    def pump(seconds, idle=0.3):
        end = time.time() + seconds
        last = time.time()
        while time.time() < end:
            r, _, _ = select.select([fd], [], [], 0.02)
            if r:
                try:
                    d = os.read(fd, 65536)
                except OSError:
                    return
                if not d:
                    return
                out.extend(d)
                note(d)
                last = time.time()
            elif time.time() - last > idle:
                return

    def wait_for(pred, seconds):
        end = time.time() + seconds
        while time.time() < end:
            if pred():
                return True
            r, _, _ = select.select([fd], [], [], 0.03)
            if r:
                try:
                    d = os.read(fd, 65536)
                except OSError:
                    return False
                if not d:
                    return False
                out.extend(d)
                note(d)
        return pred()

    def has_prompt():
        return "继续？".encode("utf-8") in bytes(out)

    def prompt_count():
        return bytes(out).count("继续？".encode("utf-8"))

    pump(1.2)

    # ---- 第一轮：连写两个 ----
    base_ready = state["ready"]
    os.write(fd, "WRITE2_TEST 写两个文件\n".encode())
    wait_for(has_prompt, 8.0)
    os.write(fd, b"Always\n")
    # 等这一轮彻底跑完（回到提示符），别用固定时长猜
    wait_for(lambda: state["ready"] > base_ready, 10.0)
    pump(0.3)
    first_turn_prompts = prompt_count()
    both_written = os.path.exists(PROBE1) and os.path.exists(PROBE2)


    # ---- 第二轮：换一条用户消息，必须重新问 ----
    base_ready = state["ready"]
    before = prompt_count()
    os.write(fd, "WRITE_LAST_TEST 再来一轮\n".encode())
    got_again = wait_for(lambda: prompt_count() > before, 8.0)
    second_prompts = prompt_count() - before
    os.write(fd, b"N\n")
    pump(1.0)

    os.write(fd, b"\x03")
    pump(0.3)
    try:
        os.close(fd)
    except OSError:
        pass
    try:
        os.waitpid(pid, 0)
    except OSError:
        pass

    for p in (PROBE1, PROBE2):
        try:
            os.unlink(p)
        except OSError:
            pass
    return first_turn_prompts, both_written, got_again, second_prompts


def main():
    if len(sys.argv) < 3:
        print("用法: confirm_check.py <二进制> <base-url>")
        return 2
    binary, base = sys.argv[1], sys.argv[2]
    os.makedirs(os.path.dirname(PROBE), exist_ok=True)

    fails = 0
    checks = 0

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

    # 三种回答：(说明, 发送的字节, 期望文件是否被写)
    cases = [
        ("在确认提示上按 Ctrl-C", b"\x03", False),
        ("答 n", b"n\n", False),
        ("答 y", b"y\n", True),
    ]

    for label, answer, expect_written in cases:
        got, written = run_case(binary, base, answer)
        check("%s：等到了确认提示" % label, got, "没等到确认提示")
        check("%s：文件%s" % (label, "被写入" if expect_written else "没被写入"),
              written == expect_written,
              "实际 %s" % ("写了" if written else "没写"))

    # 作用域：Always 只活这一轮
    n1, both, again, n2 = run_turn_scope(binary, base)
    check("Always：同一轮里只问一次", n1 == 1, "实际问了 %d 次" % n1)
    check("Always：同轮的第二次改动被放行", both, "第二个文件没写")
    check("新一轮会重新询问（作用域仅限本轮）", again and n2 >= 1,
          "新一轮没有询问")

    for p in (PROBE, PROBE1, PROBE2):
        try:
            os.unlink(p)
        except OSError:
            pass
    # 这里**不要** rmtree("build/tmp")：那是 e2e 各段共用的临时目录，删掉它
    # 会连带清掉别人放进去的东西（我就因此丢过两次日志），而它自己的三个文件
    # 上面已经删干净了。测试辅助脚本只该收拾自己造的东西。

    print("\n%d 项检查，%d 项失败" % (checks, fails))
    return 1 if fails else 0


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