"""
沙盒化 Python 代码解释器，供 MCP 工具调用。

通过 subprocess.run 在独立子进程中执行用户代码，具备以下安全措施：
  - 硬超时限制（默认 30 秒）
  - 输出长度截断（默认 8000 字符）
  - 进程隔离（避免影响主进程状态）

使用 subprocess 而非 multiprocessing.Process 的原因：
  在 daemon worker 环境中，multiprocessing.Process 会触发
  "daemonic processes are not allowed to have children" 错误。

预导入环境：
  数学/科学: math, cmath, decimal, fractions, statistics, numpy
  工具/容器: collections, itertools, functools, operator, bisect, heapq
  字符串/IO: string, re, json, io, sys, os, copy, typing
  常用符号:  sqrt, pi, e, inf, log, exp, sin, cos, tan, ceil, floor, factorial, gcd 等
"""

import ast
import subprocess
import sys
import textwrap

# 注入到用户代码前面的预导入模块代码
_CODE_PREAMBLE = textwrap.dedent("""\
    import math, cmath, decimal, fractions, statistics, random
    import string, re, json, collections, itertools, functools, operator
    import io, sys, os, copy, bisect, heapq, typing
    import numpy as np
    from math import (sqrt, pi, e, inf, nan, log, log2, log10, exp,
                      sin, cos, tan, asin, acos, atan, atan2,
                      ceil, floor, factorial, gcd, comb, perm,
                      isclose, isinf, isnan)
    from fractions import Fraction
    from decimal import Decimal
    sys.setrecursionlimit(10**5)
""")

# 单次执行的输出上限（字符数），超出部分被截断
_MAX_OUTPUT_CHARS = 8_000

_FENCE_RE = None  # 延迟编译的正则表达式


def _strip_markdown_fences(code: str) -> str:
    """去除 LLM 有时会添加的 markdown 代码围栏（```python ... ```）。

    仅当代码首尾均有围栏时才剥离，避免误伤正常代码。

    Args:
        code: 原始代码字符串。

    Returns:
        str: 去除围栏后的纯代码。
    """
    import re
    global _FENCE_RE
    if _FENCE_RE is None:
        _FENCE_RE = re.compile(
            r"^\s*```[\w]*\s*\n(.*?)\n\s*```\s*$", re.DOTALL
        )
    m = _FENCE_RE.match(code)
    return m.group(1) if m else code


def _auto_print_last_expr(code: str) -> str:
    """若代码最后一条语句是裸表达式，自动包装为 print(repr(...))。

    模拟交互式 REPL 行为：即使用户忘记 print()，也能返回表达式的值。
    正确处理同一行上分号分隔的多条语句。

    Args:
        code: 用户代码字符串。

    Returns:
        str: 可能被修改后的代码（最后的裸表达式被包装为 print）。
    """
    try:
        tree = ast.parse(code)
    except SyntaxError:
        return code
    if not tree.body:
        return code
    last = tree.body[-1]
    if not isinstance(last, ast.Expr):
        return code

    lines = code.splitlines(keepends=True)

    pre = "".join(lines[: last.lineno - 1])

    same_line_prefix = lines[last.lineno - 1][: last.col_offset]
    if same_line_prefix.strip():
        pre += same_line_prefix.rstrip() + "\n"

    expr_src = ast.get_source_segment(code, last.value)
    if expr_src is None:
        expr_src = "".join(lines[last.lineno - 1 :]).strip()
        if same_line_prefix.strip():
            expr_src = expr_src[len(same_line_prefix):].strip().lstrip(";").strip()

    return f"{pre}_result_ = {expr_src}\nif _result_ is not None:\n    print(repr(_result_))\n"


def execute_python(code: str, timeout: int = 30) -> str:
    """在子进程中执行 Python 代码并返回输出。

    Args:
        code: 待执行的 Python 代码字符串。
        timeout: 执行超时时间（秒），超时后子进程被终止。

    Returns:
        str: 执行结果，格式为 stdout（+ 可选的 [stderr] 区段）。
             超时返回错误提示，无输出返回 "(no output)"。
    """
    code = _strip_markdown_fences(code)
    full_code = _CODE_PREAMBLE + _auto_print_last_expr(code)

    try:
        result = subprocess.run(
            [sys.executable, "-c", full_code],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
    except subprocess.TimeoutExpired:
        return f"[Error] Code execution timed out after {timeout}s."
    except Exception as e:
        return f"[Error] Failed to launch subprocess: {e}"

    stdout = result.stdout or ""
    stderr = result.stderr or ""

    if len(stdout) > _MAX_OUTPUT_CHARS:
        stdout = stdout[:_MAX_OUTPUT_CHARS] + "\n...(output truncated)..."
    if len(stderr) > _MAX_OUTPUT_CHARS:
        stderr = stderr[:_MAX_OUTPUT_CHARS] + "\n...(error truncated)..."

    parts = []
    if stdout:
        parts.append(stdout)
    if stderr:
        parts.append(f"[stderr]\n{stderr}")
    return "\n".join(parts) if parts else "(no output)"
