from typing import Any, Dict
from smolagents.tools import Tool
import io
import contextlib
import traceback
import signal

class TimeoutException(Exception):
    pass

def timeout_handler(signum, frame):
    raise TimeoutException("Execution timed out.")

class PythonRunnerTool(Tool):
    name = "python_runner"
    description = "Executes Python code safely and returns stdout or errors."
    inputs = {'code': {'type': 'string', 'description': 'Python code to execute'}}
    output_type = "string"

    def forward(self, code: str) -> str:
        stdout_buffer = io.StringIO()
        safe_globals = {
            "__builtins__": {
                "print": print,
                "range": range,
                "len": len,
                "int": int,
                "float": float,
                "str": str,
                "list": list,
                "dict": dict,
                "sum": sum,
                "min": min,
                "max": max,
                "abs": abs,
            }
        }
        try:
            signal.signal(signal.SIGALRM, timeout_handler)
            signal.alarm(5)  # 5 second timeout

            with contextlib.redirect_stdout(stdout_buffer):
                exec(code, safe_globals, {})

            signal.alarm(0)
            return stdout_buffer.getvalue()

        except TimeoutException as e:
            return f"Error: {str(e)}"

        except Exception:
            return traceback.format_exc()