"""
Safe code executor for Natural Language Query system.

Executes LLM-generated Python code in a sandboxed environment with validation
and timeout mechanisms.

 Implementation:
- Separate CodeValidator for safety checks
- AST + regex-based validation
- Timeout mechanism with signal handling
- Sandboxed execution namespace
- Self-correction with LLM integration
- Comprehensive error handling
"""

from typing import Any, Dict, Optional, Tuple
import sys
import io
import signal
from contextlib import contextmanager
import logging

# Import CodeValidator for 
try:
    from .code_validator import CodeValidator
    VALIDATOR_AVAILABLE = True
except ImportError:
    VALIDATOR_AVAILABLE = False
    CodeValidator = None

logger = logging.getLogger(__name__)


class ExecutionTimeout(Exception):
    """Raised when code execution exceeds timeout."""

    pass


class CodeValidationError(Exception):
    """Raised when code fails validation checks."""

    pass


@contextmanager
def timeout(seconds: int):
    """
    Context manager for timeout.

    Args:
        seconds: Timeout in seconds

    Raises:
        ExecutionTimeout: If code execution exceeds timeout
    """

    def timeout_handler(signum, frame):
        raise ExecutionTimeout(f"Code execution exceeded {seconds} seconds")

    # Set the signal handler
    old_handler = signal.signal(signal.SIGALRM, timeout_handler)
    signal.alarm(seconds)

    try:
        yield
    finally:
        signal.alarm(0)
        signal.signal(signal.SIGALRM, old_handler)


