"""Agent-based task execution pipeline that wraps PatchPal to provide autonomous task completion capabilities."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/04_pipelines.agent.base.ipynb.

# %% auto #0
__all__ = ['AgentExecutor']

# %% ../../../nbs/04_pipelines.agent.base.ipynb #3a4d43ce
"""Agent executor that wraps patchpal-sandbox for sandboxed autonomous execution."""

import inspect
import os
import subprocess
import tempfile
import textwrap
from pathlib import Path
from typing import Callable, Dict, List, Optional

__all__ = ['AgentExecutor']


class AgentExecutor:
    """
    Sandboxed agent executor using patchpal-sandbox.

    This executor is a Python wrapper around patchpal-sandbox, providing:
    - Resource isolation and security via Docker/Podman containers
    - Support for both cloud and local LLMs
    - Custom tool integration via ~/.patchpal/tools/
    - API key management via .env files

    Args:
        model (str): LiteLLM model identifier (e.g., 'anthropic/claude-sonnet-4-5', 'ollama_chat/gpt-oss-120b')
        agent_type (str): Type of agent ('function_calling' or 'react')
        max_iterations (int): Maximum number of autopilot iterations
        env_file (str): Path to .env file with API keys
        sandbox (bool): Run in container sandbox (default: False). Set True for isolated/secure execution.
        image (str): Container image to use (default: python:3.11-slim) [sandbox only]
        network (str): Network mode ('bridge', 'host', 'none') [sandbox only]
        memory (str): Memory limit (e.g., '2g', '4g') [sandbox only]
        cpus (float): CPU limit (e.g., 2, 4) [sandbox only]
        custom_tools (list): List of Python functions to use as custom tools. Each function should have
                            type hints and a docstring. These will be written to .patchpal/tools/ directory
                            in the working directory and automatically discovered by PatchPal.
                            
                            Example:
                                def calculate_sum(a: int, b: int) -> int:
                                    '''Add two numbers together.
                                    
                                    Args:
                                        a: First number
                                        b: Second number
                                    
                                    Returns:
                                        Sum of a and b
                                    '''
                                    return a + b
                                
                                executor = AgentExecutor(
                                    model='anthropic/claude-sonnet-4-5',
                                    custom_tools=[calculate_sum]
                                )
                            
        enabled_tools (list): List of tool names to enable. If None, uses DEFAULT_TOOLS:
                             ['read_file', 'read_lines', 'edit_file', 'write_file',
                              'grep', 'find', 'run_shell', 'web_search', 'web_fetch']
                             Pass an empty list [] to use all available patchpal tools.
        disable_shell (bool): If True, remove 'run_shell' from default tools (default: False).
                             Only applies when enabled_tools=None. Useful for security when you
                             want defaults but no shell access.
        completion_promise (str): String that signals task completion (default: 'COMPLETE')
        verbose (bool): Show progress output in real-time (default: True). When False, output is
                       captured but not displayed (useful for automated scripts/logging).
    """
    
    # Default tools for code-focused agent tasks
    DEFAULT_TOOLS = [
        'read_file',      # Read complete files
        'read_lines',     # Read specific line ranges
        'edit_file',      # Edit files via find/replace
        'write_file',     # Write complete files
        'grep',           # Search for patterns in files
        'find',           # Find files by glob pattern
        'run_shell',      # Execute shell commands
        'web_search',     # Search the web for information
        'web_fetch',      # Fetch content from URLs
    ]

    def __init__(
        self,
        model: str,
        agent_type: str = "function_calling",
        max_iterations: int = 50,
        env_file: Optional[str] = None,
        sandbox: bool = False,
        image: str = "python:3.11-slim",
        network: str = "bridge",
        memory: Optional[str] = None,
        cpus: Optional[float] = None,
        custom_tools: Optional[List[Callable]] = None,
        enabled_tools: Optional[List[str]] = None,
        disable_shell: bool = False,
        completion_promise: str = "COMPLETE",
        verbose: bool = True,
    ):
        # Normalize ollama/ to ollama_chat/ for LiteLLM compatibility
        if model.startswith("ollama/"):
            model = model.replace("ollama/", "ollama_chat/", 1)

        self.model = model
        self.agent_type = agent_type
        self.max_iterations = max_iterations
        self.env_file = env_file
        self.sandbox = sandbox
        self.image = image
        self.network = network
        self.memory = memory
        self.cpus = cpus
        self.custom_tools = custom_tools or []
        
        # Handle tool configuration
        if enabled_tools is not None:
            # User explicitly specified tools
            self.enabled_tools = enabled_tools
        else:
            # Use defaults, but remove run_shell if disable_shell=True
            self.enabled_tools = self.DEFAULT_TOOLS.copy()
            if disable_shell and 'run_shell' in self.enabled_tools:
                self.enabled_tools.remove('run_shell')
        
        self.completion_promise = completion_promise
        self.verbose = verbose

    @classmethod
    def print_default_tools(cls):
        """
        Pretty-print the default tools available in AgentExecutor.
        
        This shows the tools that are used when enabled_tools=None (the default).
        Users can customize by passing a different list to enabled_tools parameter.
        """
        print("=" * 70)
        print("AgentExecutor Default Tools")
        print("=" * 70)
        print("\nThese tools are used by default when enabled_tools=None:\n")
        
        tool_descriptions = {
            'read_file': 'Read complete file contents',
            'read_lines': 'Read specific line ranges from files',
            'edit_file': 'Edit files via find/replace',
            'write_file': 'Write complete file contents',
            'grep': 'Search for patterns in files',
            'find': 'Find files by glob pattern',
            'run_shell': 'Execute shell commands',
            'web_search': 'Search the web for information',
            'web_fetch': 'Fetch and read content from URLs',
        }
        
        for i, tool in enumerate(cls.DEFAULT_TOOLS, 1):
            desc = tool_descriptions.get(tool, 'No description available')
            print(f"  {i:2d}. {tool:15s} - {desc}")
        
        print("\n" + "=" * 70)
        print("Customization Examples:")
        print("=" * 70)
        print("\n# Use defaults (all tools including shell):")
        print("executor = AgentExecutor(model='anthropic/claude-sonnet-4-5')")
        print("\n# Defaults but no shell access (safer):")
        print("executor = AgentExecutor(")
        print("    model='openai/gpt-5-mini',")
        print("    disable_shell=True")
        print(")")
        print("\n# Minimal tools:")
        print("executor = AgentExecutor(")
        print("    model='openai/gpt-5-mini',")
        print("    enabled_tools=['read_file', 'write_file']")
        print(")")
        print("\n# Web research only:")
        print("executor = AgentExecutor(")
        print("    model='openai/gpt-5-mini',")
        print("    enabled_tools=['web_search', 'web_fetch']")
        print(")")
        print()

    def _write_custom_tools(self, tools_dir: Path):
        """Write custom tool functions to Python files in the tools directory."""
        if not self.custom_tools:
            return
        
        tools_dir.mkdir(parents=True, exist_ok=True)
        
        for func in self.custom_tools:
            # Get function name
            func_name = func.__name__
            
            # Try to get source code
            try:
                source = inspect.getsource(func)
            except (OSError, TypeError):
                if self.verbose:
                    print(f"⚠️  Warning: Could not extract source for custom tool '{func_name}', skipping")
                continue
            
            # Dedent the source code
            source = textwrap.dedent(source)
            
            # Try to get imports from the source file (works for .py files)
            imports = []
            try:
                source_file = inspect.getsourcefile(func)
                if source_file and Path(source_file).exists():
                    with open(source_file, 'r') as f:
                        for line in f:
                            stripped = line.strip()
                            if stripped.startswith('import ') or stripped.startswith('from '):
                                if 'onprem' not in stripped:
                                    imports.append(stripped)
            except Exception:
                pass
            
            # If no imports found (e.g., Jupyter notebook), reconstruct from globals
            if not imports:
                func_globals = func.__globals__
                
                # Parse the source to find what names are actually used
                import ast
                try:
                    # Parse the function source to find referenced names
                    tree = ast.parse(source)
                    used_names = set()
                    for node in ast.walk(tree):
                        if isinstance(node, ast.Name):
                            used_names.add(node.id)
                except:
                    # If parsing fails, we'll include everything (fallback)
                    used_names = None
                
                # Group objects by their source module
                from_imports = {}  # module -> [names]
                direct_imports = set()
                
                # Skip these modules/names
                skip_modules = {'IPython', '__main__', 'builtins', '__builtin__', 'onprem', 'io'}
                skip_names = {'open', 'exit', 'quit', 'get_ipython', '__builtins__'}
                
                for name, obj in func_globals.items():
                    if name.startswith('_') or name in skip_names:
                        continue
                    
                    # Only include if the name is actually used in the function
                    if used_names is not None and name not in used_names:
                        continue
                    
                    # Skip onprem imports
                    if hasattr(obj, '__module__'):
                        obj_module = str(obj.__module__)
                        if any(skip in obj_module for skip in skip_modules):
                            continue
                    
                    # Check if it's a module (import xxx)
                    if hasattr(obj, '__name__') and hasattr(obj, '__file__'):
                        module_name = obj.__name__
                        if module_name and not any(skip in module_name for skip in skip_modules):
                            direct_imports.add(f"import {module_name}")
                    
                    # Check if it's from typing module
                    elif str(type(obj)).startswith("<class 'typing."):
                        if 'typing' not in from_imports:
                            from_imports['typing'] = []
                        from_imports['typing'].append(name)
                    
                    # Check if it's from another module (from xxx import yyy)
                    elif hasattr(obj, '__module__'):
                        module = obj.__module__
                        if module and not any(skip in module for skip in skip_modules):
                            # Only include standard library or well-known packages
                            if '.' not in name:  # Skip complex nested objects
                                if module not in from_imports:
                                    from_imports[module] = []
                                from_imports[module].append(name)
                
                # Build import statements
                for module, names in sorted(from_imports.items()):
                    imports.append(f"from {module} import {', '.join(sorted(set(names)))}")
                
                imports.extend(sorted(direct_imports))
            
            # Build file content with imports at the top
            file_content = ""
            if imports:
                # Remove duplicates while preserving order
                seen = set()
                unique_imports = []
                for imp in imports:
                    if imp not in seen:
                        seen.add(imp)
                        unique_imports.append(imp)
                file_content = "\n".join(unique_imports) + "\n\n"
            file_content += source
            
            # Write to file
            tool_file = tools_dir / f"{func_name}.py"
            tool_file.write_text(file_content)
            
            if self.verbose:
                print(f"✓ Wrote custom tool '{func_name}' to {tool_file}")

    def _build_sandbox_command(self, task_file: Path) -> List[str]:
        """Build the patchpal-sandbox command."""
        cmd = ['patchpal-sandbox']
        
        # Add sandbox options
        if self.image != "python:3.11-slim":
            cmd.extend(['--image', self.image])
        
        if self.network == 'host':
            cmd.append('--host-network')
        elif self.network == 'none':
            cmd.append('--no-network')
        # 'bridge' is default, no flag needed
        
        if self.memory:
            cmd.extend(['--memory', self.memory])
        
        if self.cpus:
            cmd.extend(['--cpus', str(self.cpus)])
        
        if self.env_file:
            cmd.extend(['--env-file', self.env_file])
        
        # Separator between sandbox args and patchpal args
        cmd.append('--')
        
        # Add patchpal autopilot command
        cmd.append('autopilot')
        cmd.extend(['--model', self.model])
        cmd.extend(['--prompt-file', str(task_file)])
        cmd.extend(['--completion-promise', self.completion_promise])
        cmd.extend(['--max-iterations', str(self.max_iterations)])
        
        return cmd
    
    def _build_direct_command(self, task_file: Path) -> List[str]:
        """Build the direct patchpal-autopilot command (no sandbox)."""
        cmd = ['patchpal-autopilot']
        cmd.extend(['--model', self.model])
        cmd.extend(['--prompt-file', str(task_file)])
        cmd.extend(['--completion-promise', self.completion_promise])
        cmd.extend(['--max-iterations', str(self.max_iterations)])

        return cmd
    
    def run(
        self,
        task: str,
        working_dir: Optional[str] = None,
    ) -> str:
        """
        Run the agent on a given task in a sandboxed environment.
        
        Args:
            task (str): The task description/prompt
            working_dir (str): Directory to run in (default: current directory)
            
        Returns:
            str: The agent's response/output
            
        Raises:
            RuntimeError: If agent execution fails
            FileNotFoundError: If patchpal-sandbox is not installed
        """
        # Save current directory
        original_dir = os.getcwd()
        
        # Track if we created a custom tools directory
        custom_tools_dir = None
        
        try:
            # Change to working directory if specified
            if working_dir:
                working_path = Path(working_dir)
                working_path.mkdir(parents=True, exist_ok=True)
                os.chdir(working_dir)
            
            # Write custom tools to .patchpal/tools/ directory
            if self.custom_tools:
                custom_tools_dir = Path('.patchpal/tools')
                self._write_custom_tools(custom_tools_dir)
            
            # Create task file in current directory
            task_file = Path('task_prompt.md')
            
            # Add completion promise instruction
            task_with_completion = (
                f"{task}\n\n"
                f"When you have successfully completed the task, output the following exactly:\n"
                f"<promise>{self.completion_promise}</promise>"
            )
            task_file.write_text(task_with_completion)
            
            try:
                # Set environment variables for agent configuration
                env = os.environ.copy()

                # Skip autopilot confirmation prompt (we're running non-interactively)
                env['PATCHPAL_AUTOPILOT_CONFIRMED'] = 'true'

                # Disable token-by-token streaming but still show action summaries
                # (e.g., "🤔 Thinking...", "🌐 Searching web:", etc.)
                env['PATCHPAL_STREAM_OUTPUT'] = 'false'

                # Set agent type
                if self.agent_type == "react":
                    env['PATCHPAL_REACT_MODE'] = 'true'

                # Set enabled tools if specified
                if self.enabled_tools:
                    env['PATCHPAL_ENABLED_TOOLS'] = ','.join(self.enabled_tools)

                # Load .env file if specified (for non-sandboxed execution)
                if self.env_file and not self.sandbox:
                    # Load environment variables from .env file
                    env_path = Path(self.env_file)
                    if env_path.exists():
                        with open(env_path) as f:
                            for line in f:
                                line = line.strip()
                                if line and not line.startswith('#') and '=' in line:
                                    key, value = line.split('=', 1)
                                    env[key.strip()] = value.strip()

                # Build command (sandbox or direct)
                if self.sandbox:
                    cmd = self._build_sandbox_command(task_file)
                else:
                    cmd = self._build_direct_command(task_file)

                #print(cmd)

                # Run patchpal-autopilot with real-time output
                process = subprocess.Popen(
                    cmd,
                    env=env,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                )

                output_lines = []
                for line in process.stdout:
                    if self.verbose:
                        print(line, end='')
                    output_lines.append(line)

                process.wait()

                if process.returncode != 0:
                    raise RuntimeError(
                        f"Agent execution failed with return code {process.returncode}"
                    )

                return ''.join(output_lines)
                    
            except FileNotFoundError:
                if self.sandbox:
                    raise FileNotFoundError(
                        "patchpal-sandbox command not found. Install PatchPal with: pip install onprem[agent]"
                    )
                else:
                    raise FileNotFoundError(
                        "patchpal command not found. Install PatchPal with: pip install onprem[agent]"
                    )
            except subprocess.TimeoutExpired:
                raise RuntimeError(f"Agent execution timed out after 1 hour")
            finally:
                # Clean up task file
                if task_file.exists():
                    task_file.unlink()
                
                # Clean up custom tools directory if we created it
                if custom_tools_dir and custom_tools_dir.exists():
                    import shutil
                    shutil.rmtree(custom_tools_dir)
        
        finally:
            # Restore original directory
            os.chdir(original_dir)
