"""
clankerDiffusion — local tool executor for the agent loop.

Tools the model can call via  <tool name="...">args</tool> :
  calc(expr)      safe arithmetic / math expression
  python(code)    run Python, stdout captured (threaded timeout)
  read_file(p)    read a text file
  list_dir(p)     list a directory
  retrieve(q)     RAG search over the local knowledge base (returns passages)

NOTE: python() executes locally on YOUR machine. It is sandboxed with a
restricted builtins set and a wall-clock timeout, but treat it as you would
any local code. Don't point clanker at untrusted input.
"""
import ast, math, io, contextlib, os, time, threading, builtins
from rag import retrieve as _rag_retrieve

SAFE_NAMES = {k: getattr(math, k) for k in dir(math) if not k.startswith("_")}
SAFE_NAMES.update({"abs": abs, "min": min, "max": max, "round": round,
                   "pow": pow, "len": len, "sum": sum, "sorted": sorted,
                   "int": int, "float": float, "str": str, "list": list,
                   "range": range, "True": True, "False": False, "None": None})


def _safe_eval(expr):
    node = ast.parse(expr, mode="eval").body

    def ev(n):
        if isinstance(n, ast.BinOp):
            a, b = ev(n.left), ev(n.right)
            if isinstance(n.op, ast.Add): return a + b
            if isinstance(n.op, ast.Sub): return a - b
            if isinstance(n.op, ast.Mult): return a * b
            if isinstance(n.op, ast.Div): return a / b
            if isinstance(n.op, ast.FloorDiv): return a // b
            if isinstance(n.op, ast.Mod): return a % b
            if isinstance(n.op, ast.Pow): return a ** b
            raise ValueError("op")
        if isinstance(n, ast.UnaryOp):
            v = ev(n.operand)
            return -v if isinstance(n.op, ast.USub) else +v
        if isinstance(n, ast.Constant) and isinstance(n.value, (int, float)):
            return n.value
        if isinstance(n, ast.Call) and isinstance(n.func, ast.Name):
            fn = SAFE_NAMES.get(n.func.id)
            if fn is None: raise ValueError(f"no {n.func.id}")
            return fn(*[ev(a) for a in n.args])
        if isinstance(n, ast.Name):
            if n.id in SAFE_NAMES: return SAFE_NAMES[n.id]
            raise ValueError(f"name {n.id}")
        raise ValueError("bad expr")
    return ev(node)


def _run_python(code, timeout=10.0):
    ns = {"__builtins__": {k: getattr(builtins, k)
                            for k in ("print", "len", "range", "list", "dict",
                                      "tuple", "set", "str", "int", "float",
                                      "bool", "min", "max", "sum", "sorted",
                                      "abs", "enumerate", "zip", "map", "filter",
                                      "open", "round", "pow") if hasattr(builtins, k)}}
    ns["math"] = math
    buf = io.StringIO()
    result = {}

    def target():
        try:
            with contextlib.redirect_stdout(buf):
                exec(code, ns)
            result["out"] = buf.getvalue()
        except Exception as e:
            result["out"] = buf.getvalue() + f"\nERROR: {type(e).__name__}: {e}"

    t = threading.Thread(target=target, daemon=True)
    t.start(); t.join(timeout)
    if t.is_alive():
        return f"ERROR: execution timed out after {timeout}s"
    return result.get("out", "").strip() or "(no output)"


def execute_tool(name, arg):
    try:
        if name == "calc":
            val = _safe_eval(arg.strip())
            return f"{arg.strip()} = {val}"
        if name == "python":
            return _run_python(arg)
        if name == "read_file":
            p = arg.strip().strip('"\'')
            if not os.path.exists(p): return f"ERROR: no such file: {p}"
            with open(p, "r", errors="replace") as f:
                return f.read()[:4000]
        if name == "list_dir":
            p = arg.strip().strip('"\'') or "."
            if not os.path.isdir(p): return f"ERROR: no such dir: {p}"
            return "\n".join(sorted(os.listdir(p))[:200])
        if name == "retrieve":
            try:
                passages = _rag_retrieve(arg.strip(), k=4)
                if not passages:
                    return "(no relevant context found in the knowledge base)"
                return "\n---\n".join(passages)
            except Exception as e:
                return f"ERROR: retrieve failed: {e}"
        return f"ERROR: unknown tool '{name}'"
    except Exception as e:
        return f"ERROR: {type(e).__name__}: {e}"


def parse_tool_calls(text):
    """Return list of (name, argstring) from <tool name=\"...\">..</tool> blocks."""
    calls = []
    i = 0
    while True:
        s = text.find("<tool", i)
        if s < 0: break
        e = text.find("</tool>", s)
        if e < 0: break
        head = text[s:text.find(">", s) + 1]
        # head like: <tool name="calc">
        import re
        m = re.search(r'name="([^"]+)"', head)
        name = m.group(1) if m else "?"
        arg = text[s + len(head):e].strip()
        calls.append((name, arg))
        i = e + len("</tool>")
    return calls