#!/usr/bin/env python3
"""Fail if ALL_CAPS names are loaded without being defined in scope (refactor footgun).

Catches the class of bug where ``agent_loop_run`` referenced constants that were never
imported after a split (runtime ``NameError`` inside nested functions — py_compile alone
does not catch those).

This is not a full type checker; it models module + function (+ class) scopes, ``global``,
and comprehensions. Skips names that look like ``Enum`` members on class attributes.

Usage:
    python3 scripts/check_screaming_snake_loads.py [file.py ...]

Default file list: repo-root ``agent_*.py`` plus ``subagent_runner.py``.
Set ``MUGI_CHECK_SCREAMING_SNAKE_BACKEND=1`` to also include ``mugi_backend.py``.

Exit 0 = ok, 1 = undefined loads found, 2 = script error.
"""
from __future__ import annotations

import ast
import builtins
import os
import re
import sys
from typing import Iterable

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

_SCREAM = re.compile(r"^[A-Z][A-Z0-9_]+$")
_BUILTINS = set(vars(builtins))


def _default_targets() -> list[str]:
    """Orchestration entrypoints where a missed import surfaces as runtime NameError."""
    roots: list[str] = []
    for name in sorted(os.listdir(REPO_ROOT)):
        if not name.endswith(".py"):
            continue
        if name.startswith("test_") or name == "conftest.py":
            continue
        if name.startswith("agent_") or name == "subagent_runner.py":
            roots.append(os.path.join(REPO_ROOT, name))
    # mugi_backend.py uses a huge top-level try-import block; include only when opted-in via env.
    mb = os.path.join(REPO_ROOT, "mugi_backend.py")
    if os.getenv("MUGI_CHECK_SCREAMING_SNAKE_BACKEND", "").strip() == "1":
        roots.append(mb)
    return roots


