"""AST-based analysis for Python files.

Catches evasion techniques that regex cannot:
- Import aliasing (from os import system as s)
- Nested/chained attribute access
- Dynamic function construction
- Control flow to dangerous calls
"""
import ast
from dataclasses import dataclass
from pathlib import Path
from typing import List, Set, Dict, Optional, Tuple

from .patterns import Finding, Severity


# Dangerous functions we track through the AST
DANGEROUS_FUNCTIONS = {
    'eval', 'exec', 'compile', '__import__',
    'getattr', 'setattr', 'delattr',
    'open',  # Track for path traversal
}

DANGEROUS_CALLABLES = {
    # builtins
    ('builtins', 'eval'), ('builtins', 'exec'), ('builtins', 'compile'),
    ('builtins', '__import__'), ('builtins', 'open'),
    # os module
    ('os', 'system'), ('os', 'popen'), ('os', 'spawn'),
    ('os', 'spawnl'), ('os', 'spawnle'), ('os', 'spawnlp'), ('os', 'spawnlpe'),
    ('os', 'spawnv'), ('os', 'spawnve'), ('os', 'spawnvp'), ('os', 'spawnvpe'),
    ('os', 'execl'), ('os', 'execle'), ('os', 'execlp'), ('os', 'execlpe'),
    ('os', 'execv'), ('os', 'execve'), ('os', 'execvp'), ('os', 'execvpe'),
    ('os', 'remove'), ('os', 'unlink'), ('os', 'rmdir'),
    # subprocess
    ('subprocess', 'run'), ('subprocess', 'call'), ('subprocess', 'Popen'),
    ('subprocess', 'check_call'), ('subprocess', 'check_output'),
    # pickle
    ('pickle', 'load'), ('pickle', 'loads'),
    ('_pickle', 'load'), ('_pickle', 'loads'),
    # marshal
    ('marshal', 'load'), ('marshal', 'loads'),
    # socket
    ('socket', 'socket'), ('socket', 'create_connection'),
    # requests
    ('requests', 'get'), ('requests', 'post'), ('requests', 'put'),
    ('requests', 'delete'), ('requests', 'patch'),
    # urllib
    ('urllib.request', 'urlopen'), ('urllib.request', 'urlretrieve'),
}

# Modules that when imported suggest suspicious activity
SUSPICIOUS_MODULES = {
    'ctypes',  # Memory manipulation
    'mmap',    # Memory mapping
    'pty',     # Pseudo-terminal (reverse shells)
    'fcntl',   # File control
    'resource', # Resource limits manipulation
}


@dataclass
class ImportedName:
    """Tracks an imported name and its origin."""
    name: str           # Local name (e.g., 's')
    module: str         # Source module (e.g., 'os')
    original: str       # Original name (e.g., 'system')
    line: int
    
    @property
    def qualified(self) -> Tuple[str, str]:
        return (self.module, self.original)


