"""
Lightweight Python sandbox using RestrictedPython.
Fallback when Docker is not available.
"""

from typing import Dict, Any, Optional
import ast
from dataclasses import dataclass
import time
import sys
from io import StringIO


@dataclass
class ExecutionResult:
    """Result of code execution in sandbox."""
    success: bool
    output: str
    error: Optional[str]
    execution_time: float
    memory_used: int
    exit_code: int
    return_value: Any


class RestrictedExecutor:
    """
    Execute Python code in restricted environment.
    
    Security features:
    - No file I/O operations
    - No network access
    - No dangerous imports (os, sys, subprocess, etc.)
    - Limited built-in functions
    - No access to __builtins__
    
    Note: Less secure than Docker, use only for trusted code.
    """
    
    # Allowed built-in functions
    SAFE_BUILTINS = {
        'abs': abs,
        'all': all,
        'any': any,
        'bool': bool,
        'dict': dict,
        'enumerate': enumerate,
        'filter': filter,
        'float': float,
        'int': int,
        'len': len,
        'list': list,
        'map': map,
        'max': max,
        'min': min,
        'range': range,
        'reversed': reversed,
        'round': round,
        'set': set,
        'sorted': sorted,
        'str': str,
        'sum': sum,
        'tuple': tuple,
        'zip': zip,
        'print': print,
        'isinstance': isinstance,
        'issubclass': issubclass,
        'Exception': Exception,
        'ValueError': ValueError,
        'TypeError': TypeError,
        'IndexError': IndexError,
        'KeyError': KeyError,
        'ImportError': ImportError,
        'AttributeError': AttributeError,
        'NameError': NameError,
        'RuntimeError': RuntimeError,
        'ZeroDivisionError': ZeroDivisionError,
        'AssertionError': AssertionError,
        'AttributeError': AttributeError,
        'EOFError': EOFError,
        'FloatingPointError': FloatingPointError,
        'GeneratorExit': GeneratorExit,
        'ImportError': ImportError,
        'IndexError': IndexError,
        'KeyError': KeyError,
        'KeyboardInterrupt': KeyboardInterrupt,
        'MemoryError': MemoryError,
        'NotImplementedError': NotImplementedError,
        'OSError': OSError,
        'OverflowError': OverflowError,
        'ReferenceError': ReferenceError,
        'StopIteration': StopIteration,
        'SyntaxError': SyntaxError,
        'IndentationError': IndentationError,
        'TabError': TabError,
        'SystemError': SystemError,
        'SystemExit': SystemExit,
        'TypeError': TypeError,
        'UnboundLocalError': UnboundLocalError,
        'UnicodeError': UnicodeError,
        'UnicodeEncodeError': UnicodeEncodeError,
        'UnicodeDecodeError': UnicodeDecodeError,
        'UnicodeTranslateError': UnicodeTranslateError,
        'ValueError': ValueError,
        'EnvironmentError': EnvironmentError,
        'IOError': IOError,
        'LookupError': LookupError,
        'ArithmeticError': ArithmeticError,
        'BufferError': BufferError,
        'BaseException': BaseException,
    }
    
    # Blocked imports
    BLOCKED_MODULES = {
        'os', 'sys', 'subprocess', 'socket', 'urllib',
        'requests', 'httpx', 'ftplib', 'telnetlib',
        'pickle', 'shelve', 'marshal', 'importlib',
        '__import__', 'eval', 'exec', 'compile',
    }
    
    def __init__(
        self,
        timeout: int = 10,
        max_output_size: int = 10000,  # characters
        allowed_imports: list = None,  # List of allowed modules
        allow_file_read: bool = True,  # Allow read-only file operations
    ):
        """
        Initialize restricted executor.
        
        Args:
            timeout: Execution timeout in seconds
            max_output_size: Maximum output size in characters
            allowed_imports: List of allowed modules (e.g. ['numpy', 'math'])
            allow_file_read: Allow read-only file operations (default: True)
        """
        self.timeout = timeout
        self.max_output_size = max_output_size
        self.allowed_imports = set(allowed_imports or [])
        self.allow_file_read = allow_file_read
    
    def execute(
        self,
        code: str,
        context: Optional[Dict[str, Any]] = None
    ) -> ExecutionResult:
        """
        Execute Python code in restricted environment.
        
        Args:
            code: Python code to execute
            context: Optional context variables
        
        Returns:
            ExecutionResult with output, errors, and metrics
        """
        start_time = time.time()
        
        # Custom safe import function
        def safe_import(name, globals=None, locals=None, fromlist=(), level=0):
            # Extract base package names from allowed_imports (handle "pandas as pd" format)
            allowed_base_packages = set()
            for imp_spec in self.allowed_imports:
                if ' as ' in imp_spec:
                    # Extract the module name before "as"
                    module_name = imp_spec.split(' as ')[0].strip()
                    # Get base package (e.g., "plotly.express" -> "plotly")
                    base_pkg = module_name.split('.')[0]
                    allowed_base_packages.add(base_pkg)
                    # Also add the full module name
                    allowed_base_packages.add(module_name)
                else:
                    # Simple import without alias
                    base_pkg = imp_spec.strip().split('.')[0]
                    allowed_base_packages.add(base_pkg)
                    allowed_base_packages.add(imp_spec.strip())
            
            # Check if the import is allowed
            base_name = name.split('.')[0]
            if base_name in allowed_base_packages or name in allowed_base_packages:
                return __import__(name, globals, locals, fromlist, level)
            
            raise ImportError(f"Import of '{name}' is not allowed")

        # Safe open function (read-only)
        def safe_open(file, mode='r', *args, **kwargs):
            if self.allow_file_read:
                # Only allow read modes
                if mode not in ['r', 'rb', 'rt']:
                    raise SecurityError(f"File mode '{mode}' is not allowed. Only read modes ('r', 'rb', 'rt') are permitted.")
                return open(file, mode, *args, **kwargs)
            else:
                raise SecurityError("File operations are not allowed")

        # Prepare execution environment
        safe_globals = {
            '__builtins__': self.SAFE_BUILTINS.copy(),
        }
        
        # Add safe functions to builtins
        safe_globals['__builtins__']['open'] = safe_open
        safe_globals['__builtins__']['__import__'] = safe_import
        
        # Add context if provided
        if context:
            safe_globals.update(context)
            
        # Inject allowed imports into globals (convenience)
        for import_spec in self.allowed_imports:
            try:
                # Handle "module as alias" syntax
                if ' as ' in import_spec:
                    module_name, alias = import_spec.split(' as ', 1)
                    module_name = module_name.strip()
                    alias = alias.strip()
                    
                    # Handle submodule imports like "plotly.express as px"
                    if '.' in module_name:
                        parts = module_name.split('.')
                        module = __import__(module_name, fromlist=[parts[-1]])
                    else:
                        module = __import__(module_name)
                    
                    safe_globals[alias] = module
                else:
                    # Simple import without alias
                    module_name = import_spec.strip()
                    if '.' in module_name:
                        # For submodules like "plotly.express", import and add to globals
                        parts = module_name.split('.')
                        module = __import__(module_name, fromlist=[parts[-1]])
                        safe_globals[module_name] = module
                    else:
                        module = __import__(module_name)
                        safe_globals[module_name] = module
            except ImportError:
                pass 
        
        # Capture output
        old_stdout = sys.stdout
        sys.stdout = output_buffer = StringIO()
        
        try:
            self._check_code_safety(code)

            return_value = None
            try:
                tree = ast.parse(code, mode='exec')
                if tree.body and isinstance(tree.body[-1], ast.Expr):
                    exec_body = tree.body[:-1]
                    last_expr = tree.body[-1].value
                    if exec_body:
                        compiled_exec = compile(ast.Module(body=exec_body, type_ignores=[]), '<restricted>', 'exec')
                        exec(compiled_exec, safe_globals, {})
                    compiled_eval = compile(ast.Expression(last_expr), '<restricted>', 'eval')
                    return_value = eval(compiled_eval, safe_globals, {})
                else:
                    exec(code, safe_globals, {})
            except Exception:
                exec(code, safe_globals, {})
                return_value = None

            output = output_buffer.getvalue()
            if len(output) > self.max_output_size:
                output = output[:self.max_output_size] + "\n... (output truncated)"
            execution_time = time.time() - start_time
            return ExecutionResult(
                success=True,
                output=output,
                error=None,
                execution_time=execution_time,
                memory_used=0,
                exit_code=0,
                return_value=return_value
            )
        
        except Exception as e:
            execution_time = time.time() - start_time
            return ExecutionResult(
                success=False,
                output=output_buffer.getvalue(),
                error=str(e),
                execution_time=execution_time,
                memory_used=0,
                exit_code=1,
                return_value=None
            )
        
        finally:
            sys.stdout = old_stdout
    
    def _check_code_safety(self, code: str):
        """
        Check if code contains blocked operations.
        
        Args:
            code: Python code to check
        
        Raises:
            SecurityError: If code contains blocked operations
        """
        code_lower = code.lower()
        
        # Check for blocked imports
        for module in self.BLOCKED_MODULES:
            # Skip if explicitly allowed
            if module in self.allowed_imports:
                continue
                
            if f"import {module}" in code_lower or f"from {module}" in code_lower:
                raise SecurityError(f"Import of '{module}' is not allowed")
        
        # Check for dangerous file operations (write/delete)
        # Note: read operations are now allowed via safe_open
        dangerous_file_ops = [
            'remove(', 'unlink(', 'rmdir(', 'removedirs(',
            'mkdir(', 'makedirs(', 'rename(', 'replace(',
            "mode='w", 'mode="w', "mode='a", 'mode="a',
            "mode='x", 'mode="x', "mode='wb", 'mode="wb',
            "mode='ab", 'mode="ab', "mode='xb", 'mode="xb'
        ]
        if any(op in code for op in dangerous_file_ops):
            raise SecurityError("Write/delete file operations are not allowed. Only read mode is permitted.")
        
        # Check for dangerous built-ins
        if any(fn in code for fn in ['__import__', 'eval', 'exec', 'compile']):
            raise SecurityError("Dangerous built-in functions are not allowed")


class SecurityError(Exception):
    """Raised when code violates security restrictions."""
    pass
