"""托管 shell 脚本：管道末级不得是提前退出的读取器（grep -q/-m、head、sed q）。

pipefail 下读取器先退出，生产者写满 64 KiB 管道缓冲后收到 SIGPIPE，管道状态报 141，
`if producer | grep -q PAT` 与 `if ! ...` 的判定被静默翻转。改用 here-string、进程替换
或 sed 打印即退出，读取器的结论才不会丢。
"""
import re
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SCRIPTS = (
    "scripts/svc.sh",
    "scripts/task.sh",
    "scripts/e2e.sh",
    "scripts/lint-boundaries.sh",
    "docker/verify-env.sh",
    "tests/load/collaboration/run_pgbench.sh",
)

HEREDOC = re.compile(r"<<(-?)\s*(['\"]?)([A-Za-z_][A-Za-z0-9_]*)\2")
SED_QUIT = re.compile(r"(?:^|[\s;{}'\"/$0-9])[qQ][0-9]*\s*(?:$|[\s;}'\"])")
GREP_ARG_OPTS = set("efABCdDX")
BOUNDARY = set(" \t\n;&|()<>")


class _Frame:
    """一层命令上下文（顶层、$(...)、(...)、<(...)）内正在累积的管道。"""

    def __init__(self, line):
        self.piped, self.stage, self.start, self.cases = False, [], line, []


def pipeline_tails(src):
    """词法扫描 bash 源码，返回每条多级管道末级的 (起始行号, 文本)；跳过注释与 heredoc 正文。"""
    tails, frames, modes, heredocs = [], [_Frame(1)], ["cmd"], []
    line, word, i, n = 1, "", 0, len(src)

    def emit(s):
        for f in frames:
            f.stage.append(s)

    def finish(f):
        if f.piped:
            tails.append((f.start, "".join(f.stage).strip()))
        f.piped, f.stage, f.start = False, [], line

    def case_state():
        cases = frames[-1].cases
        return cases[-1] if cases else None

    def keyword(w):
        cases = frames[-1].cases
        if w == "case":
            cases.append("in")
        elif w == "in" and case_state() == "in":
            cases[-1] = "pattern"
        elif w == "esac" and cases:
            cases.pop()

    def push_cmd(opener):
        emit(opener)
        modes.append("cmd")
        frames.append(_Frame(line))

    while i < n:
        c, nxt = src[i], src[i + 1 : i + 2]
        mode = modes[-1]
        if mode in ("sq", "ansi"):
            if c == "\\" and mode == "ansi":
                emit(src[i : i + 2])
                i += 2
                continue
            emit(c)
            line += c == "\n"
            if c == "'":
                modes.pop()
            i += 1
            continue
        if mode in ("dq", "brace"):
            if c == "\\":
                emit(src[i : i + 2])
                line += nxt == "\n"
                i += 2
            elif c == "$" and nxt == "(":
                push_cmd("$(")
                i += 2
            elif c == "$" and nxt == "{":
                emit("${")
                modes.append("brace")
                i += 2
            elif c == '"':
                emit(c)
                if mode == "dq":
                    modes.pop()
                else:
                    modes.append("dq")
                i += 1
            elif c == "}" and mode == "brace":
                emit(c)
                modes.pop()
                i += 1
            else:
                emit(c)
                line += c == "\n"
                i += 1
            continue

        if c in BOUNDARY and word:
            keyword(word)
            word = ""
        top = frames[-1]
        if c == "\\":
            if nxt == "\n":
                emit(" ")
                line += 1
            else:
                emit(src[i : i + 2])
                word += "\0"
            i += 2
        elif c == "#" and (i == 0 or src[i - 1] in " \t\n;&|("):
            j = src.find("\n", i)
            i = n if j < 0 else j
        elif c == "'" or (c == "$" and nxt == "'"):
            emit(src[i : i + 2] if c == "$" else c)
            modes.append("ansi" if c == "$" else "sq")
            word += "\0"
            i += 2 if c == "$" else 1
        elif c == '"':
            emit(c)
            modes.append("dq")
            word += "\0"
            i += 1
        elif c == "$" and nxt == "{":
            emit("${")
            modes.append("brace")
            word += "\0"
            i += 2
        elif c == "$" and nxt == "(":
            push_cmd("$(")
            word += "\0"
            i += 2
        elif c == "(":
            if case_state() == "pattern":
                emit(c)
            else:
                push_cmd(c)
            i += 1
        elif c == ")":
            if case_state() == "pattern":
                top.cases[-1] = "body"
            elif len(frames) > 1:
                finish(top)
                frames.pop()
                modes.pop()
            emit(c)
            i += 1
        elif c == "|":
            if nxt == "|":
                finish(top)
                emit("||")
                i += 2
            elif case_state() == "pattern":
                emit(c)
                i += 1
            else:
                emit(c)
                top.piped, top.stage, top.start = True, [], line
                i += 2 if nxt == "&" else 1
        elif c == ";":
            if src.startswith(";;", i) or src.startswith(";&", i):
                if case_state() == "body":
                    top.cases[-1] = "pattern"
                k = 3 if src.startswith(";;&", i) else 2
            else:
                k = 1
            finish(top)
            emit(src[i : i + k])
            i += k
        elif c == "&":
            if nxt == "&":
                finish(top)
                emit("&&")
                i += 2
            elif nxt == ">" or (i > 0 and src[i - 1] in "<>"):
                emit(c)
                i += 1
            else:
                finish(top)
                emit(c)
                i += 1
        elif c == "<" and src.startswith("<<", i):
            m = None if src.startswith("<<<", i) else HEREDOC.match(src, i)
            if m:
                heredocs.append((m.group(3), m.group(1) == "-"))
                emit(m.group(0))
                i = m.end()
            else:
                k = 3 if src.startswith("<<<", i) else 2
                emit(src[i : i + k])
                i += k
        elif c == "\n":
            line += 1
            i += 1
            for term, dash in heredocs:
                while i < n:
                    j = src.find("\n", i)
                    body = src[i : n if j < 0 else j]
                    i = n if j < 0 else j + 1
                    line += 1
                    if (body.lstrip("\t") if dash else body) == term:
                        break
            heredocs = []
            if not (top.piped and not "".join(top.stage).strip()):
                finish(top)
            emit("\n")
        else:
            emit(c)
            if c not in BOUNDARY:
                word += c
            i += 1

    if len(frames) != 1 or modes != ["cmd"]:
        raise ValueError(f"unbalanced shell lexing at EOF: modes={modes}")
    finish(frames[0])
    return tails


