"""Restricted stateful Python for Nemotron Cascade-2 math tool calls.

Tercet-R saw `stateful_python_code_exec` on `math/math_tool.jsonl` during
the Cascade-2 SFT stage. Observations are the Jupyter-style stdout / last
expression value, not a JSON envelope.

The worker process only allows a math-oriented import whitelist and rejects
dunder access, so a tool call cannot read the filesystem or start a shell.
"""

from __future__ import annotations

import ast
import io
import json
import math
import os
import subprocess
import sys
import threading
from collections.abc import Mapping
from pathlib import Path
from queue import Empty, Queue
from typing import Any


NEMOTRON_PYTHON_EXEC_TOOL: dict[str, Any] = {
    "type": "function",
    "function": {
        "name": "stateful_python_code_exec",
        "description": (
            "Call this function to execute Python code in a stateful Jupyter "
            "notebook environment. Python will respond with the output of the "
            "execution or time out after 120.0 seconds."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Code to execute",
                }
            },
            "required": ["code"],
        },
    },
}

CALCULATOR_TOOL: dict[str, Any] = {
    "type": "function",
    "function": {
        "name": "calculator",
        "description": "Evaluate a math expression and return the numeric result.",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "Math expression, for example 12.5 * (3 + 4)",
                }
            },
            "required": ["expression"],
        },
    },
}

PYTHON_EXEC_TOOL_NAMES = frozenset(
    {
        "stateful_python_code_exec",
        "python_code_exec",
        "code_interpreter",
        "python_exec",
        "python",
    }
)
CALCULATOR_TOOL_NAMES = frozenset({"calculator", "calc"})
DEFAULT_EXEC_TIMEOUT_SECONDS = 30.0
_SRC_ROOT = Path(__file__).resolve().parents[1]

_session_lock = threading.Lock()
_default_session: StatefulPythonSession | None = None


def _normalize_tool_name(name: str) -> str:
    return name.strip().lower().replace("-", "_")


def is_python_exec_tool_name(name: str) -> bool:
    return _normalize_tool_name(name) in PYTHON_EXEC_TOOL_NAMES


def is_calculator_tool_name(name: str) -> bool:
    return _normalize_tool_name(name) in CALCULATOR_TOOL_NAMES


def is_auto_math_tool_name(name: str) -> bool:
    return is_python_exec_tool_name(name) or is_calculator_tool_name(name)


def code_from_arguments(arguments: Mapping[str, Any] | None) -> str:
    if not arguments:
        return ""
    for key in ("code", "expression", "expr", "source"):
        raw = arguments.get(key)
        if isinstance(raw, str) and raw.strip():
            return raw
        if raw is not None and key != "code":
            return str(raw)
    return ""


class CodeExecError(ValueError):
    """Rejected or failed tool code."""


