"""Static gate for candidate ModelNew sources (first layer of the kernel evaluator's anti-hack defences).

The protected evaluator also enforces everything at runtime (sandbox, fresh inputs, separate reference process,
guarded timing, identity checks, profiler check, build-source scan); this gate rejects known cheat shapes early
and gives the agent a clear reason. It is deliberately strict: KernelBench solutions need torch, triton, math and
cpp_extension.load_inline, nothing else.

check_source(src) -> {"ok": bool, "errors": [str], "kernels": [names], "backend": "cuda"|"triton"|"mixed"|None}
scan_native(text) -> [banned tokens found] (used on build dirs after compilation too)
"""
from __future__ import annotations

import ast
import re

ALLOWED_IMPORTS = {"torch", "triton", "math", "numpy", "typing", "functools", "itertools", "operator",
                   "collections", "dataclasses", "enum", "__future__"}
BANNED_SUBMODULES = ("torch._", "torch.backends", "torch.compiler", "torch.cuda.graphs", "torch.library",
                     "torch.hub", "torch.distributed", "torch.multiprocessing", "torch.profiler", "torch.jit",
                     "torch.package", "torch.utils.data", "torch.utils.benchmark", "torch.overrides")
ALLOWED_CPP_EXTENSION_NAMES = {"load_inline", "CUDA_HOME"}

BANNED_CALL_NAMES = {"exec", "eval", "compile", "__import__", "open", "globals", "locals", "vars", "setattr",
                     "delattr", "breakpoint", "input", "exit", "quit", "help", "memoryview", "classmethod_",
                     "chr", "ord", "bytes", "bytearray"}
BANNED_NAMES = {"__builtins__", "__loader__", "__spec__", "__file__", "sys", "os"}
BANNED_ATTRS = {
    # streams, graphs, events, compilation, global numerics flags
    "Stream", "stream", "ExternalStream", "set_stream", "graph", "CUDAGraph", "graphs", "make_graphed_callables",
    "Event", "compile", "backends", "set_float32_matmul_precision", "set_flush_denormal",
    "use_deterministic_algorithms", "set_device", "profiler", "library", "jit", "load_library",
    # file / serialisation
    "load", "save", "fromfile", "tofile", "memmap", "loadtxt", "savetxt", "from_file",
    # introspection / monkeypatch helpers
    "f_back", "f_globals", "f_locals", "gi_frame", "tb_frame", "co_code", "func_globals",
}
OK_DUNDER = {"__init__"}

# Native (C++/CUDA) tokens that have no place in a KernelBench kernel and enable hacks: Python C API and dynamic
# loading (tamper with the harness from native code), stream/graph/event APIs (timing tricks), host system calls.
NATIVE_BANNED = [
    r"Python\.h", r"pybind11", r"\bpy::", r"\bPyObject", r"\bPyImport", r"\bPy_[A-Za-z]", r"\bPyEval",
    r"\bdlopen\b", r"\bdlsym\b", r"\bRTLD_", r"__attribute__\s*\(\(\s*constructor",
    r"cudaStreamCreate", r"getStreamFromPool", r"cudaStreamNonBlocking", r"cudaStreamPerThread",
    r"cudaStreamBeginCapture", r"cudaGraph", r"cudaEvent", r"cudaLaunchHostFunc", r"cudaStreamAddCallback",
    r"cudaSetDevice", r"cudaDeviceReset", r"\bcuCtx", r"\bcuStream", r"\bcuGraph", r"\bcuEvent",
    r"CUDAStreamGuard", r"setCurrentCUDAStream", r"\bfopen\b", r"\bfwrite\b", r"\bfread\b", r"ofstream", r"ifstream",
    r"\bfstream\b", r"<fstream>", r"\bsystem\s*\(", r"\bpopen\b", r"\bfork\s*\(", r"\bexecv", r"\bexecl",
    r"\bgetenv\b", r"\bsetenv\b", r"std::thread", r"\bpthread", r"<thread>", r"std::async", r"<unistd\.h>",
    r"<sys/", r"\bsocket\s*\(", r"\bmmap\s*\(", r"\bptrace\b", r"\batexit\b", r"clock_gettime", r"gettimeofday",
    r"at::globalContext", r"setAllowTF32", r"c10::cuda::set",
]
_NATIVE_RE = [re.compile(p) for p in NATIVE_BANNED]
ALLOWED_LDFLAGS = {"-lcublas", "-lcublasLt", "-lcudnn", "-lcuda", "-lcurand", "-lcusparse", "-lcufft"}
BANNED_CFLAG_PARTS = ("-include", "-imacros", "-wrapper", "-fplugin", "-Xlinker", "-Wl,", "--compiler-bindir",
                      "-ccbin", "-B", "-specs", "--linker-options", "-L", "-l")


