import io
import sys
from typing import List

from langroid.agent.tool_message import ToolMessage


class RunPythonCodeTool(ToolMessage):
    """
    Tool/function to run code generated by the LLM.
    The code is assumed to be self-contained, i.e. it contains all necessary imports
    and does not depend on any external variables.
    """

    request: str = "run_python_code"
    purpose: str = """
            To run python <code> and return the results to answer a question.
            """
    code: str

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        return [
            cls(code="import numpy as np\nnp.square(9)"),
        ]

    def handle(self) -> str:
        """
        Handle a RunPythonCodeTool message by running the code and returning the result.
        Returns:
            str: The result of running the code along with any print output.
        """
        code = self.code

        # Create a string-based I/O stream
        code_out = io.StringIO()

        # Temporarily redirect standard output to our string-based I/O stream
        sys.stdout = code_out

        try:
            eval_result = eval(code, {})
        except Exception as e:
            eval_result = f"ERROR: {type(e)}: {e}"

        if eval_result is None:
            eval_result = ""

        # Always restore the original standard output
        sys.stdout = sys.__stdout__

        # Get the resulting string from the I/O stream
        print_result = code_out.getvalue() or ""
        sep = "\n" if print_result else ""
        # Combine the print and eval results
        result = f"{print_result}{sep}{eval_result}"
        if result == "":
            result = "No result"
        # Return the result
        return result
