"""Python execution tool implementation."""

from __future__ import annotations

import asyncio
import contextlib
import logging
import os
import sys
import tempfile
from pathlib import Path
from typing import Any

from hermes.tools.base.tool import BaseTool, ToolSchema

logger = logging.getLogger(__name__)


class PythonExecutionTool(BaseTool):
    """Tool for executing Python code with safety measures."""

    def __init__(self, timeout: int = 30, max_output_size: int = 10000, max_memory_mb: int = 256) -> None:
        super().__init__()
        self.timeout = timeout
        self.max_output_size = max_output_size
        self.max_memory_mb = max_memory_mb

    def _define_schema(self) -> ToolSchema:
        return ToolSchema(
            name="python_executor",
            description="Execute Python code in an isolated environment with resource limits",
            parameters={
                "code": {
                    "type": "string",
                    "description": "Python code to execute",
                },
                "timeout": {
                    "type": "integer",
                    "description": "Execution timeout in seconds",
                    "default": 30,
                },
            },
            required=["code"],
            category="execution",
            tags=["python", "execute", "code"],
        )

    async def execute(self, **kwargs: Any) -> dict[str, Any]:
        """Execute Python code with isolation."""
        code = kwargs["code"]
        timeout = kwargs.get("timeout", self.timeout)

        restricted_code = self._inject_safety(code)

        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".py", delete=False, encoding="utf-8"
        ) as f:
            f.write(restricted_code)
            temp_path = f.name

        try:
            env = os.environ.copy()
            sensitive_keys = [k for k in env if any(word in k.upper() for word in
                ["API_KEY", "APIKEY", "SECRET", "TOKEN", "PASSWORD", "CREDENTIAL", "AUTH"])]
            for k in sensitive_keys:
                del env[k]

            process = await asyncio.create_subprocess_exec(
                sys.executable,
                "-I",
                temp_path,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=env,
                cwd=tempfile.gettempdir(),
            )

            try:
                stdout, stderr = await asyncio.wait_for(
                    process.communicate(), timeout=timeout
                )
            except TimeoutError:
                process.kill()
                await process.wait()
                return {
                    "success": False,
                    "error": f"Execution timed out after {timeout}s",
                    "stdout": "",
                    "stderr": "",
                    "return_code": -1,
                }

            stdout_str = stdout.decode("utf-8", errors="replace")[:self.max_output_size]
            stderr_str = stderr.decode("utf-8", errors="replace")[:self.max_output_size]

            return {
                "success": process.returncode == 0,
                "stdout": stdout_str,
                "stderr": stderr_str,
                "return_code": process.returncode,
            }

        except Exception as e:
            logger.error(f"Python execution error: {e}")
            return {
                "success": False,
                "error": str(e),
                "stdout": "",
                "stderr": "",
                "return_code": -1,
            }
        finally:
            with contextlib.suppress(Exception):
                Path(temp_path).unlink(missing_ok=True)

    def _inject_safety(self, code: str) -> str:
        """Inject safety restrictions into code — blocks imports and subprocess access."""
        banner = """import sys as _sys, builtins as _builtins

# Block __import__ — prevents all module imports
_builtins.__import__ = None
# Nullify dangerous modules so even if __import__ is somehow recovered they fail
for _mod in ['os', 'subprocess', 'shutil', 'socket', 'ctypes', 'pickle',
             'marshal', 'compile', 'codeop', 'code', 'inspect',
             'importlib', 'pkgutil', 'mmap', 'webbrowser', 'antigravity',
             'ctypes', '_ctypes', '_winapi', 'posix', 'grp', 'pwd']:
    _sys.modules[_mod] = None
del _sys, _builtins
"""
        return banner + code