class ScreamChecker(ast.NodeVisitor):
    """Track local scopes; flag ALL_CAPS Name loads with no binding in scope chain."""

    def __init__(self, path: str) -> None:
        self.path = path
        self.errors: list[tuple[int, str]] = []
        # Stack of local scopes (each: set of names bound in that block).
        # scopes[0] is always the module scope.
        self._scopes: list[set[str]] = [set()]
        # Per-function: names declared ``global`` (resolve against scopes[0] only).
        self._func_globals: list[set[str]] = []

    # -- scope helpers -------------------------------------------------

    def _define(self, names: Iterable[str]) -> None:
        for n in names:
            if n:
                self._scopes[-1].add(n)

    def _push(self) -> None:
        self._scopes.append(set())

    def _pop(self) -> None:
        self._scopes.pop()

    def _resolve(self, name: str) -> bool:
        if name in _BUILTINS:
            return True
        if self._func_globals and name in self._func_globals[-1]:
            return name in self._scopes[0]
        for scope in reversed(self._scopes):
            if name in scope:
                return True
        return False

    # -- visitors ------------------------------------------------------

    def visit_Module(self, node: ast.Module) -> None:
        for stmt in node.body:
            self._visit_toplevel(stmt)

    def _visit_toplevel(self, node: ast.stmt) -> None:
        if isinstance(node, ast.Import):
            for alias in node.names:
                self._define([alias.asname or alias.name.split(".")[0]])
        elif isinstance(node, ast.ImportFrom):
            for alias in node.names:
                if alias.name == "*":
                    continue
                self._define([alias.asname or alias.name])
        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            self._define([node.name])
            self._visit_function(node, is_async=isinstance(node, ast.AsyncFunctionDef))
        elif isinstance(node, ast.ClassDef):
            self._define([node.name])
            self._visit_class(node)
        elif isinstance(node, ast.Assign):
            for t in node.targets:
                self._store_target(t)
            if node.value:
                self.visit(node.value)
        elif isinstance(node, ast.AnnAssign):
            if node.target:
                self._store_target(node.target)
            if node.value:
                self.visit(node.value)
        elif isinstance(node, ast.TypeAlias):
            self._define([node.name.id])  # py3.12+
        elif isinstance(node, ast.Try):
            for s in node.body + node.orelse + node.finalbody:
                self._visit_toplevel(s)
            for handler in node.handlers:
                for s in handler.body:
                    self._visit_toplevel(s)
        elif isinstance(node, ast.If):
            if node.test:
                self.visit(node.test)
            for s in node.body + node.orelse:
                self._visit_toplevel(s)
        elif isinstance(node, ast.While):
            if node.test:
                self.visit(node.test)
            for s in node.body + node.orelse:
                self._visit_toplevel(s)
        elif isinstance(node, (ast.With, ast.AsyncWith)):
            for item in node.items:
                self.visit(item.context_expr)
            for s in node.body:
                self._visit_toplevel(s)
        else:
            self.generic_visit(node)

    def _store_target(self, target: ast.expr) -> None:
        if isinstance(target, ast.Name):
            self._define([target.id])
        elif isinstance(target, (ast.Tuple, ast.List)):
            for elt in target.elts:
                self._store_target(elt)

    def _visit_class(self, node: ast.ClassDef) -> None:
        self._push()
        for stmt in node.body:
            if isinstance(stmt, ast.Assign):
                for t in stmt.targets:
                    self._store_target(t)
            elif isinstance(stmt, ast.AnnAssign) and stmt.target:
                self._store_target(stmt.target)
            elif isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)):
                self._define([stmt.name])
                self._visit_function(stmt, is_async=isinstance(stmt, ast.AsyncFunctionDef))
            elif isinstance(stmt, ast.ClassDef):
                self._define([stmt.name])
                self._visit_class(stmt)
        self._visit_body_exprs(node.body)
        self._pop()

    def _visit_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef, *, is_async: bool) -> None:
        args = node.args
        params: list[str] = []
        for a in args.posonlyargs + args.args + ([args.vararg] if args.vararg else []) + args.kwonlyargs:
            if a:
                params.append(a.arg)
        if args.kwarg and args.kwarg.arg:
            params.append(args.kwarg.arg)

        glob: set[str] = set()
        for stmt in node.body:
            if isinstance(stmt, ast.Global):
                glob.update(stmt.names)

        self._func_globals.append(glob)
        self._push()
        self._define(params)
        for stmt in node.body:
            self._visit_stmt(stmt)
        self._pop()
        self._func_globals.pop()

    def _visit_stmt(self, node: ast.stmt) -> None:
        if isinstance(node, ast.Global):
            return
        if isinstance(node, ast.Nonlocal):
            return
        if isinstance(node, ast.Import):
            for alias in node.names:
                self._define([alias.asname or alias.name.split(".")[0]])
            return
        if isinstance(node, ast.ImportFrom):
            for alias in node.names:
                if alias.name == "*":
                    continue
                self._define([alias.asname or alias.name])
            return
        if isinstance(node, ast.Assign):
            for t in node.targets:
                self._store_target(t)
            self.visit(node.value)
            return
        if isinstance(node, ast.AnnAssign):
            if node.target:
                self._store_target(node.target)
            if node.value:
                self.visit(node.value)
            return
        if isinstance(node, ast.AugAssign):
            self._store_target(node.target)
            self.visit(node.value)
            return
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            self._define([node.name])
            self._visit_function(node, is_async=isinstance(node, ast.AsyncFunctionDef))
            return
        if isinstance(node, ast.ClassDef):
            self._define([node.name])
            self._visit_class(node)
            return
        if isinstance(node, ast.For):
            self._visit_pattern_or_expr(node.target)
            self.visit(node.iter)
            for s in node.body + node.orelse:
                self._visit_stmt(s)
            return
        if isinstance(node, ast.AsyncFor):
            self._visit_pattern_or_expr(node.target)
            self.visit(node.iter)
            for s in node.body + node.orelse:
                self._visit_stmt(s)
            return
        if isinstance(node, ast.While):
            self.visit(node.test)
            for s in node.body + node.orelse:
                self._visit_stmt(s)
            return
        if isinstance(node, ast.If):
            self.visit(node.test)
            for s in node.body + node.orelse:
                self._visit_stmt(s)
            return
        if isinstance(node, ast.With):
            for item in node.items:
                if item.optional_vars:
                    self._visit_pattern_or_expr(item.optional_vars)
                self.visit(item.context_expr)
            for s in node.body:
                self._visit_stmt(s)
            return
        if isinstance(node, ast.AsyncWith):
            for item in node.items:
                if item.optional_vars:
                    self._visit_pattern_or_expr(item.optional_vars)
                self.visit(item.context_expr)
            for s in node.body:
                self._visit_stmt(s)
            return
        if isinstance(node, ast.Try):
            for s in node.body:
                self._visit_stmt(s)
            for handler in node.handlers:
                if handler.name:
                    self._push()
                    self._define([handler.name])
                    for s in handler.body:
                        self._visit_stmt(s)
                    self._pop()
                else:
                    for s in handler.body:
                        self._visit_stmt(s)
            for s in node.orelse:
                self._visit_stmt(s)
            for s in node.finalbody:
                self._visit_stmt(s)
            return
        if isinstance(node, ast.Match):  # py3.10+
            self.visit(node.subject)
            for case in node.cases:
                self._push()
                self._visit_pattern(case.pattern)
                if case.guard:
                    self.visit(case.guard)
                for s in case.body:
                    self._visit_stmt(s)
                self._pop()
            return
        self.generic_visit(node)

    def _visit_body_exprs(self, body: list[ast.stmt]) -> None:
        for stmt in body:
            self._visit_stmt(stmt)

    def _visit_pattern_or_expr(self, node: ast.expr | ast.pattern) -> None:
        if isinstance(node, ast.Name):
            self._define([node.id])
        elif isinstance(node, (ast.Tuple, ast.List)):
            for elt in node.elts:
                self._visit_pattern_or_expr(elt)
        elif hasattr(ast, "MatchAs") and isinstance(node, ast.MatchAs):  # type: ignore[attr-defined]
            if node.name:
                self._define([node.name])
            if node.pattern:
                self._visit_pattern_or_expr(node.pattern)
        elif hasattr(ast, "MatchMapping") and isinstance(node, ast.MatchMapping):  # type: ignore[attr-defined]
            return
        elif hasattr(ast, "MatchSequence") and isinstance(node, ast.MatchSequence):  # type: ignore[attr-defined]
            for p in node.patterns:
                self._visit_pattern_or_expr(p)

    def _visit_pattern(self, node: ast.pattern) -> None:
        self._visit_pattern_or_expr(node)  # best-effort

    def visit_Name(self, node: ast.Name) -> None:
        if isinstance(node.ctx, ast.Load):
            if _SCREAM.match(node.id) and len(node.id) >= 4:
                if not self._resolve(node.id):
                    self.errors.append((node.lineno or 0, node.id))
        self.generic_visit(node)

    def visit_Attribute(self, node: ast.Attribute) -> None:
        self.visit(node.value)


def check_file(path: str) -> list[tuple[int, str]]:
    with open(path, "r", encoding="utf-8") as f:
        src = f.read()
    tree = ast.parse(src, filename=path)
    chk = ScreamChecker(path)
    chk.visit(tree)
    return chk.errors


def main(argv: list[str]) -> int:
    paths = argv[1:] or _default_targets()
    bad: list[tuple[str, int, str]] = []
    for p in paths:
        if not os.path.isfile(p):
            print(f"[check_screaming_snake_loads] skip missing: {p}", file=sys.stderr)
            continue
        for line, name in check_file(p):
            bad.append((p, line, name))
    if not bad:
        print(f"[check_screaming_snake_loads] ok ({len(paths)} file(s))")
        return 0
    print("[check_screaming_snake_loads] undefined ALL_CAPS loads:", file=sys.stderr)
    for p, line, name in sorted(bad):
        rel = os.path.relpath(p, REPO_ROOT)
        print(f"  {rel}:{line}: {name}", file=sys.stderr)
    return 1


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