#!/usr/bin/env python3
"""Fail CI when production code launches Claude/Pi outside the lease gate."""

from __future__ import annotations

import argparse
import json
import re
import shlex
import sys
from pathlib import Path
from typing import Final, NamedTuple, Sequence

SCANNED_ROOTS: Final = ("packages", "apps", "plugins", "tools")
SCANNED_SUFFIXES: Final = {
    ".bash",
    ".cjs",
    ".js",
    ".json",
    ".mjs",
    ".py",
    ".sh",
    ".ts",
    ".tsx",
    ".yaml",
    ".yml",
    ".zsh",
}
SKIPPED_DIRECTORIES: Final = {
    ".git",
    ".next",
    ".turbo",
    "coverage",
    "dist",
    "node_modules",
}
DANGEROUS_PRIMITIVE: Final = "--dangerously-" + "skip-permissions"
CHOKE_POINT_SUFFIX: Final = "framework/tools/lease-broker/launch-runtime.py"
SHELL_SUFFIXES: Final = {".bash", ".sh", ".zsh"}

# A launch line is gated only when it CALLS the common wrapper in command
# position, invokes the TypeScript adapter, or constructs/runs a `mosaic`
# runtime command. A marker in a comment, string argument, echo, or unrelated
# variable can never satisfy these invocation-shaped patterns.
GATED_PATTERNS: Final = (
    re.compile(r"(?:^\s*|=>\s*)execLeaseGatedRuntime\s*\("),
    re.compile(
        r"\bexecRuntime\s*\(\s*[\"']python3[\"']\s*,\s*"
        r"\[\s*launcher\s*,.*[\"']--runtime[\"']\s*,\s*runtime\s*,\s*[\"']--[\"']"
    ),
    re.compile(
        r"(?:^|[;&|]\s*|\bexec\s+|\b(?:LAUNCH_COMMAND|launch_cmd)\s*=\s*\(?|\becho\s+[\"'])"
        r"mosaic\s+(?:yolo\s+)?(?:claude|pi|claudex|[\"']?\$\{?runtime\}?[\"']?|[\"']?\$MOSAIC_AGENT_RUNTIME[\"']?)\b"
    ),
    re.compile(r"\[\s*[\"']mosaic[\"']\s*,\s*(?:[\"']yolo[\"']\s*,\s*)?(?:runtime|[\"'](?:claude|pi|claudex)[\"'])"),
)

DIRECT_PATTERNS: Final = (
    # Shell/process command forms, including here-doc command examples that an
    # operator could execute verbatim.
    re.compile(r"^\s*(?:claude|pi)(?:\s|$)"),
    re.compile(r"(?:^|[;&|]\s*|\bexec\s+|\bcommand\s+)(?:claude|pi)\s+(?:-p\b|--dangerously\b|--print\b)"),
    re.compile(r"\bexec\s+(?:claude|pi)(?:\s|$)"),
    re.compile(r"\bexec\s+(?:/[^\s/]+)+/(?:claude|pi)(?:\s|$)"),
    re.compile(r"\bexec\s+[\"']?\$(?:\{?runtime\}?|MOSAIC_AGENT_RUNTIME)\b"),
    # JS/TS and Python process APIs with a literal runtime binary.
    re.compile(
        r"\b(?:spawn|spawnSync|exec|execSync|execFile|execFileSync|execv|execvp|execvpe|Popen|run|call|system|check_call|check_output)\s*\(\s*(?:\[\s*)?[\"'](?:claude|pi)(?:[\"']|\s)"
    ),
    re.compile(
        r"\b(?:spawn|spawnSync|exec|execSync|execFile|execFileSync|execv|execvp|execvpe|Popen|run|call|system|check_call|check_output)\s*\(\s*(?:\[\s*)?[\"'](?:/[^\"'/]+)+/(?:claude|pi)[\"']"
    ),
    # Launch-command arrays and the prior @mosaicstack/coord dynamic default.
    re.compile(r"\b(?:spawn|spawnSync|exec|execSync|execFile|execFileSync)\s*\(\s*runtime\b"),
    re.compile(r"\b(?:command|launchCommand|LAUNCH_COMMAND)\s*=\s*(?:\(|\[)\s*[\"']?(?:claude|pi)\b"),
    re.compile(r"\b(?:command|launchCommand|LAUNCH_COMMAND)\s*=\s*(?:\(|\[)\s*[\"']?\$(?:\{?runtime\}?|MOSAIC_AGENT_RUNTIME)\b"),
    re.compile(r"\breturn\s*\[\s*runtime\s*,\s*[\"'](?:-p|--dangerously)"),
    re.compile(r"\$\(\s*(?:claude|pi)(?:\s|$)"),
    re.compile(r"\beval\s+[\"'](?:claude|pi)(?:\s|[\"'])"),
)
RUNTIME_ASSIGNMENT: Final = re.compile(
    r"^\s*([A-Za-z_][A-Za-z0-9_]*)\s*=\s*[\"']?(?:claude|pi)(?:\s|[\"']|$)"
)
TYPESCRIPT_WRAPPER_INVOCATION: Final = re.compile(
    r"(?m)^\s*execRuntime\s*\(\s*[\"']python3[\"']\s*,\s*"
    r"\[\s*launcher\s*,[^\]]*[\"']--runtime[\"']\s*,\s*runtime\s*,\s*"
    r"[\"']--[\"']\s*,\s*runtime(?:\s*,|\s*\])",
    re.DOTALL,
)


