"""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 itertools
import json
import math
import os
import re
import statistics
import subprocess
import sys
import threading
from collections.abc import Iterable, Iterator, 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 result.",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": (
                        "A math expression in standard calculator notation, for example "
                        "(3 + 4)^2 * sqrt(2) / ln(10). Use + - * / ^ (power), % (remainder), "
                        "! (factorial) and parentheses; the constants pi and e; and functions "
                        "such as sqrt, cbrt, exp, ln, log (natural, or log(x, base)), log10, "
                        "log2, sin, cos, tan, asin, acos, atan (radians), abs, round, floor, "
                        "ceil, factorial, comb, perm, gcd, lcm, min, max, sum, mean and median."
                    ),
                }
            },
            "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 is_calculator_expression(source: str) -> bool:
    """A one-line calculator call runs in the calculator; several statements run as notebook code."""
    return "\n" not in source.strip() and ";" not in source


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 CalculatorError(CodeExecError):
    """An expression the calculator cannot evaluate. The message says how to fix it."""


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()


MAX_EXPRESSION_CHARS = 10_000
# Exact integer powers past this many bits are computed as floats, and overflow there.
_MAX_EXACT_POWER_BITS = 100_000
_MAX_FACTORIAL = 100_000
_MAX_LISTED_ITEMS = 1_000
# About 3,600 digits, inside the 4,300 that Python's int-to-str conversion allows.
_MAX_EXACT_INTEGER_BITS = 12_000
_POWER = "_calculator_power"
_SYMBOLS = str.maketrans(
    {"×": "*", "·": "*", "÷": "/", "−": "-", "–": "-", "π": "pi", "∞": "inf", "²": "**2", "³": "**3"}
)


def calculate(expression: str, variables: Mapping[str, Any] | None = None) -> str:
    """Evaluate one expression in calculator notation and format its result.

    `^` is a power and `n!` a factorial, as on a calculator, and functions
    such as sqrt, ln and sin need no `math.` prefix. Python's `**`, `math.`
    names and comprehensions work as well. `variables` are names that a
    notebook in the same conversation defined.
    """
    text = _calculator_text(expression)
    try:
        tree = ast.parse(text, mode="eval")
    except SyntaxError as error:
        raise CalculatorError(
            f"could not read the expression ({error.msg}). "
            "Write every operator out, for example 2*pi or 3*(4 + 5)."
        ) from None
    _reject_unsafe_ast(tree)
    tree = ast.fix_missing_locations(_GuardedPower().visit(tree))
    # One dict as globals, so generator expressions see the functions too.
    namespace: dict[str, Any] = {"__builtins__": {}, **_calculator_functions(), **(variables or {}), _POWER: _power}
    try:
        value = eval(compile(tree, "<calculator>", "eval"), namespace)  # noqa: S307 — AST already rejected unsafe nodes
        return _calculator_result(value)
    except CodeExecError:
        raise
    except NameError as error:
        raise CalculatorError(
            f"unknown name {error.name or str(error)!r}. Call functions by name, such as sqrt(2), "
            "ln(5) or sin(pi/6), and use the constants pi and e."
        ) from None
    except ZeroDivisionError:
        raise CalculatorError("division by zero") from None
    except OverflowError:
        raise CalculatorError("the result is too large") from None
    except ValueError as error:
        if "domain" in str(error):
            raise CalculatorError(
                "math domain error: an input is outside the function's domain, "
                "such as the square root or log of a negative number"
            ) from None
        raise CalculatorError(str(error)) from None
    except Exception as error:  # noqa: BLE001 — every failure goes back to the model as text
        raise CalculatorError(str(error) or type(error).__name__) from None


def execute_calculator(expression: str) -> str:
    return calculate(expression)


def execute_math_tool(name: str, arguments: Mapping[str, Any] | None) -> str:
    source = code_from_arguments(arguments)
    if is_calculator_tool_name(name) and is_calculator_expression(source):
        return calculate(source)
    return default_python_session().run(source)


def _calculator_text(expression: str) -> str:
    """Calculator notation as Python: ^ as **, n! as factorial(n), √, π, × and ÷ spelled out."""
    text = expression.strip().translate(_SYMBOLS)
    if text.endswith("=") and not text.endswith("=="):
        text = text[:-1].rstrip()
    if not text:
        raise CalculatorError("the expression is empty")
    if len(text) > MAX_EXPRESSION_CHARS:
        raise CalculatorError(f"the expression is longer than {MAX_EXPRESSION_CHARS:,} characters")
    text = re.sub(r"√\s*(?=\()", "sqrt", text)
    text = re.sub(r"√\s*(\d+(?:\.\d*)?|[A-Za-z_]\w*)", r"sqrt(\1)", text)
    text = re.sub(r"\bmod\b", "%", text)
    return _expand_factorials(text.replace("^", "**"))


def _expand_factorials(text: str) -> str:
    """Rewrite postfix factorials, 5! or (n + 1)!, as factorial() calls. != stays a comparison."""
    position = 0
    while (index := text.find("!", position)) >= 0:
        if text.startswith("!=", index):
            position = index + 2
            continue
        start = _operand_start(text, index)
        if start is None:
            raise CalculatorError("! needs a number, a name or a parenthesized expression before it, as in 5! or (n + 1)!")
        text = f"{text[:start]}factorial({text[start:index].strip()}){text[index + 1:]}"
        position = start
    return text


def _operand_start(text: str, end: int) -> int | None:
    """Where the operand that ends just before `end` starts: a number, a name, a call or a parenthesized group."""
    index = end
    while index > 0 and text[index - 1].isspace():
        index -= 1
    stop = index
    if index > 0 and text[index - 1] == ")":
        depth = 0
        cursor = index - 1
        while cursor >= 0:
            if text[cursor] == ")":
                depth += 1
            elif text[cursor] == "(":
                depth -= 1
                if depth == 0:
                    break
            cursor -= 1
        if cursor < 0:
            return None
        index = cursor
    while index > 0 and (text[index - 1].isalnum() or text[index - 1] in "_."):
        index -= 1
    return index if index < stop else None


class _GuardedPower(ast.NodeTransformer):
    """Send every power through `_power`, which keeps huge exact integers from stalling the worker."""

    def visit_BinOp(self, node: ast.BinOp) -> ast.AST:
        self.generic_visit(node)
        if isinstance(node.op, ast.Pow):
            return ast.Call(func=ast.Name(id=_POWER, ctx=ast.Load()), args=[node.left, node.right], keywords=[])
        return node


def _power(base: Any, exponent: Any, modulus: Any = None) -> Any:
    if modulus is not None:
        return pow(base, exponent, modulus)
    if (
        type(base) is int
        and type(exponent) is int
        and exponent > 0
        and abs(base) > 1
        and exponent * math.log2(abs(base)) > _MAX_EXACT_POWER_BITS
    ):
        return float(base) ** exponent
    return base**exponent


def _factorial(value: Any) -> Any:
    if isinstance(value, float) and value.is_integer():
        value = int(value)
    if isinstance(value, int) and not isinstance(value, bool):
        if value < 0:
            raise CalculatorError("factorial is not defined for negative integers")
        if value > _MAX_FACTORIAL:
            raise CalculatorError(f"factorial({value}) is too large; lgamma({value} + 1) gives its natural log")
        return math.factorial(value)
    return math.gamma(value + 1)


def _values(arguments: tuple[Any, ...]) -> list[Any]:
    """An aggregate's numbers, given separately or as one list: mean(1, 2, 3) or mean([1, 2, 3])."""
    if len(arguments) == 1 and isinstance(arguments[0], Iterable) and not isinstance(arguments[0], (str, bytes)):
        return list(arguments[0])
    return list(arguments)