class StatefulPythonSession:
    """One long-lived restricted interpreter, matching the training tool."""

    def __init__(self, *, timeout_seconds: float = DEFAULT_EXEC_TIMEOUT_SECONDS) -> None:
        self.timeout_seconds = timeout_seconds
        self._lock = threading.Lock()
        self._process: subprocess.Popen[str] | None = None

    def close(self) -> None:
        with self._lock:
            self._kill_locked()

    def reset(self) -> None:
        with self._lock:
            self._kill_locked()

    def run(self, code: str) -> str:
        text = code.strip()
        if not text:
            raise CodeExecError("code is empty")
        with self._lock:
            return self._run_locked(text)

    def _run_locked(self, code: str) -> str:
        process = self._ensure_process_locked()
        try:
            process.stdin.write(json.dumps({"code": code}, ensure_ascii=False) + "\n")
            process.stdin.flush()
            line = self._readline_locked(process, self.timeout_seconds)
        except (BrokenPipeError, OSError) as error:
            self._kill_locked()
            raise CodeExecError(f"python worker died: {error}") from error
        if not line:
            stderr = ""
            if process.stderr is not None:
                try:
                    stderr = process.stderr.read()
                except OSError:
                    stderr = ""
            self._kill_locked()
            detail = stderr.strip() or "python worker closed stdout"
            raise CodeExecError(detail)
        try:
            payload = json.loads(line)
        except json.JSONDecodeError as error:
            self._kill_locked()
            raise CodeExecError(f"python worker returned invalid JSON: {line!r}") from error
        if not isinstance(payload, dict):
            raise CodeExecError("python worker returned a non-object")
        output = str(payload.get("output") or "")
        if payload.get("ok") is True:
            return output if output else "None"
        raise CodeExecError(output or "execution failed")

    def _readline_locked(self, process: subprocess.Popen[str], timeout: float) -> str:
        if process.stdout is None:
            raise CodeExecError("python worker has no stdout")
        lines: Queue[str] = Queue()

        def _read() -> None:
            lines.put(process.stdout.readline() if process.stdout is not None else "")

        reader = threading.Thread(target=_read, daemon=True)
        reader.start()
        reader.join(timeout)
        if reader.is_alive():
            self._kill_locked()
            raise CodeExecError(f"timed out after {timeout:g}s")
        try:
            return lines.get_nowait()
        except Empty:
            return ""

    def _ensure_process_locked(self) -> subprocess.Popen[str]:
        process = self._process
        if process is not None and process.poll() is None:
            return process
        env = os.environ.copy()
        pythonpath = env.get("PYTHONPATH", "")
        env["PYTHONPATH"] = (
            str(_SRC_ROOT) if not pythonpath else f"{_SRC_ROOT}{os.pathsep}{pythonpath}"
        )
        self._process = subprocess.Popen(
            [sys.executable, "-m", "tiny_gdn.code_exec"],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
            env=env,
        )
        if self._process.stdin is None or self._process.stdout is None:
            self._kill_locked()
            raise CodeExecError("failed to start python worker")
        return self._process

    def _kill_locked(self) -> None:
        process = self._process
        self._process = None
        if process is None:
            return
        try:
            process.kill()
            process.wait(timeout=2)
        except (OSError, subprocess.TimeoutExpired):
            pass


def default_python_session() -> StatefulPythonSession:
    global _default_session
    with _session_lock:
        if _default_session is None:
            _default_session = StatefulPythonSession()
        return _default_session


def reset_default_python_session() -> None:
    global _default_session
    with _session_lock:
        session = _default_session
        _default_session = None
    if session is not None:
        session.close()


def execute_calculator(expression: str) -> str:
    tree = ast.parse(expression, mode="eval")
    _reject_unsafe_ast(tree)
    value = eval(  # noqa: S307 — AST already rejected unsafe nodes
        compile(tree, "<calculator>", "eval"),
        {"__builtins__": {}},
        _calculator_namespace(),
    )
    return _format_result(value)


def execute_math_tool(name: str, arguments: Mapping[str, Any] | None) -> str:
    source = code_from_arguments(arguments)
    if is_calculator_tool_name(name) and "\n" not in source and ";" not in source:
        try:
            return execute_calculator(source)
        except (SyntaxError, CodeExecError, TypeError, ValueError, ZeroDivisionError) as error:
            raise CodeExecError(str(error)) from error
    return default_python_session().run(source)


def _calculator_namespace() -> dict[str, Any]:
    names = {
        key: getattr(math, key)
        for key in dir(math)
        if not key.startswith("_")
    }
    names.update(
        {
            "abs": abs,
            "min": min,
            "max": max,
            "round": round,
            "pow": pow,
            "pi": math.pi,
            "e": math.e,
            "tau": math.tau,
            "inf": math.inf,
        }
    )
    return names


def _reject_unsafe_ast(tree: ast.AST) -> None:
    for node in ast.walk(tree):
        if isinstance(node, (ast.ClassDef, ast.AsyncFunctionDef)):
            raise CodeExecError("class definitions are blocked")
        if isinstance(node, ast.Attribute) and node.attr.startswith("_"):
            raise CodeExecError("dunder attribute access is blocked")
        if isinstance(node, ast.Name) and node.id.startswith("_"):
            raise CodeExecError("dunder names are blocked")
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
            if node.func.id in _BANNED_CALLS:
                raise CodeExecError(f"{node.func.id}() is blocked")


_BANNED_CALLS = frozenset(
    {
        "eval",
        "exec",
        "compile",
        "open",
        "input",
        "breakpoint",
        "getattr",
        "setattr",
        "delattr",
        "globals",
        "locals",
        "vars",
        "dir",
        "help",
        "__import__",
        "memoryview",
        "exit",
        "quit",
    }
)


def _format_result(value: Any) -> str:
    if value is None:
        return "None"
    if isinstance(value, bool):
        return str(value)
    if isinstance(value, float):
        if value.is_integer() and abs(value) < 1e15:
            return str(int(value))
        return format(value, ".12g")
    return str(value)