class DangerousCallVisitor(ast.NodeVisitor):
    """AST visitor that tracks imports and finds dangerous calls."""
    
    def __init__(self, filename: str):
        self.filename = filename
        self.findings: List[Finding] = []
        
        # Track imported names: local_name -> ImportedName
        self.imports: Dict[str, ImportedName] = {}
        
        # Track module aliases: alias -> module_name
        self.module_aliases: Dict[str, str] = {}
        
        # Track variable assignments that might hold dangerous refs
        self.dangerous_vars: Dict[str, Tuple[str, int]] = {}  # name -> (reason, line)
    
    def visit_Import(self, node: ast.Import):
        """Handle: import os, import subprocess as sp"""
        for alias in node.names:
            module = alias.name
            local_name = alias.asname or alias.name
            
            self.module_aliases[local_name] = module
            
            # Check for suspicious module imports
            base_module = module.split('.')[0]
            if base_module in SUSPICIOUS_MODULES:
                self.findings.append(Finding(
                    severity=Severity.WARN,
                    pattern_id='suspicious_module_import',
                    description=f'Import of suspicious module: {module}',
                    file=self.filename,
                    line=node.lineno,
                    match=f'import {module}',
                    context=f'import {module}' + (f' as {alias.asname}' if alias.asname else ''),
                ))
        
        self.generic_visit(node)
    
    def visit_ImportFrom(self, node: ast.ImportFrom):
        """Handle: from os import system, from os import system as s"""
        module = node.module or ''
        
        for alias in node.names:
            original_name = alias.name
            local_name = alias.asname or alias.name
            
            # Track the import
            self.imports[local_name] = ImportedName(
                name=local_name,
                module=module,
                original=original_name,
                line=node.lineno,
            )
            
            # Check if importing something dangerous
            if (module, original_name) in DANGEROUS_CALLABLES:
                severity = Severity.DANGER if original_name in ('system', 'popen', 'eval', 'exec') else Severity.WARN
                
                # If aliased, it's more suspicious (trying to hide)
                if alias.asname:
                    severity = Severity.DANGER
                    desc = f'Dangerous function imported with alias: {module}.{original_name} as {local_name}'
                else:
                    desc = f'Dangerous function imported: {module}.{original_name}'
                
                self.findings.append(Finding(
                    severity=severity,
                    pattern_id='dangerous_import',
                    description=desc,
                    file=self.filename,
                    line=node.lineno,
                    match=f'from {module} import {original_name}',
                    context=ast.unparse(node) if hasattr(ast, 'unparse') else '',
                ))
            
            # Check for suspicious module
            base_module = module.split('.')[0] if module else ''
            if base_module in SUSPICIOUS_MODULES:
                self.findings.append(Finding(
                    severity=Severity.WARN,
                    pattern_id='suspicious_module_import',
                    description=f'Import from suspicious module: {module}',
                    file=self.filename,
                    line=node.lineno,
                    match=f'from {module} import {original_name}',
                ))
        
        self.generic_visit(node)
    
    def visit_Assign(self, node: ast.Assign):
        """Track assignments that might store dangerous references."""
        # Check for: x = eval, x = __builtins__.eval, etc.
        if isinstance(node.value, ast.Name):
            value_name = node.value.id
            if value_name in ('eval', 'exec', 'compile', 'system', 'open'):
                for target in node.targets:
                    if isinstance(target, ast.Name):
                        self.dangerous_vars[target.id] = (f'alias for {value_name}', node.lineno)
                        self.findings.append(Finding(
                            severity=Severity.DANGER,
                            pattern_id='dangerous_alias',
                            description=f'Variable assigned dangerous builtin: {target.id} = {value_name}',
                            file=self.filename,
                            line=node.lineno,
                            match=f'{target.id} = {value_name}',
                        ))
        
        # Check for: x = getattr(module, 'dangerous')
        if isinstance(node.value, ast.Call):
            call = node.value
            if isinstance(call.func, ast.Name) and call.func.id == 'getattr':
                if len(call.args) >= 2 and isinstance(call.args[1], ast.Constant):
                    attr_name = call.args[1].value
                    if attr_name in ('eval', 'exec', 'compile', 'system', 'popen'):
                        for target in node.targets:
                            if isinstance(target, ast.Name):
                                self.dangerous_vars[target.id] = (f'getattr result for {attr_name}', node.lineno)
        
        self.generic_visit(node)
    
    def visit_Call(self, node: ast.Call):
        """Check function calls for dangerous patterns."""
        
        # Direct call to dangerous name: eval(...), exec(...)
        if isinstance(node.func, ast.Name):
            func_name = node.func.id
            
            # Check if it's an imported dangerous function
            if func_name in self.imports:
                imp = self.imports[func_name]
                if imp.qualified in DANGEROUS_CALLABLES:
                    self.findings.append(Finding(
                        severity=Severity.DANGER,
                        pattern_id='dangerous_call_aliased',
                        description=f'Call to aliased dangerous function: {func_name} (imported from {imp.module}.{imp.original})',
                        file=self.filename,
                        line=node.lineno,
                        match=f'{func_name}(...)',
                    ))
            
            # Check if it's a variable holding a dangerous reference
            if func_name in self.dangerous_vars:
                reason, _ = self.dangerous_vars[func_name]
                self.findings.append(Finding(
                    severity=Severity.DANGER,
                    pattern_id='dangerous_var_call',
                    description=f'Call via dangerous variable: {func_name} ({reason})',
                    file=self.filename,
                    line=node.lineno,
                    match=f'{func_name}(...)',
                ))
        
        # Attribute call: os.system(...), subprocess.run(..., shell=True)
        if isinstance(node.func, ast.Attribute):
            self._check_attribute_call(node)
        
        # Check subprocess calls for shell=True
        self._check_subprocess_shell(node)
        
        # Check for eval/exec with dynamic input
        self._check_dynamic_code_exec(node)
        
        self.generic_visit(node)
    
    def _check_attribute_call(self, node: ast.Call):
        """Check attribute-based calls like os.system(), module.func()"""
        attr = node.func
        if not isinstance(attr, ast.Attribute):
            return
        
        # Get the full attribute chain
        parts = []
        current = attr
        while isinstance(current, ast.Attribute):
            parts.append(current.attr)
            current = current.value
        
        if isinstance(current, ast.Name):
            parts.append(current.id)
        
        parts.reverse()
        
        if len(parts) >= 2:
            # Check if first part is an alias
            base = parts[0]
            if base in self.module_aliases:
                module = self.module_aliases[base]
            elif base in self.imports:
                module = self.imports[base].module
            else:
                module = base
            
            func = parts[-1]
            
            # Check against dangerous callables
            if (module, func) in DANGEROUS_CALLABLES:
                self.findings.append(Finding(
                    severity=Severity.WARN,  # WARN because might be legitimate
                    pattern_id='dangerous_module_call',
                    description=f'Call to potentially dangerous function: {module}.{func}',
                    file=self.filename,
                    line=node.lineno,
                    match='.'.join(parts) + '(...)',
                ))
    
    def _check_subprocess_shell(self, node: ast.Call):
        """Check for subprocess calls with shell=True."""
        # Get function name
        func_name = None
        if isinstance(node.func, ast.Name):
            if node.func.id in self.imports:
                imp = self.imports[node.func.id]
                if imp.module == 'subprocess':
                    func_name = imp.original
        elif isinstance(node.func, ast.Attribute):
            if isinstance(node.func.value, ast.Name):
                base = node.func.value.id
                module = self.module_aliases.get(base, base)
                if module == 'subprocess':
                    func_name = node.func.attr
        
        if func_name in ('run', 'call', 'Popen', 'check_call', 'check_output'):
            # Check for shell=True in keyword arguments
            for kw in node.keywords:
                if kw.arg == 'shell':
                    if isinstance(kw.value, ast.Constant) and kw.value.value is True:
                        self.findings.append(Finding(
                            severity=Severity.DANGER,
                            pattern_id='subprocess_shell_true',
                            description='subprocess call with shell=True - command injection risk',
                            file=self.filename,
                            line=node.lineno,
                            match=f'subprocess.{func_name}(..., shell=True)',
                        ))
                    elif not isinstance(kw.value, ast.Constant):
                        # shell= with a variable - suspicious
                        self.findings.append(Finding(
                            severity=Severity.WARN,
                            pattern_id='subprocess_shell_dynamic',
                            description='subprocess call with dynamic shell parameter',
                            file=self.filename,
                            line=node.lineno,
                            match=f'subprocess.{func_name}(..., shell=?)',
                        ))
    
    def _check_dynamic_code_exec(self, node: ast.Call):
        """Check if eval/exec is called with non-literal input."""
        func_name = None
        if isinstance(node.func, ast.Name):
            func_name = node.func.id
        
        if func_name in ('eval', 'exec', 'compile'):
            if node.args:
                first_arg = node.args[0]
                # If first arg is not a string literal, it's dynamic
                if not isinstance(first_arg, ast.Constant):
                    self.findings.append(Finding(
                        severity=Severity.DANGER,
                        pattern_id='dynamic_code_execution',
                        description=f'{func_name}() called with dynamic input - high risk',
                        file=self.filename,
                        line=node.lineno,
                        match=f'{func_name}(<dynamic>)',
                    ))


def analyze_python_ast(content: str, filename: str) -> List[Finding]:
    """Analyze Python code using AST for dangerous patterns."""
    try:
        tree = ast.parse(content, filename=filename)
    except SyntaxError as e:
        # Return a finding about syntax error (might be obfuscation)
        return [Finding(
            severity=Severity.WARN,
            pattern_id='python_syntax_error',
            description=f'Python syntax error (might be obfuscated or Python 2): {e}',
            file=filename,
            line=e.lineno or 1,
            match=str(e),
        )]
    except Exception as e:
        return [Finding(
            severity=Severity.INFO,
            pattern_id='ast_parse_error',
            description=f'Could not parse Python AST: {e}',
            file=filename,
            line=1,
            match=str(e),
        )]
    
    visitor = DangerousCallVisitor(filename)
    visitor.visit(tree)
    
    return visitor.findings


def is_python_file(filename: str) -> bool:
    """Check if a file is a Python file."""
    return filename.endswith(('.py', '.pyw', '.pyi'))