def scan_native(text: str) -> list[str]:
    return sorted({p.pattern for p in _NATIVE_RE if p.search(text)})


def _strings(tree) -> list[str]:
    out = []
    for n in ast.walk(tree):
        if isinstance(n, ast.Constant) and isinstance(n.value, str):
            out.append(n.value)
    return out


def _attr_root(n):
    while isinstance(n, ast.Attribute):
        n = n.value
    return n


def _dotted(n) -> str:
    parts = []
    while isinstance(n, ast.Attribute):
        parts.append(n.attr)
        n = n.value
    if isinstance(n, ast.Name):
        parts.append(n.id)
    return ".".join(reversed(parts))


def check_source(src: str) -> dict:
    errors: list[str] = []
    kernels: list[str] = []
    backend = set()
    try:
        tree = ast.parse(src)
    except SyntaxError as e:
        return {"ok": False, "errors": [f"syntax error: {e}"], "kernels": [], "backend": None}

    def err(node, msg):
        errors.append(f"line {getattr(node, 'lineno', '?')}: {msg}")

    has_modelnew = False
    triton_aliases = set()
    for n in ast.walk(tree):
        if isinstance(n, ast.Import):
            for a in n.names:
                if a.name.split(".")[0] == "triton":
                    triton_aliases.add(a.asname or a.name.split(".")[0])
        elif isinstance(n, ast.ImportFrom) and (n.module or "").split(".")[0] == "triton":
            for a in n.names:
                triton_aliases.add(a.asname or a.name)
    for n in ast.walk(tree):
        # ---- imports
        if isinstance(n, ast.Import):
            for a in n.names:
                top = a.name.split(".")[0]
                if top not in ALLOWED_IMPORTS or any(a.name.startswith(b) for b in BANNED_SUBMODULES):
                    err(n, f"import of '{a.name}' is not allowed")
                if a.name == "torch.utils.cpp_extension":
                    pass
        elif isinstance(n, ast.ImportFrom):
            mod = n.module or ""
            top = mod.split(".")[0]
            if n.level or top not in ALLOWED_IMPORTS or any(mod.startswith(b) for b in BANNED_SUBMODULES):
                err(n, f"import from '{mod}' is not allowed")
            for a in n.names:
                if a.name == "*":
                    err(n, "star imports are not allowed")
                if mod == "torch.utils.cpp_extension" and a.name not in ALLOWED_CPP_EXTENSION_NAMES:
                    err(n, f"only load_inline may be imported from torch.utils.cpp_extension (got {a.name})")
                if mod == "torch.utils" and a.name == "cpp_extension":
                    pass
                if a.name in BANNED_ATTRS or a.name.startswith("_"):
                    err(n, f"import of '{a.name}' is not allowed")
                if mod == "torch" and a.name in ("cuda",):
                    pass
        # ---- classes
        elif isinstance(n, ast.ClassDef):
            if n.name == "Model":
                err(n, "do not define a class named Model (the reference); define ModelNew only")
            if n.name == "ModelNew":
                has_modelnew = True
                for b in n.bases:
                    d = _dotted(b)
                    if d not in ("nn.Module", "torch.nn.Module", "Module"):
                        err(n, f"ModelNew must subclass nn.Module directly (got base '{d}')")
        # ---- functions: find triton kernels
        elif isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)):
            for dec in n.decorator_list:
                d = _dotted(dec.func if isinstance(dec, ast.Call) else dec)
                if d.endswith("jit") and d.split(".")[0] in ("triton",) or d in ("jit",):
                    kernels.append(n.name)
                    backend.add("triton")
                if d.endswith("autotune") or d.endswith("heuristics"):
                    pass
        # ---- exception handling, global
        elif isinstance(n, (ast.Try,)) or type(n).__name__ == "TryStar":
            err(n, "try/except is not allowed (no fallback paths)")
        elif isinstance(n, ast.Global):
            err(n, "global statements are not allowed")
        # ---- names
        elif isinstance(n, ast.Name):
            if n.id in BANNED_NAMES:
                err(n, f"use of '{n.id}' is not allowed")
        # ---- attributes
        elif isinstance(n, ast.Attribute):
            a = n.attr
            root = _attr_root(n)
            is_self = isinstance(root, ast.Name) and (root.id == "self" or root.id in triton_aliases)
            if a in OK_DUNDER:
                pass
            elif a.startswith("__"):
                err(n, f"dunder attribute '{a}' is not allowed")
            elif a.startswith("_") and not is_self:
                err(n, f"private attribute '{_dotted(n)}' is not allowed")
            elif a in BANNED_ATTRS and not is_self:
                err(n, f"'{_dotted(n)}' is not allowed (streams/graphs/events/compile/flags/file IO)")
        # ---- calls
        if isinstance(n, ast.Call):
            f = n.func
            if isinstance(f, ast.Name) and f.id in BANNED_CALL_NAMES:
                err(n, f"call to '{f.id}' is not allowed")
            if isinstance(f, ast.Name) and f.id == "getattr":
                if len(n.args) < 2 or not (isinstance(n.args[1], ast.Constant) and isinstance(n.args[1].value, str)
                                           and not n.args[1].value.startswith("_")):
                    err(n, "getattr needs a literal, non-private attribute name")
            fname = _dotted(f)
            if fname.endswith("load_inline"):
                backend.add("cuda")
                for kw in n.keywords:
                    if kw.arg in ("build_directory",):
                        err(n, "load_inline(build_directory=...) is not allowed")
                    if kw.arg == "extra_ldflags":
                        vals = _literal_list(kw.value)
                        if vals is None or any(v not in ALLOWED_LDFLAGS for v in vals):
                            err(n, f"extra_ldflags must be a literal list from {sorted(ALLOWED_LDFLAGS)}")
                    if kw.arg in ("extra_cflags", "extra_cuda_cflags"):
                        vals = _literal_list(kw.value)
                        if vals is None:
                            err(n, f"{kw.arg} must be a literal list of strings")
                        else:
                            for v in vals:
                                if any(v.startswith(b) for b in BANNED_CFLAG_PARTS):
                                    err(n, f"compiler flag '{v}' is not allowed")
        # ---- assignments (monkeypatching)
        targets = []
        if isinstance(n, ast.Assign):
            targets = n.targets
        elif isinstance(n, (ast.AugAssign, ast.AnnAssign)):
            targets = [n.target]
        elif isinstance(n, ast.Delete):
            targets = n.targets
        for t in targets:
            for tt in (t.elts if isinstance(t, (ast.Tuple, ast.List)) else [t]):
                if isinstance(tt, ast.Attribute):
                    root = _attr_root(tt)
                    if not (isinstance(root, ast.Name) and root.id == "self"):
                        err(n, f"assignment to attribute '{_dotted(tt)}' is not allowed (only self.<x>)")

    if not has_modelnew:
        errors.append("no class ModelNew defined")
    text = "".join(_strings(tree))
    text_nl = "\n".join(_strings(tree))
    for tok in sorted(set(scan_native(text)) | set(scan_native(text_nl))):
        errors.append(f"native source contains banned token /{tok}/")
    for m in re.finditer(r"__global__\s+(?:__launch_bounds__\s*\([^)]*\)\s*)?(?:static\s+)?void\s+([A-Za-z_]\w*)", text_nl):
        kernels.append(m.group(1))
    if "cuda" in backend and not re.search(r"__global__", text_nl):
        errors.append("load_inline is used but no __global__ kernel is defined in the sources")
    if not kernels:
        errors.append("no custom kernel: define a CUDA __global__ kernel via load_inline or a @triton.jit kernel")
    b = None if not backend else ("mixed" if len(backend) > 1 else next(iter(backend)))
    return {"ok": not errors, "errors": errors[:30], "kernels": sorted(set(kernels)), "backend": b}


def _literal_list(node):
    if isinstance(node, (ast.List, ast.Tuple)) and all(isinstance(e, ast.Constant) and isinstance(e.value, str)
                                                        for e in node.elts):
        return [e.value for e in node.elts]
    return None


if __name__ == "__main__":
    import json
    import sys
    print(json.dumps(check_source(open(sys.argv[1]).read()), indent=1))