class CodeExecutor:
    """
    Execute LLM-generated code safely with validation and sandboxing.

    Features:
    - AST-based validation (no file I/O, no imports, no exec/eval)
    - Execution timeout
    - Output capture (stdout/stderr)
    - Error handling with detailed messages
    """

    # Dangerous functions/modules to block
    BLOCKED_BUILTINS = {
        "__import__",
        "eval",
        "exec",
        "compile",
        "open",
        "input",
        "raw_input",
        "execfile",
    }

    # Dangerous modules to block
    BLOCKED_MODULES = {
        "os",
        "sys",
        "subprocess",
        "socket",
        "urllib",
        "requests",
        "shutil",
        "glob",
        "pickle",
        "shelve",
    }

    def __init__(
        self,
        timeout_seconds: int = 30,
        max_output_size: int = 10000,
        use_validator: bool = True,
        strict_mode: bool = True,
    ):
        """
        Initialize code executor.

        Args:
            timeout_seconds: Maximum execution time in seconds
            max_output_size: Maximum output size in characters
            use_validator: Use CodeValidator for enhanced validation 
            strict_mode: Enable strict validation mode
        """
        self.timeout_seconds = timeout_seconds
        self.max_output_size = max_output_size
        self.use_validator = use_validator and VALIDATOR_AVAILABLE
        self.strict_mode = strict_mode

        # Initialize CodeValidator 
        if self.use_validator:
            self.validator = CodeValidator(strict_mode=strict_mode)
            logger.debug("CodeExecutor initialized with CodeValidator ")
        else:
            self.validator = None
            if use_validator and not VALIDATOR_AVAILABLE:
                logger.warning("CodeValidator requested but not available, using basic validation")
            logger.debug("CodeExecutor initialized with basic validation")

    def validate_code(self, code: str) -> bool:
        """
        Validate code using AST analysis.

        Uses CodeValidator  if available, otherwise falls back to basic validation.

        Args:
            code: Python code to validate

        Returns:
            bool: True if code is safe

        Raises:
            CodeValidationError: If code contains dangerous operations
        """
        # Use enhanced CodeValidator if available 
        if self.use_validator and self.validator:
            is_valid, errors = self.validator.validate(code)
            if not is_valid:
                error_msg = "Code validation failed:\n" + "\n".join(f"  - {e}" for e in errors)
                raise CodeValidationError(error_msg)
            return True

        # Fallback to basic AST validation (backward compatibility)
        try:
            import ast
            tree = ast.parse(code)
        except SyntaxError as e:
            raise CodeValidationError(f"Syntax error: {e}")

        # Check for dangerous operations
        for node in ast.walk(tree):
            # Block imports
            if isinstance(node, (ast.Import, ast.ImportFrom)):
                if isinstance(node, ast.ImportFrom):
                    module = node.module
                    if module and any(
                        blocked in module for blocked in self.BLOCKED_MODULES
                    ):
                        raise CodeValidationError(
                            f"Blocked import: {module}"
                        )

            # Block dangerous built-ins
            if isinstance(node, ast.Name):
                if node.id in self.BLOCKED_BUILTINS:
                    raise CodeValidationError(
                        f"Blocked built-in function: {node.id}"
                    )

            # Block file operations
            if isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name):
                    if node.func.id in ["open", "file"]:
                        raise CodeValidationError(
                            "File operations are not allowed"
                        )

        return True

    def execute(
        self, code: str, context: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Execute code in a sandboxed environment.

        Args:
            code: Python code to execute
            context: Variables to make available in execution context

        Returns:
            dict: Execution result with 'result', 'stdout', 'stderr', 'error'
        """
        # Validate code first
        try:
            self.validate_code(code)
        except CodeValidationError as e:
            return {"success": False, "error": str(e), "error_type": "validation"}

        # Prepare execution context
        exec_globals = {
            "__builtins__": {
                k: v
                for k, v in __builtins__.items()
                if k not in self.BLOCKED_BUILTINS
            }
        }

        # Add user-provided context
        if context:
            exec_globals.update(context)

        # Capture stdout/stderr
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()

        old_stdout = sys.stdout
        old_stderr = sys.stderr

        result = None
        error = None
        error_type = None

        try:
            sys.stdout = stdout_capture
            sys.stderr = stderr_capture

            # Execute with timeout
            with timeout(self.timeout_seconds):
                exec(code, exec_globals)

                # Get result if assigned
                result = exec_globals.get("result")

            success = True

        except ExecutionTimeout as e:
            success = False
            error = str(e)
            error_type = "timeout"

        except Exception as e:
            success = False
            error = f"{type(e).__name__}: {str(e)}"
            error_type = type(e).__name__

        finally:
            # Restore stdout/stderr
            sys.stdout = old_stdout
            sys.stderr = old_stderr

        # Get captured output
        stdout_text = stdout_capture.getvalue()
        stderr_text = stderr_capture.getvalue()

        # Truncate if too long
        if len(stdout_text) > self.max_output_size:
            stdout_text = (
                stdout_text[: self.max_output_size] + "\n... (output truncated)"
            )

        if len(stderr_text) > self.max_output_size:
            stderr_text = (
                stderr_text[: self.max_output_size] + "\n... (output truncated)"
            )

        return {
            "success": success,
            "result": result,
            "stdout": stdout_text,
            "stderr": stderr_text,
            "error": error,
            "error_type": error_type,
        }

    def execute_with_retry(
        self,
        code: str,
        context: Optional[Dict[str, Any]] = None,
        max_retries: int = 3,
        on_error=None,
    ) -> Dict[str, Any]:
        """
        Execute code with retry on failure.

        Basic retry method with callback support (backward compatible).

        Args:
            code: Python code to execute
            context: Execution context
            max_retries: Maximum number of retries
            on_error: Callback function for error correction (code, error) -> corrected_code

        Returns:
            dict: Execution result
        """
        last_result = None

        for attempt in range(max_retries):
            result = self.execute(code, context)

            if result["success"]:
                return result

            last_result = result

            # Try to fix error if callback provided
            if on_error and attempt < max_retries - 1:
                try:
                    code = on_error(code, result["error"])
                    logger.info(f"Self-correction attempt {attempt + 1}/{max_retries - 1}")
                except Exception as e:
                    logger.warning(f"Error correction failed: {e}")
                    break

        # All retries failed
        return last_result

    def execute_with_llm_retry(
        self,
        query: str,
        initial_code: str,
        stream,  # Eventstream object
        llm_agent=None,  # LLMAgent for error correction
        max_retries: int = 3,
        verbose: bool = False,
    ) -> Tuple[bool, Dict[str, Any], Optional[str]]:
        """
        Execute code with automatic retry and LLM-based self-correction .

        This is the recommended method for production use with LLM integration.

        Args:
            query: Original user query
            initial_code: Generated code to execute
            stream: Eventstream object to operate on
            llm_agent: LLMAgent instance for error correction (optional)
            max_retries: Maximum retry attempts
            verbose: Print debug information

        Returns:
            Tuple of (success: bool, results: dict, error: Optional[str])
            - success: True if execution succeeded
            - results: Dictionary with 'result', 'stdout', 'stderr', etc.
            - error: Error message if failed, None otherwise

        Example:
            >>> from outhad_edge.ai import CodeExecutor, LLMAgent
            >>> executor = CodeExecutor()
            >>> agent = LLMAgent(model='gpt-4')
            >>> success, results, error = executor.execute_with_llm_retry(
            ...     query="What is conversion rate?",
            ...     initial_code=code,
            ...     stream=eventstream,
            ...     llm_agent=agent
            ... )
        """
        code = initial_code
        attempt = 0
        context = {'stream': stream}

        while attempt < max_retries:
            if verbose:
                print(f"\n[Attempt {attempt + 1}/{max_retries}] Executing code...")

            # Execute code
            result = self.execute(code, context)

            if result["success"]:
                if verbose:
                    print("✓ Execution successful")
                return True, result, None

            # Execution failed
            error = result.get("error", "Unknown error")
            error_type = result.get("error_type", "unknown")

            if verbose:
                print(f"✗ Execution failed: {error_type}")
                print(f"  Error: {error}")

            attempt += 1

            # Try self-correction with LLM if available and retries left
            if llm_agent and attempt < max_retries:
                if verbose:
                    print(f"\n[Self-Correction {attempt}] Using LLM to fix error...")

                try:
                    # Use LLMAgent to correct the error
                    corrected_code = llm_agent.correct_error(
                        query=query,
                        code=code,
                        error=error
                    )

                    if verbose:
                        print(f"✓ Generated corrected code ({len(corrected_code)} chars)")

                    code = corrected_code

                except Exception as e:
                    if verbose:
                        print(f"✗ Self-correction failed: {e}")
                    # Can't fix, break retry loop
                    break
            else:
                # No LLM agent or no retries left
                if verbose and not llm_agent:
                    print("  (No LLM agent provided for self-correction)")
                break

        # All retries exhausted
        final_error = f"Failed after {attempt} attempt(s). Last error:\n{error}"
        return False, {}, final_error


def test_executor():
    """Test the code executor with sample code."""
    executor = CodeExecutor(timeout_seconds=5)

    # Test 1: Valid code
    print("Test 1: Valid code")
    code1 = """
x = 10
y = 20
result = x + y
print(f"Sum: {result}")
"""
    result1 = executor.execute(code1)
    print(f"Success: {result1['success']}")
    print(f"Result: {result1['result']}")
    print(f"Output: {result1['stdout']}")
    print()

    # Test 2: Code with timeout
    print("Test 2: Timeout")
    code2 = """
import time
time.sleep(10)  # Will be blocked by import check
"""
    result2 = executor.execute(code2)
    print(f"Success: {result2['success']}")
    print(f"Error: {result2['error']}")
    print()

    # Test 3: Code with blocked function
    print("Test 3: Blocked function")
    code3 = """
f = open('test.txt', 'w')
"""
    result3 = executor.execute(code3)
    print(f"Success: {result3['success']}")
    print(f"Error: {result3['error']}")
    print()


if __name__ == "__main__":
    test_executor()
