#!/usr/bin/env python3
"""
tea-guard — PreToolUse(Bash) hook for the `tea` plugin.

Enforces, deterministically, the one rule prose cannot: every `tea` command
that touches Gitea runs under the login the OPERATOR pinned — never one Claude
chose. It does this by *resolving and rewriting* the command rather than just
checking it:

  Claude must write:   tea ... --login "$GITEA_LOGIN" ...
  The guard rewrites:  tea ... --login <operator-pinned-login> ...

The pin is read from .claude/settings.local.json (env.GITEA_LOGIN) at call
time — from the FILE, not the environment — so a freshly pinned login works in
the same session with no restart. WHERE that file is looked for is not decided
here: skills/auth/scripts/pin.py holds the search order, and the sync and wiki
scripts resolve the pin through the same module. One order, one copy of it. The
guard and the scripts disagreeing about a directory is a bug by construction,
and was one: in a git worktree `tea` worked and every script said "no login
pinned".

Rules:
  - not a `tea` command ............................. allow (passthrough)
  - tea logins list/ls, tea --version/--help ........ allow (no identity used)
  - no --login / -l ................................. BLOCK
  - --login <literal> or --login "$OTHER_VAR" ....... BLOCK (Claude may not pick)
  - --login "$GITEA_LOGIN", pin found ............... REWRITE to the pin, allow
  - --login "$GITEA_LOGIN", no pin .................. BLOCK (run /tea:auth)

"A `tea` command" means the shell would RUN `tea`, not that the string contains
the word. The guard used to ask the second question — a substring search over
the whole command line — and in a repository whose subject *is* the CLI that is
a different question with the same answer far too often: an issue title, a
commit message, `grep -rn " tea " docs/` and `echo tea` were all blocked, with
a message telling the operator to add `--login` to `git commit`. Worse, the
advice was unfollowable: the only way past the guard was to reword the prose.

So the command is tokenized (heredoc bodies dropped, line continuations
folded, backticks and newlines treated as boundaries) and only words in
*command position* count — the first word, and the first word after `;`, `&&`,
`||`, `|`, `&`, `(`, `)`, `{`, `}`, past any VAR=value assignments and prefix
words like `env`/`sudo`/`xargs`. Quoting is what saves the prose: a title or a
`-m` message is one token, and one token is never a command. Compound commands
stay guarded segment by segment, substitutions included, and every `tea` in the
line is checked — not just the first.

If the line cannot be tokenized at all (unbalanced quotes), the old substring
test decides. That direction fails closed: it over-matches, and over-matching
blocks.

Output protocol: exit 0 + JSON {hookSpecificOutput:{updatedInput,...}} to
rewrite; exit 2 + stderr to block.
"""
import sys, os, re, json, shlex

# The identity layer, reached by the plugin's own layout — the one thing a hook
# may assume about where it lives. Import failure is not fatal on its own: a
# command that is not `tea` still passes through untouched (see main), and only
# a command that needs a login is blocked.
sys.path.append(os.path.abspath(os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    os.pardir, "skills", "auth", "scripts")))
try:
    import pin
except Exception:
    pin = None

PLACEHOLDERS = {"$GITEA_LOGIN", "${GITEA_LOGIN}"}

# Operators after which the next word is a command again.
SEPARATORS = {";", ";;", "&", "&&", "|", "|&", "||", "(", ")", "{", "}"}
# Words that stand in front of a command without being one.
TRANSPARENT = {"env", "command", "exec", "nohup", "time", "sudo", "xargs",
               "if", "then", "else", "elif", "while", "until", "do", "!"}

ASSIGNMENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=")
REDIRECT = re.compile(r"^\d*[<>]+&?\d*-?$")
HEREDOC = re.compile(r"<<-?\s*(['\"]?)([A-Za-z_][A-Za-z0-9_]*)\1")
# A login flag and its value, in the ORIGINAL text — this is what gets
# rewritten, so it works on the raw string rather than on tokens.
LOGIN_FLAG = re.compile(r"(--login|(?<![\w-])-l)(\s+|=)(\S+)")
# The pre-tokenizer test, kept for the one case tokenizing cannot serve.
LOOKS_LIKE_TEA = re.compile(r"(^|[;&|(]|\s)tea(\s|$)")

NO_LOGIN = ('every `tea` command must include --login "$GITEA_LOGIN" '
            '(the guard substitutes the operator-pinned login). '
            'Run /tea:auth if no login is pinned.')


def named_login(raw):
    return ('do not name the login yourself (got `%s`). Write exactly '
            '--login "$GITEA_LOGIN"; the guard replaces it with the login '
            'the operator pinned via /tea:auth. This prevents acting under '
            'the wrong identity.' % raw)


def unquote(value):
    for q in ('"', "'"):
        if len(value) >= 2 and value[0] == q and value[-1] == q:
            return value[1:-1]
    return value


def strip_heredocs(cmd):
    """Drop heredoc bodies. They are data the shell feeds to a command, not
    commands — and a commit message quoting a raw `tea api` call is exactly the
    thing that used to be unwritable."""
    lines, kept, i = cmd.split("\n"), [], 0
    while i < len(lines):
        line = lines[i]
        kept.append(line)
        i += 1
        for m in HEREDOC.finditer(line):
            delim, dash = m.group(2), m.group(0).startswith("<<-")
            while i < len(lines):
                probe = lines[i].strip() if dash else lines[i].rstrip()
                i += 1
                if probe == delim:
                    break
    return "\n".join(kept)