_ALLOWED_IMPORT_ROOTS = frozenset(
    {
        "cmath",
        "collections",
        "copy",
        "decimal",
        "fractions",
        "functools",
        "itertools",
        "json",
        "math",
        "mpmath",
        "numbers",
        "numpy",
        "operator",
        "re",
        "statistics",
        "string",
        "sympy",
        "textwrap",
        "unicodedata",
    }
)

_WORKER_NAMESPACE: dict[str, Any] | None = None


def _allowed_import(
    name: str,
    globals: dict[str, Any] | None = None,
    locals: dict[str, Any] | None = None,
    fromlist: tuple[str, ...] = (),
    level: int = 0,
) -> Any:
    root = name.split(".")[0]
    if root not in _ALLOWED_IMPORT_ROOTS:
        raise ImportError(f"import of {name!r} is blocked")
    return __import__(name, globals, locals, fromlist, level)


def _worker_namespace() -> dict[str, Any]:
    global _WORKER_NAMESPACE
    if _WORKER_NAMESPACE is None:
        builtins = {
            "abs": abs,
            "all": all,
            "any": any,
            "bin": bin,
            "bool": bool,
            "bytes": bytes,
            "chr": chr,
            "complex": complex,
            "dict": dict,
            "divmod": divmod,
            "enumerate": enumerate,
            "filter": filter,
            "float": float,
            "format": format,
            "frozenset": frozenset,
            "hex": hex,
            "int": int,
            "isinstance": isinstance,
            "issubclass": issubclass,
            "iter": iter,
            "len": len,
            "list": list,
            "map": map,
            "max": max,
            "min": min,
            "next": next,
            "oct": oct,
            "ord": ord,
            "pow": pow,
            "print": print,
            "range": range,
            "repr": repr,
            "reversed": reversed,
            "round": round,
            "set": set,
            "slice": slice,
            "sorted": sorted,
            "str": str,
            "sum": sum,
            "tuple": tuple,
            "zip": zip,
            "True": True,
            "False": False,
            "None": None,
            "__import__": _allowed_import,
        }
        _WORKER_NAMESPACE = {
            "__builtins__": builtins,
            "__name__": "__tool__",
            "math": math,
        }
    return _WORKER_NAMESPACE


def run_cell(code: str) -> str:
    tree = ast.parse(code)
    _reject_unsafe_import_roots(tree)
    _reject_unsafe_ast(tree)
    namespace = _worker_namespace()
    buffer = io.StringIO()
    previous = sys.stdout
    sys.stdout = buffer
    try:
        if tree.body and isinstance(tree.body[-1], ast.Expr):
            body = tree.body[:-1]
            last = tree.body[-1]
            if body:
                exec(  # noqa: S102
                    compile(ast.Module(body, type_ignores=[]), "<tool>", "exec"),
                    namespace,
                    namespace,
                )
            value = eval(  # noqa: S307
                compile(ast.Expression(last.value), "<tool>", "eval"),
                namespace,
                namespace,
            )
            printed = buffer.getvalue()
            if value is None:
                return printed if printed else "None"
            rendered = _format_result(value)
            return f"{printed}{rendered}" if printed else rendered
        exec(compile(tree, "<tool>", "exec"), namespace, namespace)  # noqa: S102
        printed = buffer.getvalue()
        return printed if printed else "None"
    finally:
        sys.stdout = previous


def _reject_unsafe_import_roots(tree: ast.AST) -> None:
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                root = alias.name.split(".")[0]
                if root not in _ALLOWED_IMPORT_ROOTS:
                    raise CodeExecError(f"import of {alias.name!r} is blocked")
        elif isinstance(node, ast.ImportFrom):
            root = (node.module or "").split(".")[0]
            if root not in _ALLOWED_IMPORT_ROOTS:
                raise CodeExecError(f"import of {node.module!r} is blocked")


def _worker_loop() -> None:
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            request = json.loads(line)
            code = str(request.get("code") or "")
            output = run_cell(code)
            sys.stdout.write(json.dumps({"ok": True, "output": output}, ensure_ascii=False) + "\n")
        except Exception as error:
            sys.stdout.write(
                json.dumps({"ok": False, "output": str(error)}, ensure_ascii=False) + "\n"
            )
        sys.stdout.flush()


if __name__ == "__main__":
    _worker_loop()