def _root(value: Any, degree: Any) -> Any:
    """The real nth root, so root(-8, 3) is -2."""
    if isinstance(value, (int, float)) and value < 0 and float(degree).is_integer() and int(degree) % 2 == 1:
        return -((-value) ** (1 / degree))
    return value ** (1 / degree)


def _calculator_functions() -> dict[str, Any]:
    names: dict[str, Any] = {key: getattr(math, key) for key in dir(math) if not key.startswith("_")}
    names.update(
        {
            "ln": lambda value: math.log(value),
            "arcsin": math.asin,
            "arccos": math.acos,
            "arctan": math.atan,
            "arctan2": math.atan2,
            "sec": lambda value: 1 / math.cos(value),
            "csc": lambda value: 1 / math.sin(value),
            "cot": lambda value: 1 / math.tan(value),
            "sign": lambda value: (value > 0) - (value < 0),
            "sgn": lambda value: (value > 0) - (value < 0),
            "root": _root,
            "factorial": _factorial,
            "fact": _factorial,
            "nCr": math.comb,
            "choose": math.comb,
            "nPr": math.perm,
            "pow": _power,
            "abs": abs,
            "round": round,
            "int": int,
            "float": float,
            "divmod": divmod,
            "range": range,
            "len": len,
            "sorted": sorted,
            "list": list,
            "tuple": tuple,
            "sum": lambda *values: sum(_values(values)),
            "min": lambda *values: min(_values(values)),
            "max": lambda *values: max(_values(values)),
            "prod": lambda *values: math.prod(_values(values)),
            "product": lambda *values: math.prod(_values(values)),
            "mean": lambda *values: statistics.fmean(_values(values)),
            "avg": lambda *values: statistics.fmean(_values(values)),
            "average": lambda *values: statistics.fmean(_values(values)),
            "median": lambda *values: statistics.median(_values(values)),
            "mode": lambda *values: statistics.mode(_values(values)),
            "stdev": lambda *values: statistics.stdev(_values(values)),
            "std": lambda *values: statistics.stdev(_values(values)),
            "pstdev": lambda *values: statistics.pstdev(_values(values)),
            "variance": lambda *values: statistics.variance(_values(values)),
            "var": lambda *values: statistics.variance(_values(values)),
            "pvariance": lambda *values: statistics.pvariance(_values(values)),
            "phi": (1 + math.sqrt(5)) / 2,
            "deg": math.pi / 180,
            "math": math,
        }
    )
    return names