def shell_words(cmd):
    """Tokens, with operators as tokens of their own and quotes honored.

    Backticks and newlines become separators before tokenizing: shlex knows
    neither, and both start a command. Inside quotes that substitution is
    harmless — the token still spans the quotes, and a token is never a
    command."""
    text = strip_heredocs(cmd)
    text = re.sub(r"\\\n", " ", text)
    text = text.replace("`", " ; ").replace("\n", " ; ")
    lex = shlex.shlex(text, posix=True, punctuation_chars=True)
    lex.whitespace_split = True
    return list(lex)


def tea_invocations(words):
    """The argument list of every `tea` the shell would actually run."""
    found, current, expect, skip = [], None, True, False
    for w in words:
        if skip:
            skip = False
            continue
        if REDIRECT.match(w):
            skip = True          # the target of a redirection is not a command
            continue
        if w in SEPARATORS:
            current, expect = None, True
            continue
        if expect:
            if ASSIGNMENT.match(w) or w in TRANSPARENT:
                continue
            expect = False
            if w.rsplit("/", 1)[-1] == "tea":
                current = []
                found.append(current)
            continue
        if current is not None:
            current.append(w)
    return found


def is_meta(args):
    """Login enumeration and `--version`/`--help`: no identity is used, and
    /tea:auth needs `tea logins list` while no pin exists yet."""
    if not args:
        return False
    if args[0] in ("--version", "-v", "--help", "-h", "help"):
        return True
    return args[0] in ("logins", "login") and len(args) > 1 \
        and args[1] in ("list", "ls")


def login_value(args):
    """The login as written, or None if the flag is absent."""
    for i, a in enumerate(args):
        if a in ("--login", "-l"):
            return args[i + 1] if i + 1 < len(args) else ""
        if a.startswith("--login=") or a.startswith("-l="):
            return a.split("=", 1)[1]
    return None


def substitute(cmd, login):
    """Every placeholder login in the line, replaced by the pin. Every one:
    a command may run `tea` twice, and half a rewrite leaves the second call
    with an unset variable and no login at all."""
    def repl(m):
        if unquote(m.group(3)) in PLACEHOLDERS:
            return m.group(1) + m.group(2) + shlex.quote(login)
        return m.group(0)
    return LOGIN_FLAG.sub(repl, cmd)


def block(msg):
    sys.stderr.write("tea-guard: BLOCKED — " + msg + "\n")
    sys.exit(2)


def allow_passthrough():
    # exit 0 with no stdout → tool runs unchanged
    sys.exit(0)


def rewrite(tool_input, new_cmd, note):
    updated = dict(tool_input)
    updated["command"] = new_cmd
    print(json.dumps({
        "hookSpecificOutput": {
            "hookEventName": "PreToolUse",
            "updatedInput": updated,
            "additionalContext": note,
        }
    }))
    sys.exit(0)


def main():
    try:
        payload = json.load(sys.stdin)
    except Exception:
        # Can't parse the hook payload — fail open for non-tea safety, but we
        # can't even read the command, so don't block arbitrary Bash.
        allow_passthrough()

    tool_input = payload.get("tool_input") or {}
    cmd = tool_input.get("command") or ""

    try:
        runs = tea_invocations(shell_words(cmd))
    except ValueError:
        # Unbalanced quotes: what the shell would run is not knowable here.
        # Fall back to the substring test — it over-matches, and over-matching
        # blocks rather than lets an unpinned call through.
        runs = None

    if runs is None:
        if not LOOKS_LIKE_TEA.search(cmd):
            allow_passthrough()
        m = LOGIN_FLAG.search(cmd)
        if not m:
            block(NO_LOGIN)
        if unquote(m.group(3)) not in PLACEHOLDERS:
            block(named_login(m.group(3)))
    else:
        # The word appears but nothing runs it → not our concern. This is the
        # branch that lets prose about the CLI be written at all.
        if not runs:
            allow_passthrough()
        for args in runs:
            if is_meta(args):
                continue
            raw = login_value(args)
            if raw is None:
                block(NO_LOGIN)
            if unquote(raw) not in PLACEHOLDERS:
                block(named_login(raw))
        if all(is_meta(args) for args in runs):
            allow_passthrough()

    if pin is None:
        block('cannot import skills/auth/scripts/pin.py, so the pinned login '
              'cannot be resolved. The plugin tree is incomplete; reinstall it.')

    # The hint is the directory the Bash command will run in; the rest of the
    # order (CLAUDE_PROJECT_DIR first, cwd last, and the worktree branch of the
    # search) is pin.py's, and is the same order the scripts get.
    login, src = pin.find_pin(payload.get("cwd"))
    if not login:
        block('no login is pinned. Run /tea:auth to choose one (writes '
              '.claude/settings.local.json env.GITEA_LOGIN). The guard reads '
              'the file at call time, so it takes effect with no restart.')

    rewrite(tool_input, substitute(cmd, login),
            'tea-guard: resolved --login -> %s (pinned in %s)' % (login, src))


if __name__ == "__main__":
    main()