class LaunchSite(NamedTuple):
    path: Path
    line_number: int
    line: str
    classification: str


def is_test_path(path: Path) -> bool:
    name = path.name.lower()
    return (
        "__tests__" in path.parts
        or ".spec." in name
        or ".test." in name
        or name.endswith("_unittest.py")
        or name.startswith("test-")
        or name.startswith("test_")
    )


def strip_comments(path: Path, line: str, in_block_comment: bool = False) -> tuple[str, bool]:
    suffix = path.suffix.lower()
    hash_comments = suffix in {".bash", ".py", ".sh", ".yaml", ".yml", ".zsh"}
    slash_comments = suffix in {".cjs", ".js", ".mjs", ".ts", ".tsx"}
    output: list[str] = []
    quote: str | None = None
    escaped = False
    index = 0

    while index < len(line):
        if in_block_comment:
            end = line.find("*/", index)
            if end < 0:
                return "".join(output), True
            in_block_comment = False
            index = end + 2
            continue

        character = line[index]
        following = line[index + 1] if index + 1 < len(line) else ""
        if quote is not None:
            output.append(character)
            if escaped:
                escaped = False
            elif character == "\\":
                escaped = True
            elif character == quote:
                quote = None
            index += 1
            continue

        if character in {"'", '"', "`"}:
            quote = character
            output.append(character)
            index += 1
            continue
        if slash_comments and character == "/" and following == "/":
            break
        if slash_comments and character == "/" and following == "*":
            in_block_comment = True
            index += 2
            continue
        if hash_comments and character == "#":
            if suffix == ".py" or index == 0 or line[index - 1].isspace():
                break
        output.append(character)
        index += 1

    return "".join(output), in_block_comment


def is_choke_point(path: Path) -> bool:
    return path.as_posix().endswith(CHOKE_POINT_SUFFIX)


def shell_commands(line: str) -> list[list[str]]:
    candidate = line.rstrip().removesuffix("\\").rstrip()
    try:
        lexer = shlex.shlex(candidate, posix=True, punctuation_chars=";&|")
        lexer.whitespace_split = True
        lexer.commenters = ""
        tokens = list(lexer)
    except ValueError:
        return []

    commands: list[list[str]] = []
    current: list[str] = []
    for token in tokens:
        if token and all(character in ";&|" for character in token):
            if current:
                commands.append(current)
                current = []
        else:
            current.append(token)
    if current:
        commands.append(current)
    return commands


def is_wrapper_invocation(path: Path, line: str) -> bool:
    if path.suffix.lower() not in SHELL_SUFFIXES:
        return False
    separator = re.search(r"\s--(?:\s|$)", line)
    if separator is None:
        return False
    # Everything after the wrapper separator is opaque runtime argv and may
    # contain an open quote continued on later physical lines. Parse only the
    # complete command-position prefix through the separator.
    wrapper_prefix = line[: separator.end()]
    for command in shell_commands(wrapper_prefix):
        if command and command[0] == "exec":
            command = command[1:]
        if not command:
            continue
        if command[0] == "python3":
            if len(command) < 2 or not command[1].endswith("launch-runtime.py"):
                continue
            arguments = command[2:]
        elif command[0].endswith(("launch-runtime.py", "launch-runtime.sh")):
            arguments = command[1:]
        else:
            continue
        try:
            runtime_index = arguments.index("--runtime")
            separator_index = arguments.index("--")
        except ValueError:
            continue
        if runtime_index + 1 < len(arguments) and runtime_index < separator_index:
            return True
    return False


def is_gated_line(path: Path, line: str) -> bool:
    return is_wrapper_invocation(path, line) or any(
        pattern.search(line) for pattern in GATED_PATTERNS
    )


def is_shell_direct_invocation(path: Path, line: str) -> bool:
    if path.suffix.lower() not in SHELL_SUFFIXES:
        return False
    assignment = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=")
    case_arm = re.match(r"^\s*[^;&()]+\)\s*", line)
    command_source = line[case_arm.end() :] if case_arm is not None else line
    for command in shell_commands(command_source):
        index = 0
        while index < len(command) and assignment.match(command[index]):
            index += 1
        while index < len(command) and command[index] in {"command", "exec", "nohup"}:
            index += 1
        if index < len(command) and command[index] == "env":
            index += 1
            while index < len(command) and (
                command[index].startswith("-") or assignment.match(command[index])
            ):
                index += 1
            while index < len(command) and command[index] in {"command", "exec", "nohup"}:
                index += 1
        if index < len(command) and Path(command[index]).name in {"claude", "pi"}:
            return True
    return False


