"""
Python Execution Tool

This tool allows the agent to execute Python code in a controlled environment.
"""

import subprocess
import tempfile
import os
import sys
from typing import Dict, Any, List
from tools.tool_registry import BaseTool


class PythonTool(BaseTool):
    """Tool for executing Python code"""
    
    def __init__(self):
        super().__init__(
            name="python_tool",
            description="Execute Python code and return results. Supports data analysis, calculations, file operations, and more."
        )
        self.temp_dir = tempfile.mkdtemp()
    
    def execute(self, query: str, context: Dict[str, Any]) -> Any:
        """
        Execute Python code from the query
        
        Args:
            query: User query containing Python code or request
            context: Task analysis and other context
            
        Returns:
            Execution result including output and any errors
        """
        # Extract Python code from query
        python_code = self._extract_python_code(query)
        
        if not python_code:
            # If no explicit code, try to generate code based on the query
            python_code = self._generate_code_from_query(query)
        
        if not python_code:
            return "No Python code found or could be generated from the query."
        
        return self._execute_python_code(python_code)
    
    def _extract_python_code(self, query: str) -> str:
        """Extract Python code from query text"""
        # Look for code blocks
        import re
        
        # Check for code blocks with ```python
        python_blocks = re.findall(r'```python\n(.*?)\n```', query, re.DOTALL)
        if python_blocks:
            return python_blocks[0]
        
        # Check for code blocks with ```
        code_blocks = re.findall(r'```\n(.*?)\n```', query, re.DOTALL)
        if code_blocks:
            # Try to determine if it's Python code
            code = code_blocks[0]
            if any(keyword in code for keyword in ['print(', 'import ', 'def ', 'for ', 'if ', 'while ']):
                return code
        
        # Check for inline code with backticks
        inline_code = re.findall(r'`([^`]+)`', query)
        for code in inline_code:
            if any(keyword in code for keyword in ['print(', 'import ', '=']):
                return code
        
        return ""
    
    def _generate_code_from_query(self, query: str) -> str:
        """Generate Python code based on natural language query"""
        import re
        import math
        query_lower = query.lower()
        # Improved number extraction (supports negative and decimal numbers)
        numbers = re.findall(r'-?\d+(?:\.\d+)?', query)
        # Addition
        if any(word in query_lower for word in ['add', 'sum', 'plus', 'total', 'together']) and len(numbers) >= 2:
            return f"result = {numbers[0]} + {numbers[1]}\nprint(f'Result: {{result}}')"
        # Subtraction
        if any(word in query_lower for word in ['subtract', 'minus', 'difference', 'less']) and len(numbers) >= 2:
            return f"result = {numbers[0]} - {numbers[1]}\nprint(f'Result: {{result}}')"
        # Multiplication
        if any(word in query_lower for word in ['multiply', 'times', 'product', 'multiplied by']) and len(numbers) >= 2:
            return f"result = {numbers[0]} * {numbers[1]}\nprint(f'Result: {{result}}')"
        # Division
        if any(word in query_lower for word in ['divide', 'divided by', 'quotient', 'over']) and len(numbers) >= 2:
            return f"result = {numbers[0]} / {numbers[1]}\nprint(f'Result: {{result}}')"
        # Power
        if any(word in query_lower for word in ['power', 'raised to', 'exponent']) and len(numbers) >= 2:
            return f"result = {numbers[0]} ** {numbers[1]}\nprint(f'Result: {{result}}')"
        # Factorial
        if 'factorial' in query_lower and numbers:
            return "import math\nresult = math.factorial({})\nprint(f'Result: {{result}}')".format(numbers[0])
        # Square root
        if any(word in query_lower for word in ['square root', 'sqrt']) and numbers:
            return "import math\nresult = math.sqrt({})\nprint(f'Result: {{result}}')".format(numbers[0])
        # Hello World
        if 'hello world' in query_lower:
            return "print('Hello, World!')"
        # Current time
        if 'current time' in query_lower or 'time now' in query_lower:
            return "import datetime\nprint(f'Current time: {datetime.datetime.now()}')"
        # List files
        if 'list files' in query_lower:
            return "import os\nfiles = os.listdir('.')\nfor file in files:\n    print(file)"
        # Fallback
        return ""
    
    def _execute_python_code(self, code: str) -> Dict[str, Any]:
        """Execute Python code and capture output"""
        # Create a temporary file for the code
        temp_file = os.path.join(self.temp_dir, f"temp_code_{os.getpid()}.py")
        
        try:
            # Write code to temporary file
            with open(temp_file, 'w') as f:
                f.write(code)
            
            # Execute the code
            result = subprocess.run(
                [sys.executable, temp_file],
                capture_output=True,
                text=True,
                timeout=30,  # 30 second timeout
                cwd=self.temp_dir
            )
            
            execution_result = {
                "code": code,
                "stdout": result.stdout,
                "stderr": result.stderr,
                "return_code": result.returncode,
                "success": result.returncode == 0
            }
            
            # Clean up
            if os.path.exists(temp_file):
                os.remove(temp_file)
            
            return execution_result
            
        except subprocess.TimeoutExpired:
            return {
                "code": code,
                "stdout": "",
                "stderr": "Code execution timed out after 30 seconds",
                "return_code": -1,
                "success": False
            }
        except Exception as e:
            return {
                "code": code,
                "stdout": "",
                "stderr": f"Error executing code: {str(e)}",
                "return_code": -1,
                "success": False
            }
    
    def get_capabilities(self) -> List[str]:
        """Return capabilities of this tool"""
        return [
            "code_execution",
            "data_analysis", 
            "mathematical_computation",
            "file_processing",
            "text_processing",
            "automation"
        ]
    
    def validate_input(self, query: str, context: Dict[str, Any]) -> bool:
        """Validate input for Python execution"""
        # Basic safety checks
        dangerous_keywords = [
            'os.system',
            'subprocess.call',
            'eval(',
            'exec(',
            '__import__',
            'open(',  # Be careful with file operations
            'input(',  # Avoid interactive input
        ]
        
        query_lower = query.lower()
        for keyword in dangerous_keywords:
            if keyword in query_lower:
                return False
        
        return True

