"""Restricted Python executor: the tool the ReVeR judge calls for mechanical checks (counting, string and format
tests). The code runs in a separate process with a time limit, a memory limit, a whitelist of builtins and modules,
and the task (instruction, response, rubrics) bound as variables. The reply is a JSON string with the keys ok,
stdout, result and error; `result` is the repr of a variable named `result`, if the code sets one."""
from __future__ import annotations

import ast
import contextlib
import json
import multiprocessing as mp
import re
import reprlib
import signal
import time
import traceback
from dataclasses import dataclass
from multiprocessing.connection import Connection
from typing import Any


DISALLOWED_NODES = (
    ast.Global,
    ast.Nonlocal,
    ast.ClassDef,
)

ALLOWED_IMPORTS = {
    "collections",
    "datetime",
    "itertools",
    "json",
    "math",
    "re",
    "statistics",
    "string",
    "textwrap",
    "unicodedata",
}

DISALLOWED_NAMES = {
    "__import__",
    "breakpoint",
    "compile",
    "delattr",
    "dir",
    "eval",
    "exec",
    "exit",
    "getattr",
    "globals",
    "help",
    "input",
    "locals",
    "open",
    "quit",
    "setattr",
    "vars",
}


@dataclass(frozen=True)
class ToolExecutionResult:
    ok: bool
    stdout: str = ""
    result: str | None = None
    error: str | None = None

    def to_json(self, output_limit: int) -> str:
        obj = {
            "ok": self.ok,
            "stdout": self.stdout[:output_limit],
            "result": self.result[:output_limit] if self.result is not None else None,
            "error": self.error[:output_limit] if self.error is not None else None,
        }
        return json.dumps(obj, ensure_ascii=False)


class _BoundedTextBuffer:
    """A minimal text stream that never retains more than ``limit`` chars."""

    def __init__(self, limit: int) -> None:
        self.limit = max(0, limit)
        self._parts: list[str] = []
        self._length = 0

    def write(self, value: str) -> int:
        text = str(value)
        remaining = self.limit - self._length
        if remaining > 0:
            chunk = text[:remaining]
            self._parts.append(chunk)
            self._length += len(chunk)
        return len(text)

    def flush(self) -> None:
        return None

    def getvalue(self) -> str:
        return "".join(self._parts)


def _bounded_repr(value: Any, output_limit: int) -> str:
    formatter = reprlib.Repr()
    formatter.maxlevel = 4
    formatter.maxdict = 20
    formatter.maxlist = 20
    formatter.maxtuple = 20
    formatter.maxset = 20
    formatter.maxfrozenset = 20
    formatter.maxdeque = 20
    formatter.maxarray = 20
    formatter.maxstring = output_limit
    formatter.maxlong = output_limit
    formatter.maxother = output_limit
    return formatter.repr(value)[:output_limit]


def _configure_worker(memory_limit_mb: int) -> None:
    """Install worker-only signal and address-space limits."""

    signal.signal(signal.SIGINT, signal.SIG_DFL)
    signal.signal(signal.SIGTERM, signal.SIG_DFL)

    import resource

    memory_limit_bytes = memory_limit_mb * 1024 * 1024
    resource.setrlimit(resource.RLIMIT_AS, (memory_limit_bytes, memory_limit_bytes))


def _send_worker_result(connection: Connection, result: ToolExecutionResult, output_limit: int) -> None:
    try:
        connection.send(
            {
                "ok": result.ok,
                "stdout": result.stdout[:output_limit],
                "result": result.result[:output_limit] if result.result is not None else None,
                "error": result.error[:output_limit] if result.error is not None else None,
            }
        )
    except (BrokenPipeError, EOFError, OSError):
        pass
    finally:
        connection.close()


def strip_code_fence(code: str) -> str:
    code = (code or "").strip()
    code = re.sub(r"^\s*```(?:python)?\s*", "", code, flags=re.IGNORECASE)
    code = re.sub(r"\s*```\s*$", "", code)
    return code.strip()


def validate_code(code: str) -> None:
    tree = ast.parse(code, mode="exec")
    for node in ast.walk(tree):
        if isinstance(node, DISALLOWED_NODES):
            raise ValueError(f"Disallowed Python syntax: {node.__class__.__name__}")
        if isinstance(node, ast.Import):
            for alias in node.names:
                root = alias.name.split(".", 1)[0]
                if alias.name not in ALLOWED_IMPORTS or root not in ALLOWED_IMPORTS:
                    raise ValueError(f"Disallowed import: {alias.name}")
                if alias.asname and (alias.asname in DISALLOWED_NAMES or alias.asname.startswith("__")):
                    raise ValueError(f"Disallowed import alias: {alias.asname}")
        if isinstance(node, ast.ImportFrom):
            module = node.module or ""
            if node.level != 0 or module not in ALLOWED_IMPORTS:
                raise ValueError(f"Disallowed import: {module}")
            for alias in node.names:
                if alias.name == "*" or alias.name.startswith("__"):
                    raise ValueError(f"Disallowed import name: {alias.name}")
                if alias.asname and (alias.asname in DISALLOWED_NAMES or alias.asname.startswith("__")):
                    raise ValueError(f"Disallowed import alias: {alias.asname}")
        if isinstance(node, ast.Name):
            if node.id in DISALLOWED_NAMES or node.id.startswith("__"):
                raise ValueError(f"Disallowed name: {node.id}")
        if isinstance(node, ast.Attribute) and node.attr.startswith("__"):
            raise ValueError(f"Disallowed attribute: {node.attr}")