def is_direct_line(path: Path, line: str) -> bool:
    return is_shell_direct_invocation(path, line) or any(
        pattern.search(line) for pattern in DIRECT_PATTERNS
    )


def executes_runtime_variable(line: str, runtime_variables: set[str]) -> bool:
    for variable in runtime_variables:
        reference = rf"\$(?:{re.escape(variable)}|\{{{re.escape(variable)}\}})"
        if re.search(rf"\beval\s+[\"']?{reference}", line):
            return True
        if re.match(rf"^\s*[\"']?{reference}[\"']?(?:\s|$)", line):
            return True
    return False


def typescript_wrapper_lines(path: Path, source: str) -> set[int]:
    if path.suffix.lower() not in {".js", ".mjs", ".ts", ".tsx"}:
        return set()
    return {
        source.count("\n", 0, match.start()) + 1
        for match in TYPESCRIPT_WRAPPER_INVOCATION.finditer(source)
    }


def classify_text(path: Path, source: str) -> list[LaunchSite]:
    sites: list[LaunchSite] = []
    validated_typescript_wrappers = typescript_wrapper_lines(path, source)
    runtime_variables: set[str] = set()
    gated_continuation = False
    in_block_comment = False
    for line_number, physical_line in enumerate(source.splitlines(), start=1):
        line, in_block_comment = strip_comments(path, physical_line, in_block_comment)
        stripped = line.strip()
        if not stripped:
            gated_continuation = False
            continue

        assignment = RUNTIME_ASSIGNMENT.match(line)
        if assignment is not None:
            runtime_variables.add(assignment.group(1))
        primitive_violation = DANGEROUS_PRIMITIVE in line and not is_choke_point(path)
        line_is_direct = is_direct_line(path, line) or executes_runtime_variable(
            line, runtime_variables
        )
        line_is_gated = line_number in validated_typescript_wrappers or is_gated_line(path, line)

        if primitive_violation:
            sites.append(
                LaunchSite(path, line_number, physical_line.rstrip(), "dangerous-primitive")
            )
        elif line_is_direct and not gated_continuation:
            # Direct syntax always wins over a same-line marker. Only the command
            # continuation of a previously validated wrapper may contain the raw
            # runtime binary itself.
            sites.append(LaunchSite(path, line_number, physical_line.rstrip(), "direct"))
        elif line_is_gated:
            sites.append(LaunchSite(path, line_number, physical_line.rstrip(), "gated"))

        gated = line_is_gated or gated_continuation
        gated_continuation = gated and line.rstrip().endswith("\\")
    return sites


def scan_text(path: Path, source: str) -> list[LaunchSite]:
    return [site for site in classify_text(path, source) if site.classification != "gated"]


def source_files(root: Path):
    for relative_root in SCANNED_ROOTS:
        search_root = root / relative_root
        if not search_root.is_dir():
            continue
        for path in search_root.rglob("*"):
            if not path.is_file() or path.suffix.lower() not in SCANNED_SUFFIXES:
                continue
            if any(part in SKIPPED_DIRECTORIES for part in path.parts):
                continue
            if is_test_path(path):
                continue
            yield path


def scan_repository(root: Path) -> list[LaunchSite]:
    violations: list[LaunchSite] = []
    for path in source_files(root):
        try:
            source = path.read_text(encoding="utf-8")
        except UnicodeDecodeError:
            violations.append(LaunchSite(path, 0, "non-UTF-8 source", "unscannable"))
            continue
        violations.extend(scan_text(path.relative_to(root), source))
    return violations


def inventory_repository(root: Path) -> list[LaunchSite]:
    inventory: list[LaunchSite] = []
    for path in source_files(root):
        try:
            source = path.read_text(encoding="utf-8")
        except UnicodeDecodeError:
            inventory.append(LaunchSite(path.relative_to(root), 0, "non-UTF-8 source", "unscannable"))
            continue
        relative_path = path.relative_to(root)
        inventory.extend(classify_text(relative_path, source))
    return inventory


def format_violation(violation: LaunchSite) -> str:
    return f"{violation.path}:{violation.line_number}: {violation.classification}: {violation.line.strip()}"


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, default=Path.cwd())
    parser.add_argument("--json", action="store_true")
    arguments = parser.parse_args(argv)
    root = arguments.root.resolve()
    inventory = inventory_repository(root)
    violations = [site for site in inventory if site.classification != "gated"]

    if arguments.json:
        print(
            json.dumps(
                {
                    "gated": sum(site.classification == "gated" for site in inventory),
                    "total": len(inventory),
                    "sites": [
                        {
                            "path": str(site.path),
                            "line": site.line_number,
                            "classification": site.classification,
                            "source": site.line.strip(),
                        }
                        for site in inventory
                    ],
                },
                sort_keys=True,
            )
        )
    else:
        for site in inventory:
            print(format_violation(site))
        print(
            f"runtime launch inventory: "
            f"{len(inventory) - len(violations)} gated/{len(inventory)} total"
        )

    if violations:
        print("ungated consequential runtime launch detected", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())