from pydantic import create_model
import inspect, json
from inspect import Parameter
import ast
from transformers import Tool


class PythonInterpreter(Tool):
    name = "python_interpreter"
    description = (
        "This is a tool that runs python code. It takes the python code and returns the result of the python code."
    )

    inputs = ["text"]
    outputs = ["text"]
    
    def run(code):
        tree = ast.parse(code)
        last_node = tree.body[-1] if tree.body else None
    
        # If the last node is an expression, modify the AST to capture the result
        if isinstance(last_node, ast.Expr):
            tgts = [ast.Name(id='_result', ctx=ast.Store())]
            assign = ast.Assign(targets=tgts, value=last_node.value)
            tree.body[-1] = ast.fix_missing_locations(assign)

        ns = {}
        exec(compile(tree, filename='<ast>', mode='exec'), ns)
        return ns.get('_result', None)
    
    def python(code: str):
        #Return result of executing `code` using python. If execution not permitted, returns `#FAIL#`
        go = input(f'Proceed with execution?\n```\n{code}\n```\n')
        if go.lower()!='y': return '#FAIL#'
        return run(code)
    
    def __call__(self, task: str):
        return python(task)