"""Compile and run a minted probe in a restricted namespace.

A probe is Python source defining `probe(ctx) -> dict` with keys applicable (bool), flag (bool, True = the work is
WRONG), evidence (str). Restrictions (enforced, not advisory):
  - imports limited to a stdlib whitelist (no os, io, pathlib, subprocess, socket, importlib, sys, builtins ...);
  - no open/exec/eval/compile/__import__/globals/getattr-on-dunder tricks in the source (static screen);
  - forbidden identifiers: ground_truth, private_data, public_data, evaluation, report (the hidden test data);
  - per-call wall-clock timeout (SIGALRM) in the worker.
The probe receives only what state.ProbeContext exposes (instruction, supervisor, datetime, start/end models,
main_user, predicted_answer, task_status) plus pure helpers.
"""
from __future__ import annotations

import builtins
import re
import signal
from typing import Any

ALLOWED_IMPORTS = {
    "re", "json", "math", "datetime", "collections", "itertools", "functools", "string", "difflib", "unicodedata",
    "decimal", "fractions", "statistics", "operator", "dateutil", "dateutil.parser", "dateutil.relativedelta",
    "calendar", "typing", "dataclasses", "_strptime", "time",
}
FORBIDDEN_SRC = [
    r"\bopen\s*\(", r"\bexec\s*\(", r"\beval\s*\(", r"\bcompile\s*\(", r"__import__", r"\bglobals\s*\(",
    r"\blocals\s*\(", r"__builtins__", r"__subclasses__", r"__globals__", r"__code__", r"\bground_truth\b",
    r"\bprivate_data\b", r"\bpublic_data\b", r"\bimportlib\b", r"\bsubprocess\b", r"\bsocket\b", r"\bpathlib\b",
    r"\bimport\s+os\b", r"\bfrom\s+os\b", r"\bimport\s+sys\b", r"\bfrom\s+sys\b", r"\bshutil\b", r"\bpickle\b",
    r"\bsetattr\s*\(", r"\bdelattr\s*\(", r"\.save\s*\(", r"\.delete\s*\(", r"\.update_record", r"\.create\s*\(",
    r"\.session\b", r"\bsqlite3\b", r"\bexecute\s*\(",
]

_real_import = builtins.__import__


def _guarded_import(name, globals=None, locals=None, fromlist=(), level=0):  # noqa: A002
    root = name.split(".")[0]
    if name not in ALLOWED_IMPORTS and root not in ALLOWED_IMPORTS:
        raise ImportError(f"import of '{name}' is not allowed in a probe")
    return _real_import(name, globals, locals, fromlist, level)


SAFE_BUILTINS = {k: getattr(builtins, k) for k in (
    "abs", "all", "any", "bool", "dict", "enumerate", "filter", "float", "frozenset", "int", "isinstance",
    "issubclass", "iter", "len", "list", "map", "max", "min", "next", "object", "print", "range", "repr",
    "reversed", "round", "set", "slice", "sorted", "str", "sum", "tuple", "zip", "Exception", "ValueError",
    "KeyError", "TypeError", "IndexError", "AttributeError", "StopIteration", "ZeroDivisionError",
    "RuntimeError", "NotImplementedError", "AssertionError", "ArithmeticError", "LookupError", "hasattr",
    "getattr", "callable", "chr", "ord", "divmod", "pow", "hash", "id", "type", "format", "bytes", "complex",
    "property", "staticmethod", "classmethod", "super", "True", "False", "None", "OverflowError",
)}
SAFE_BUILTINS["__import__"] = _guarded_import
SAFE_BUILTINS["__build_class__"] = builtins.__build_class__
SAFE_BUILTINS["__name__"] = "probe"


def screen(src: str) -> list[str]:
    bad = []
    for pat in FORBIDDEN_SRC:
        if re.search(pat, src):
            bad.append(pat)
    if re.search(r"getattr\s*\([^)]*['\"]__", src):
        bad.append("getattr-dunder")
    if "def probe" not in src:
        bad.append("no `def probe`")
    return bad


def helpers() -> dict[str, Any]:
    from appworld.common.utils import answer_to_text, dict_of, list_of, set_of
    return {"list_of": list_of, "set_of": set_of, "dict_of": dict_of, "answer_to_text": answer_to_text}


def compile_probe(src: str):
    bad = screen(src)
    if bad:
        raise ValueError(f"probe source rejected by screen: {bad}")
    ns: dict[str, Any] = {"__builtins__": SAFE_BUILTINS, **helpers()}
    code = builtins.compile(src, "<probe>", "exec")
    builtins.exec(code, ns)  # noqa: S102  (restricted namespace)
    fn = ns.get("probe")
    if not callable(fn):
        raise ValueError("probe source defines no callable `probe`")
    return fn


class _Timeout(Exception):
    pass


def _alarm(signum, frame):  # noqa: ARG001
    raise _Timeout()


def run_probe(fn, ctx, timeout_s: int = 30) -> dict:
    old = signal.signal(signal.SIGALRM, _alarm)
    signal.alarm(timeout_s)
    try:
        out = fn(ctx)
        if not isinstance(out, dict):
            return {"status": "bad_output", "applicable": False, "flag": False, "evidence": repr(out)[:300]}
        app = bool(out.get("applicable", True))
        flag = bool(out.get("flag", False)) and app
        return {"status": "ok", "applicable": app, "flag": flag, "evidence": str(out.get("evidence", ""))[:600]}
    except _Timeout:
        return {"status": "timeout", "applicable": False, "flag": False, "evidence": f"timeout {timeout_s}s"}
    except Exception as e:  # noqa: BLE001
        return {"status": "error", "applicable": False, "flag": False,
                "evidence": f"{type(e).__name__}: {str(e)[:400]}"}
    finally:
        signal.alarm(0)
        signal.signal(signal.SIGALRM, old)