"""Deterministic tools Dinah consults before the encoder. Each one answers only what it can settle exactly.

`code`: "what does this function return" questions are answered by running the function. It fires on content,
the same way for every request (it never looks at a benchmark name, id or task string), when
  1. a state field is Python source with a top-level function,
  2. another state field is the call's arguments (`[1, 2], 'a'`), screened like the source, and
  3. every option of a `choice` question is a Python literal.
The function runs in a child process (`python -I -S`, empty environment, temporary directory, CPU, memory,
file-size and wall-time limits) after an AST screen (imports only from a standard-library allow list; no
open/exec/eval/__import__, no dunder names, no underscore attributes). The tool answers only when exactly one
option equals the result (same type and value); an error, a timeout, or zero or several matching options leave
the question to the encoder. Every option still gets a probability.
"""
from __future__ import annotations

import ast
import json
import re
import subprocess
import sys
import tempfile

MODULES = {"math", "string", "re", "collections", "itertools", "functools", "operator", "copy", "heapq", "bisect",
           "datetime", "fractions", "decimal", "statistics", "textwrap", "unicodedata", "numbers", "typing", "dataclasses",
           "enum", "array", "calendar"}
FORBIDDEN = {"open", "exec", "eval", "compile", "__import__", "breakpoint", "globals", "locals", "vars", "memoryview",
             "setattr", "delattr"}
DUNDER = re.compile(r"__\w+__")
P_MATCH = 0.97           # probability of the matching option; the rest is split evenly
TIMEOUT_S = 3.0

CHILD = r"""
import ast, json, resource, sys
for lim, v in ((resource.RLIMIT_CPU, 2), (resource.RLIMIT_FSIZE, 0), (resource.RLIMIT_AS, 1 << 30)):
    try: resource.setrlimit(lim, (v, v))
    except (ValueError, OSError): pass
p = json.load(sys.stdin)
g = {"__name__": "__tool__"}
exec(compile(p["source"], "<source>", "exec"), g)
r = eval(compile(p["fn"] + "(" + p["args"] + ")", "<call>", "eval"), g)
def same(a, b):
    if type(a) is not type(b): return False
    if isinstance(a, (list, tuple)): return len(a) == len(b) and all(same(x, y) for x, y in zip(a, b))
    if isinstance(a, dict): return a.keys() == b.keys() and all(same(a[k], b[k]) for k in a)
    return a == b
hits = []
for i, o in enumerate(p["options"]):
    try:
        if same(r, ast.literal_eval(o)): hits.append(i)
    except Exception: pass
print(json.dumps({"hits": hits}))
"""


def _safe(tree):
    for node in ast.walk(tree):
        if isinstance(node, ast.Import) and any(a.name.split(".")[0] not in MODULES for a in node.names):
            return None
        if isinstance(node, ast.ImportFrom) and (node.module or "").split(".")[0] not in MODULES:
            return None
        if isinstance(node, ast.Name) and (node.id in FORBIDDEN or node.id.startswith("__")):
            return None
        if isinstance(node, ast.Attribute) and node.attr.startswith("_"):
            return None
        if isinstance(node, ast.Constant) and isinstance(node.value, str) and DUNDER.search(node.value):
            return None
    return tree


def _function(text):
    """Name of the last top-level function if `text` is screened Python source, else None."""
    if not isinstance(text, str) or "def " not in text:
        return None
    try:
        tree = _safe(ast.parse(text))
    except (SyntaxError, ValueError, RecursionError):
        return None
    if tree is None:
        return None
    defs = [n.name for n in tree.body if isinstance(n, ast.FunctionDef)]
    return defs[-1] if defs else None


def _call_args(text, fn):
    if not isinstance(text, str) or "def " in text:
        return False
    try:
        tree = ast.parse(f"{fn}({text})", mode="eval")
    except (SyntaxError, ValueError, MemoryError, RecursionError):
        return False
    return isinstance(tree.body, ast.Call) and _safe(tree) is not None


def _literals(options):
    try:
        for o in options:
            ast.literal_eval(o)
        return True
    except (SyntaxError, ValueError, TypeError, MemoryError, RecursionError):
        return False


def _candidate(state):
    if not isinstance(state, dict):
        return None
    fields = [v for v in state.values() if isinstance(v, str)]
    for src in fields:
        fn = _function(src)
        if fn:
            for args in fields:
                if args is not src and _call_args(args, fn):
                    return src, fn, args
    return None


def _run(source, fn, args, options, python=sys.executable):
    payload = json.dumps({"source": source, "fn": fn, "args": args, "options": options})
    with tempfile.TemporaryDirectory() as d:
        try:
            r = subprocess.run([python, "-I", "-S", "-c", CHILD], input=payload, capture_output=True, text=True,
                               timeout=TIMEOUT_S, cwd=d, env={"PYTHONHASHSEED": "0"})
        except subprocess.TimeoutExpired:
            return None
    if r.returncode != 0:
        return None
    try:
        return json.loads(r.stdout.strip().splitlines()[-1])["hits"]
    except (ValueError, IndexError, KeyError):
        return None


def code(state, questions):
    """Answers {key: answer} for the questions the code tool settles; {} when it does not fire."""
    cand = _candidate(state)
    if not cand:
        return {}
    out = {}
    for key, q in questions.items():
        if q.get("type") != "choice":
            continue
        keys = list(q["criteria"])
        options = [q["criteria"][k] for k in keys]
        if len(options) < 2 or not _literals(options):
            continue
        hits = _run(*cand, options)
        if not hits or len(hits) != 1:
            continue
        rest = (1.0 - P_MATCH) / (len(keys) - 1)
        probs = {k: (P_MATCH if i == hits[0] else rest) for i, k in enumerate(keys)}
        out[key] = {"type": "choice", "choice": keys[hits[0]], "probabilities": probs, "confidence": P_MATCH}
    return out


TOOLS = {"code": code}