"""
Code Validator for Natural Language Query system.

Validates LLM-generated Python code for safety before execution using:
- AST-based static analysis
- Regex pattern matching for dangerous operations
- Import whitelist/blacklist
- Syntax checking
"""

import ast
import re
from typing import List, Tuple, Set
import logging

logger = logging.getLogger(__name__)


class CodeValidator:
    """
    Validates generated Python code for safety before execution.

    Uses multi-layered validation:
    1. Syntax validation (AST parsing)
    2. Import validation (whitelist/blacklist)
    3. Dangerous pattern detection (regex)
    4. AST node analysis (blocked operations)

    Example:
        >>> validator = CodeValidator()
        >>> is_valid, errors = validator.validate("print('hello')")
        >>> if is_valid:
        ...     print("Code is safe to execute")
    """

    # Allowed imports (whitelist)
    ALLOWED_IMPORTS: Set[str] = {
        'pandas',
        'pd',
        'numpy',
        'np',
        'outhad_edge',
        'datetime',
        'collections',
        'itertools',
        'math',
        'statistics',
        'json',
        're',
    }

    # Blocked modules (blacklist - takes priority)
    BLOCKED_MODULES: Set[str] = {
        'os',
        'sys',
        'subprocess',
        'socket',
        'urllib',
        'urllib2',
        'urllib3',
        'requests',
        'http',
        'shutil',
        'glob',
        'pickle',
        'shelve',
        'ctypes',
        'multiprocessing',
        'threading',
        'asyncio',
    }

    # Dangerous regex patterns
    DANGEROUS_PATTERNS: List[str] = [
        r'__import__',
        r'eval\s*\(',
        r'exec\s*\(',
        r'compile\s*\(',
        r'globals\s*\(',
        r'locals\s*\(',
        r'vars\s*\(',
        r'dir\s*\(',
        r'getattr\s*\(',
        r'setattr\s*\(',
        r'delattr\s*\(',
        r'hasattr\s*\(',
        r'\bos\.',
        r'\bsys\.',
        r'subprocess\.',
        r'open\s*\(',
        r'file\s*\(',
        r'input\s*\(',
        r'raw_input\s*\(',
        r'\.delete\s*\(',
        r'\.drop\s*\(',
        r'\.remove\s*\(',
        r'\.unlink\s*\(',
        r'\brm\s+',
        r'\bdel\s+',
        r'\.system\s*\(',
        r'\.popen\s*\(',
    ]

    # Blocked built-in functions
    BLOCKED_BUILTINS: Set[str] = {
        '__import__',
        'eval',
        'exec',
        'compile',
        'open',
        'input',
        'raw_input',
        'execfile',
        'file',
        'reload',
        'vars',
        'globals',
        'locals',
    }

    def __init__(
        self,
        allowed_imports: Set[str] = None,
        blocked_modules: Set[str] = None,
        additional_patterns: List[str] = None,
        strict_mode: bool = True,
    ):
        """
        Initialize code validator.

        Args:
            allowed_imports: Custom set of allowed imports (overrides default if provided)
            blocked_modules: Additional blocked modules to add to blacklist
            additional_patterns: Additional dangerous patterns to check
            strict_mode: If True, use strict validation (recommended)
        """
        # Use custom or default allowed imports
        self.allowed_imports = allowed_imports or self.ALLOWED_IMPORTS.copy()

        # Combine default and custom blocked modules
        self.blocked_modules = self.BLOCKED_MODULES.copy()
        if blocked_modules:
            self.blocked_modules.update(blocked_modules)

        # Combine default and custom patterns
        self.dangerous_patterns = self.DANGEROUS_PATTERNS.copy()
        if additional_patterns:
            self.dangerous_patterns.extend(additional_patterns)

        self.strict_mode = strict_mode

        logger.debug(f"CodeValidator initialized (strict={strict_mode})")

    def validate(self, code: str) -> Tuple[bool, List[str]]:
        """
        Validate code for safety.

        Performs multi-layered validation:
        1. Syntax check (AST parsing)
        2. Import validation
        3. Dangerous pattern detection
        4. AST node analysis

        Args:
            code: Python code string to validate

        Returns:
            Tuple of (is_valid: bool, errors: List[str])
            - is_valid: True if code passed all checks
            - errors: List of validation error messages
        """
        errors = []

        if not code or not code.strip():
            errors.append("Code is empty")
            return False, errors

        # 1. Syntax validation
        syntax_errors = self._check_syntax(code)
        errors.extend(syntax_errors)

        if syntax_errors and self.strict_mode:
            # If syntax is invalid, can't do AST analysis
            return False, errors

        # 2. Dangerous pattern detection (regex)
        pattern_errors = self._check_dangerous_patterns(code)
        errors.extend(pattern_errors)

        # 3. AST-based validation (imports, builtins, etc.)
        if not syntax_errors:  # Only if syntax is valid
            ast_errors = self._check_ast(code)
            errors.extend(ast_errors)

        is_valid = len(errors) == 0

        if not is_valid:
            logger.warning(f"Code validation failed: {len(errors)} errors")
            for error in errors:
                logger.debug(f"  - {error}")

        return is_valid, errors

    def _check_syntax(self, code: str) -> List[str]:
        """Check code syntax using AST parsing."""
        errors = []

        try:
            ast.parse(code)
        except SyntaxError as e:
            errors.append(f"Syntax error on line {e.lineno}: {e.msg}")
        except Exception as e:
            errors.append(f"Parse error: {str(e)}")

        return errors

    def _check_dangerous_patterns(self, code: str) -> List[str]:
        """Check for dangerous patterns using regex."""
        errors = []

        for pattern in self.dangerous_patterns:
            if re.search(pattern, code, re.IGNORECASE):
                # Extract pattern name for better error message
                pattern_name = pattern.replace(r'\b', '').replace(r'\s*\(', '').replace(r'\(', '')
                errors.append(f"Dangerous pattern detected: {pattern_name}")

        return errors

    def _check_ast(self, code: str) -> List[str]:
        """Check AST for dangerous operations."""
        errors = []

        try:
            tree = ast.parse(code)
        except:
            # Syntax errors already caught in _check_syntax
            return errors

        for node in ast.walk(tree):
            # Check imports
            if isinstance(node, ast.Import):
                for alias in node.names:
                    import_name = alias.name.split('.')[0]  # Get root module
                    if import_name in self.blocked_modules:
                        errors.append(f"Blocked import: {alias.name}")
                    elif self.strict_mode and import_name not in self.allowed_imports:
                        errors.append(
                            f"Import not in whitelist: {alias.name}. "
                            f"Allowed: {', '.join(sorted(self.allowed_imports))}"
                        )

            # Check from imports
            if isinstance(node, ast.ImportFrom):
                if node.module:
                    module_name = node.module.split('.')[0]  # Get root module
                    if module_name in self.blocked_modules:
                        errors.append(f"Blocked import from: {node.module}")
                    elif self.strict_mode and module_name not in self.allowed_imports:
                        errors.append(
                            f"Import not in whitelist: {node.module}. "
                            f"Allowed: {', '.join(sorted(self.allowed_imports))}"
                        )

            # Check for blocked built-in function calls
            if isinstance(node, ast.Name):
                if node.id in self.BLOCKED_BUILTINS:
                    errors.append(f"Blocked built-in function: {node.id}")

            # Check for file operations
            if isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name):
                    func_name = node.func.id
                    if func_name in ['open', 'file']:
                        errors.append(f"File operation not allowed: {func_name}()")

            # Check for attribute access to blocked modules
            if isinstance(node, ast.Attribute):
                if isinstance(node.value, ast.Name):
                    if node.value.id in self.blocked_modules:
                        errors.append(
                            f"Access to blocked module: {node.value.id}.{node.attr}"
                        )

        return errors

    def validate_and_raise(self, code: str) -> None:
        """
        Validate code and raise exception if invalid.

        Args:
            code: Python code to validate

        Raises:
            CodeValidationError: If validation fails
        """
        is_valid, errors = self.validate(code)

        if not is_valid:
            from .code_executor import CodeValidationError
            error_msg = "Code validation failed:\n" + "\n".join(f"  - {e}" for e in errors)
            raise CodeValidationError(error_msg)

    def is_safe(self, code: str) -> bool:
        """
        Quick check if code is safe (returns bool only).

        Args:
            code: Python code to check

        Returns:
            bool: True if safe, False otherwise
        """
        is_valid, _ = self.validate(code)
        return is_valid

    def get_allowed_imports(self) -> Set[str]:
        """Get set of allowed imports."""
        return self.allowed_imports.copy()

    def get_blocked_modules(self) -> Set[str]:
        """Get set of blocked modules."""
        return self.blocked_modules.copy()

    def add_allowed_import(self, module: str) -> None:
        """
        Add module to allowed imports whitelist.

        Args:
            module: Module name to allow
        """
        self.allowed_imports.add(module)
        logger.debug(f"Added {module} to allowed imports")

    def remove_allowed_import(self, module: str) -> None:
        """
        Remove module from allowed imports whitelist.

        Args:
            module: Module name to disallow
        """
        self.allowed_imports.discard(module)
        logger.debug(f"Removed {module} from allowed imports")

    def add_blocked_module(self, module: str) -> None:
        """
        Add module to blocked modules blacklist.

        Args:
            module: Module name to block
        """
        self.blocked_modules.add(module)
        logger.debug(f"Added {module} to blocked modules")

    def __repr__(self) -> str:
        return (
            f"CodeValidator(allowed_imports={len(self.allowed_imports)}, "
            f"blocked_modules={len(self.blocked_modules)}, "
            f"strict_mode={self.strict_mode})"
        )