def _safe_worker(
    code: str,
    context: dict[str, Any],
    connection: Connection,
    output_limit: int,
    memory_limit_mb: int,
) -> None:
    import collections
    import itertools
    import math
    import re as re_module
    import statistics

    try:
        _configure_worker(memory_limit_mb)
    except Exception:
        _send_worker_result(
            connection,
            ToolExecutionResult(
                ok=False,
                error="failed to configure Python tool process limits:\n" + traceback.format_exc(limit=4),
            ),
            output_limit,
        )
        return

    safe_modules = {
        "collections": collections,
        "itertools": itertools,
        "json": json,
        "math": math,
        "re": re_module,
        "statistics": statistics,
    }

    def restricted_import(name, globals=None, locals=None, fromlist=(), level=0):
        if level != 0:
            raise ImportError("relative imports are not allowed")
        if name not in safe_modules:
            raise ImportError(f"import of {name!r} is not allowed")
        return safe_modules[name]

    safe_builtins = {
        "abs": abs,
        "all": all,
        "any": any,
        "bool": bool,
        "dict": dict,
        "enumerate": enumerate,
        "filter": filter,
        "float": float,
        "int": int,
        "isinstance": isinstance,
        "len": len,
        "list": list,
        "map": map,
        "max": max,
        "min": min,
        "print": print,
        "range": range,
        "repr": repr,
        "reversed": reversed,
        "round": round,
        "set": set,
        "sorted": sorted,
        "str": str,
        "sum": sum,
        "tuple": tuple,
        "zip": zip,
        "__import__": restricted_import,
    }
    namespace = {
        "__builtins__": safe_builtins,
        "collections": collections,
        "itertools": itertools,
        "json": json,
        "math": math,
        "re": re_module,
        "statistics": statistics,
    }
    namespace.update(context)
    stdout = _BoundedTextBuffer(output_limit)
    try:
        validate_code(code)
        compiled = compile(code, "<qwen_agent_python_tool>", "exec")
        with contextlib.redirect_stdout(stdout):
            exec(compiled, namespace, namespace)
        result = None
        if "result" in namespace:
            result = _bounded_repr(namespace["result"], output_limit)
        payload = ToolExecutionResult(ok=True, stdout=stdout.getvalue(), result=result)
    except MemoryError:
        payload = ToolExecutionResult(
            ok=False,
            stdout=stdout.getvalue(),
            error=f"memory limit exceeded ({memory_limit_mb} MiB)",
        )
    except Exception:
        payload = ToolExecutionResult(
            ok=False,
            stdout=stdout.getvalue(),
            error=traceback.format_exc(limit=4),
        )
    _send_worker_result(connection, payload, output_limit)


def _terminate_process(process: mp.Process, grace_seconds: float) -> str:
    """Stop and reap a worker, escalating from SIGTERM to SIGKILL."""

    process.terminate()
    process.join(grace_seconds)
    if not process.is_alive():
        return "SIGTERM"

    process.kill()
    process.join(grace_seconds)
    if process.is_alive():
        raise RuntimeError(f"failed to reap Python tool process pid={process.pid} after SIGKILL")
    return "SIGKILL"


class RestrictedPythonExecutor:
    def __init__(
        self,
        *,
        timeout_seconds: float = 3.0,
        output_limit: int = 4000,
        memory_limit_mb: int = 512,
        termination_grace_seconds: float = 1.0,
    ) -> None:
        if timeout_seconds <= 0:
            raise ValueError("timeout_seconds must be positive")
        if output_limit <= 0:
            raise ValueError("output_limit must be positive")
        if type(memory_limit_mb) is not int or memory_limit_mb <= 0:
            raise ValueError("memory_limit_mb must be a positive integer")
        if termination_grace_seconds <= 0:
            raise ValueError("termination_grace_seconds must be positive")
        self.timeout_seconds = timeout_seconds
        self.output_limit = output_limit
        self.memory_limit_mb = memory_limit_mb
        self.termination_grace_seconds = termination_grace_seconds

    def run(self, code: str, context: dict[str, Any]) -> str:
        code = strip_code_fence(code)
        if not code:
            return ToolExecutionResult(ok=False, error="empty code").to_json(self.output_limit)
        if len(code) > 12000:
            return ToolExecutionResult(ok=False, error="code is too long").to_json(self.output_limit)

        ctx = mp.get_context("spawn")
        receive_connection, send_connection = ctx.Pipe(duplex=False)
        process = ctx.Process(
            target=_safe_worker,
            args=(code, context, send_connection, self.output_limit, self.memory_limit_mb),
        )
        try:
            process.start()
            send_connection.close()
            process.join(self.timeout_seconds)
            if process.is_alive():
                signal_used = _terminate_process(process, self.termination_grace_seconds)
                return ToolExecutionResult(
                    ok=False,
                    error=f"execution timed out after {self.timeout_seconds}s; worker reaped with {signal_used}",
                ).to_json(self.output_limit)
            if receive_connection.poll(0.1):
                data = receive_connection.recv()
                return ToolExecutionResult(**data).to_json(self.output_limit)
            if process.exitcode is not None and process.exitcode < 0:
                signal_name = signal.Signals(-process.exitcode).name
                return ToolExecutionResult(
                    ok=False,
                    error=f"execution terminated by signal {signal_name}",
                ).to_json(self.output_limit)
            return ToolExecutionResult(
                ok=False,
                error=f"execution failed without returning a result (exitcode={process.exitcode})",
            ).to_json(self.output_limit)
        finally:
            send_connection.close()
            receive_connection.close()
            if process.is_alive():
                _terminate_process(process, self.termination_grace_seconds)
            if process.exitcode is not None:
                process.close()