from transformers import Tool
import ast
import io
import sys
from contextlib import redirect_stdout
import importlib

class PythonExecutorToolWithInputs(Tool):
    name = "python_executor_with_test_inputs"
    description = (
    "This tool executes Python code with given inputs. It takes a single input containing:"
    "Python code to execute."
    "Followed by a line containing only '#inputs' (without quotes)."
    "Followed by input values, one per line, to be used when the code calls input()."
    "For example if there are 2 input funtion calls in the code with first taking 2 arguments and second taking 1:"
    "12 43\n94"
    "The tool returns the outputs(if multiple or else single) of the executed code as a string."
    )
    inputs = ["text"]
    outputs = ["text"]

    def __call__(self, task: str):
        # Split code and inputs
        parts = task.split('#inputs')
        if len(parts) != 2:
            return "Error: Input must contain '#inputs' separator."

        code = parts[0].strip()
        inputss = parts[1].strip().split('\n')
        input_iterator = iter(inputss)

        def safe_input():
            try:
                return next(input_iterator)
            except StopIteration:
                raise EOFError("No more input available")

        # Set up a safe environment for code execution
        safe_globals = {
            'print': print,
            'input': safe_input,
            'len': len,
            'int': int,
            'float': float,
            'str': str,
            'list': list,
            'dict': dict,
            'set': set,
            'tuple': tuple,
            'range': range,
            'enumerate': enumerate,
            'zip': zip,
            'sum': sum,
            'min': min,
            'max': max,
            'abs': abs,
            'round': round,
        }

        # Pre-import allowed modules
        allowed_modules = ['math', 'random', 'datetime', 're','numpy','pandas']
        for module_name in allowed_modules:
            try:
                module = importlib.import_module(module_name)
                safe_globals[module_name] = module
            except ImportError:
                pass  # Skip if module is not available

        # Capture stdout
        output_buffer = io.StringIO()

        try:
            # Parse the code to check for syntax errors
            ast.parse(code)

            # Execute the code
            with redirect_stdout(output_buffer):
                exec(code, safe_globals)

            output = output_buffer.getvalue()
        except Exception as e:
            output = f"Error: {str(e)}"

        return output