from smolagents.tools import Tool
import pdb
import smolagents
import io
import sys

class LocalPythonDebuggerTool(Tool):
    name = "local-python-debugger"
    description = """
    The Local Python Debugger Tool is designed to assist in debugging InterpreterErrors—specifically, errors related to "Code execution failed at line". 
    It is not intended for other types of InterpreterErrors, such as:
    - "Reached the max number of operations"
    - "Code parsing failed on line"

    The tool debugs code executed in the following manner:
    ```python
    expression = ast.parse(code)
    for node in expression.body:
        result = evaluate_ast(node, ...)
    ```
    where evaluate_ast processes the Abstract Syntax Tree (AST) by evaluating AST functions like: ast.Assign, ast.AugAssign, ast.Call, ast.Constant, ast.Name, ast.BinOp, etc.
    As a result, the debugger output reflects the evaluate_ast execution context rather than a direct Python runtime environment.
    
    Possible debugger commands:
    - bt - Shows the full stack trace, helping to understand where the error occurred in the call stack.
    - args - Shows the arguments passed to the function.
    - l - Lists the source code of the current file.
    - p {variable} - Prints the value of the variable. NOTE: variable names are different than in the code. They should be inspected first with `l` command.

    """
    inputs = {'code': {'type': 'string', 'description': 'code to debug'}, 'command': {'type': 'string', 'description': 'pdb command to execute'}}
    output_type = "string"

    def setup(self):
        from smolagents import LocalPythonInterpreter
        from smolagents.default_tools import FinalAnswerTool

        import io
        import pdb

        self.python_executor = LocalPythonInterpreter(
            [],
            {"final_answer": FinalAnswerTool()},
            max_print_outputs_length=None,
        )

        self.output_capture = io.StringIO()

        self.pdb_instance = pdb.Pdb(stdout=self.output_capture)

    def wrap_code_with_trycatch(self, code: str) -> str:
        from smolagents.utils import parse_code_blobs

        try:
            code = parse_code_blobs(code)
        except Exception as e:
            pass

        code_lines = code.split("\n")
        for i, line in enumerate(code_lines):
            code_lines[i] = f"\t{line}"
        code = "\n".join(code_lines)

        # add extra indentation
        code = (
            f"try:\n{code}\n" + "except Exception as e:" + "\n\t" + "out=sys.exc_info()"
        )
        code = "out=None" + "\n" + code + "\n" + "out"  # provide results in output

        return code

    def execute_code(self, code: str, state: dict):
        from smolagents.utils import parse_code_blobs

        code = parse_code_blobs(code)

        # Silencing the tool_validation error: execute_code: Name '' is undefined.
        py_executor_output = None
        py_executor_logs = None
        py_executor_is_final_answer = None

        py_executor_output, py_executor_logs, py_executor_is_final_answer = (
            self.python_executor(code, state)
        )

        return py_executor_output, py_executor_logs, py_executor_is_final_answer

    def execude_pdb_command(
        self, command: str, traceback
    ):
        output = None
        try:
            # Set up the pdb instance with the current frame and traceback
            frame = traceback.tb_frame
            while frame.f_back:
                frame = frame.f_back
            self.pdb_instance.reset()
            self.pdb_instance.setup(frame, traceback)
            # Execute the pdb command
            self.pdb_instance.onecmd(command)
            # Retrieve the output
            output = self.output_capture.getvalue()
        finally:
            # Clear the output capture for future use
            self.output_capture.truncate(0)
            self.output_capture.seek(0)

        return output

    def forward(self, code: str, command: str):
        import sys

        code = self.wrap_code_with_trycatch(code)

        state = {"sys": sys}

        # Silencing the tool_validation error: execute_code: Name '' is undefined.
        py_executor_output = None
        tb = None

        py_executor_output, py_executor_logs, py_executor_is_final_answer = (
            self.execute_code(code, state)
        )

        if py_executor_output is not None:
            try:
                exc_type, exc_value, tb = py_executor_output
            except Exception as e:
                return "No traceback found"
            pdb_result = self.execude_pdb_command(command, tb)
            return pdb_result

        return "No traceback found"

    def __init__(self, *args, **kwargs):
        self.is_initialized = False