"""
Python Executor - Persistent Python execution environment.
Similar to Jupyter kernel or smolagents LocalPythonExecutor.
"""

import io
import sys
from typing import Any, Dict

__all__ = ["PythonExecutor"]


class PythonExecutor:
    """
    Python execution environment similar to smolagents LocalPythonExecutor.
    Maintains persistent state and tool namespace across executions.
    """

    def __init__(self):
        self.namespace = {}
        self.reset_environment()

    def reset_environment(self):
        """Reset the execution environment."""
        # Start with basic Python builtins
        self.namespace = {
            "__builtins__": __builtins__,
        }
        # Add print output capture
        self.namespace["_print_outputs"] = []

        # Override print to capture outputs
        def captured_print(*args, **kwargs):
            output = " ".join(str(arg) for arg in args)
            self.namespace["_print_outputs"].append(output)
            print(*args, **kwargs)  # Also print to console

        self.namespace["print"] = captured_print

        # Pre-import commonly used libraries
        try:
            import pandas as pd
            import numpy as np
            import os
            from pathlib import Path
            self.namespace["pd"] = pd
            self.namespace["np"] = np
            self.namespace["os"] = os
            self.namespace["Path"] = Path

            # Add file I/O helper functions
            def write_text_file(filepath, content):
                """Write text content to a file."""
                with open(filepath, 'w') as f:
                    f.write(content)
                return f"File written to {filepath}"

            def write_dataframe_to_csv(df, filepath, **kwargs):
                """Write pandas DataFrame to CSV file."""
                df.to_csv(filepath, **kwargs)
                return f"DataFrame written to {filepath}"

            def write_dataframe_to_tsv(df, filepath, **kwargs):
                """Write pandas DataFrame to TSV file."""
                df.to_csv(filepath, sep='\t', **kwargs)
                return f"DataFrame written to {filepath}"

            def create_directory(dirpath):
                """Create directory if it doesn't exist."""
                Path(dirpath).mkdir(parents=True, exist_ok=True)
                return f"Directory created: {dirpath}"

            self.namespace["write_text_file"] = write_text_file
            self.namespace["write_dataframe_to_csv"] = write_dataframe_to_csv
            self.namespace["write_dataframe_to_tsv"] = write_dataframe_to_tsv
            self.namespace["create_directory"] = create_directory

        except ImportError as e:
            print(f"Warning: Could not import library: {e}")

    def send_functions(self, functions: Dict[str, Any]):
        """Inject functions into the execution namespace."""
        for func_name, func in functions.items():
            self.namespace[func_name] = func

    def send_tools(self, tools: Dict[str, Any]):
        """Legacy method for backward compatibility - redirects to send_functions."""
        self.send_functions(tools)

    def send_variables(self, variables: Dict[str, Any]):
        """Inject variables into the execution namespace."""
        if variables:
            print(f"📝 Injecting {len(variables)} variables into Python executor...")
            self.namespace.update(variables)

    def __call__(self, code: str) -> Any:
        """Execute code in the persistent namespace."""
        return self.execute(code)

    def execute(self, code: str) -> Any:
        """
        Execute Python code in the persistent environment.
        Similar to smolagents python execution.
        """
        try:
            # Clear previous print outputs
            self.namespace["_print_outputs"] = []

            # Capture stdout
            old_stdout = sys.stdout
            sys.stdout = captured_output = io.StringIO()

            try:
                # Execute the code in our persistent namespace
                exec(code, self.namespace)

                # Get any stdout output
                stdout_output = captured_output.getvalue()

                # Get print outputs from our custom print function
                print_outputs = self.namespace.get("_print_outputs", [])

                # Combine outputs
                all_outputs = []
                if stdout_output.strip():
                    all_outputs.append(stdout_output.strip())
                if print_outputs:
                    all_outputs.extend(print_outputs)

                result = "\n".join(all_outputs) if all_outputs else "Code executed successfully"

                return result

            finally:
                sys.stdout = old_stdout

        except Exception as e:
            return f"Error: {str(e)}"