def _grep_exits_early(words):
    for w in words[1:]:
        w = w.strip("'\"")
        if w == "--":
            break
        if w.startswith("--"):
            if w.split("=")[0] in ("--quiet", "--silent", "--max-count"):
                return True
        elif w.startswith("-") and len(w) > 1:
            for ch in w[1:]:
                if ch in "qm":
                    return True
                if ch in GREP_ARG_OPTS:
                    break
    return False


def early_exit_reader(stage):
    """末级命令读到结论后是否会提前关闭 stdin（grep -q/-m、head、带 q 命令的 sed）。"""
    words = stage.split()
    if not words:
        return False
    cmd = words[0]
    if cmd in ("grep", "egrep", "fgrep"):
        return _grep_exits_early(words)
    if cmd == "head":
        return True
    if cmd == "sed":
        return bool(SED_QUIT.search(stage.split(None, 1)[1] if len(words) > 1 else ""))
    return False


def hazards(src):
    return [(line, text) for line, text in pipeline_tails(src) if early_exit_reader(text)]


class ShellPipeHazardTests(unittest.TestCase):
    def test_detector_flags_synthetic_hazard(self):
        found = hazards("set -euo pipefail\nseq 1 200000 | grep -q 1\n")
        self.assertEqual([line for line, _ in found], [2], found)

    def test_detector_flags_hazard_shapes(self):
        for snippet in (
            'if ! seq 1 200000 | grep -q 1; then echo miss; fi',
            'x="$(sed -n \'s/^k=//p\' "$f" | head -1)"',
            'if ! "$bin" serve -h 2>&1 | grep -Eq -- "(^|[[:space:]])-{1,2}config"; then :; fi',
            'echo "${r:+($(printf \'%s\' "$r" | head -c 120))}"',
            'find . -type f \\\n  | grep -m1 x',
            "seq 1 200000 | sed -n '/1/{p;q;}'",
            'seq 1 200000 | grep -qiE "a|b" || exit 1',
        ):
            with self.subTest(snippet=snippet):
                self.assertTrue(hazards(snippet + "\n"), snippet)

    def test_detector_accepts_safe_forms(self):
        src = "\n".join((
            "set -euo pipefail",
            "# seq 1 200000 | grep -q 1",
            "grep -q 1 < <(seq 1 200000)",
            'grep -q x <<<"$v"',
            "x=\"$(sed -n '/^k=/{s///p;q;}' \"$f\")\"",
            "ps -eo comm | sed -n '1,12p'",
            "case \"$w\" in --auth|--config) grep -q x <<<\"$w\" ;; esac",
            "cat <<'EOF'",
            "seq 1 200000 | grep -q 1",
            "EOF",
            "",
        ))
        self.assertEqual(hazards(src), [])

    def test_managed_scripts_have_no_early_exit_pipe_tail(self):
        for rel in SCRIPTS:
            with self.subTest(script=rel):
                path = ROOT / rel
                self.assertTrue(path.is_file(), f"{rel} missing")
                found = hazards(path.read_text(encoding="utf-8"))
                self.assertEqual(
                    found, [], "\n".join(f"{rel}:{line}: | {text}" for line, text in found)
                )


if __name__ == "__main__":
    unittest.main()