import io
import contextlib
from llama_index.core.tools import FunctionTool

def execute_python_code(file_path: str) -> str:
    """
    Executes Python code from a given file path and captures its output. It only works when provided the path to a .py script
    Args:
        file_path (str): The local path to the .py file to execute.
    """
    print(f"Executing Python code from: {file_path}")
    
    # Use io.StringIO to create an in-memory text buffer to capture output
    string_io = io.StringIO()
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            code = f.read()

        # Use contextlib.redirect_stdout to temporarily redirect all print() output to our in-memory buffer.
        with contextlib.redirect_stdout(string_io):
            exec(code, {})

        # Get the content that was "printed" from the buffer
        output = string_io.getvalue()
        
        if not output.strip():
             # Fallback for scripts that don't print but end with an expression
            try:
                lines = code.strip().split('\n')
                last_line = lines[-1]
                # Safely evaluate the last line if it's a simple expression
                output = str(eval(last_line, {"__builtins__": None}, {}))
            except:
                output = "Code executed successfully with no print output."
        
        return f"Successfully executed. Output:\n---\n{output.strip()}\n---"

    except Exception as e:
        return f"Error executing Python code: {e}"

def get_code_interpreter_tool() -> FunctionTool:
    """Initializes and returns the Python Code Interpreter tool."""
    return FunctionTool.from_defaults(
        fn=execute_python_code,
        name="python_code_interpreter",
        description="A tool that can execute Python code from a local file and return its output. It takes a 'file_path' to the Python file as input. Use this to find the output of a provided .py script."
    )