def _calculator_result(value: Any) -> str:
    if value is None or isinstance(value, bool):
        return str(value)
    if isinstance(value, int):
        return _exact_integer(value)
    if isinstance(value, float):
        if math.isnan(value):
            return "nan"
        if math.isinf(value):
            return "inf" if value > 0 else "-inf"
        if value.is_integer() and abs(value) < 1e15:
            return str(int(value))
        return format(value, ".12g")
    if isinstance(value, complex):
        if value.imag == 0:
            return _calculator_result(value.real)
        if value.real == 0:
            return f"{_calculator_result(value.imag)}i"
        sign = "-" if value.imag < 0 else "+"
        return f"{_calculator_result(value.real)} {sign} {_calculator_result(abs(value.imag))}i"
    if isinstance(value, (list, tuple, set, frozenset, range, Iterator)):
        items = list(itertools.islice(iter(value), _MAX_LISTED_ITEMS + 1))
        shown = ", ".join(_calculator_result(item) for item in items[:_MAX_LISTED_ITEMS])
        if len(items) > _MAX_LISTED_ITEMS:
            hidden = f"{len(value) - _MAX_LISTED_ITEMS:,} more" if isinstance(value, (list, tuple, set, frozenset, range)) else "more"
            shown = f"{shown}, … ({hidden})"
        return f"({shown})" if isinstance(value, tuple) else f"[{shown}]"
    return str(value)


def _exact_integer(value: int) -> str:
    if value.bit_length() <= _MAX_EXACT_INTEGER_BITS:
        return str(value)
    magnitude = abs(value)
    shift = max(int(math.log10(magnitude)) - 12, 0)
    head = str(magnitude // 10**shift)
    digits = len(head) + shift
    mantissa = f"{head[0]}.{head[1:13]}".rstrip("0").rstrip(".")
    sign = "-" if value < 0 else ""
    return f"{sign}{mantissa}e+{digits - 1} (an integer with {digits:,} digits)"


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 notebook_variables() -> dict[str, Any]:
    """Names this worker's notebook defined, which calculator calls in the same conversation can use."""
    return {key: value for key, value in _worker_namespace().items() if not key.startswith("_") and key != "math"}


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)
            if "expression" in request:
                output = calculate(str(request.get("expression") or ""), notebook_variables())
            else:
                output = run_cell(str(request.get("code") or ""))
            sys.stdout.write(json.dumps({"ok": True, "output": output}, ensure_ascii=False) + "\n")
        except Exception as error:
            detail = str(error) or type(error).__name__
            sys.stdout.write(
                json.dumps({"ok": False, "output": detail}, ensure_ascii=False) + "\n"
            )
        sys.stdout.flush()


if __name__ == "__main__":
    _worker_loop()