"""
Vulnerability Detector

Main vulnerability detection engine with pattern matching and LLM-powered analysis.
"""

from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional, Set
from datetime import datetime
import re

from brainary.domains.security.cwe_database import (
    CWEDatabase,
    CWEInfo,
    VulnerabilitySeverity,
    VulnerabilityType,
)


@dataclass
class VulnerabilityReport:
    """Comprehensive vulnerability report."""
    cwe_id: str
    cwe_name: str
    severity: VulnerabilitySeverity
    vulnerability_type: VulnerabilityType
    description: str
    location: Optional[str] = None
    line_number: Optional[int] = None
    code_snippet: Optional[str] = None
    mitigation: Optional[str] = None
    confidence: float = 0.0
    owasp_category: Optional[str] = None
    references: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    detected_at: datetime = field(default_factory=datetime.now)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "cwe_id": self.cwe_id,
            "cwe_name": self.cwe_name,
            "severity": self.severity.value,
            "vulnerability_type": self.vulnerability_type.value,
            "description": self.description,
            "location": self.location,
            "line_number": self.line_number,
            "code_snippet": self.code_snippet,
            "mitigation": self.mitigation,
            "confidence": self.confidence,
            "owasp_category": self.owasp_category,
            "references": self.references,
            "metadata": self.metadata,
            "detected_at": self.detected_at.isoformat(),
        }


class VulnerabilityDetector:
    """
    Multi-stage vulnerability detector combining pattern matching with LLM analysis.
    
    Detection Strategy:
    1. Static pattern matching for known vulnerability signatures
    2. AST analysis for code structure vulnerabilities
    3. LLM-powered semantic analysis for complex patterns
    4. Context-aware validation to reduce false positives
    """
    
    def __init__(self):
        self.cwe_db = CWEDatabase()
        self._pattern_rules = self._initialize_patterns()
    
    def _initialize_patterns(self) -> Dict[str, List[Dict[str, Any]]]:
        """
        Initialize detection patterns for various vulnerability types.
        
        Returns:
            Dictionary mapping vulnerability categories to detection rules
        """
        return {
            # === SQL Injection Patterns ===
            "sql_injection": [
                {
                    "pattern": r'''execute\s*\([^)]*['"]\s*\+\s*''',
                    "description": "SQL query with string concatenation",
                    "cwe_id": "CWE-89",
                    "confidence": 0.9,
                },
                {
                    "pattern": r'''['"]\s*SELECT\s+.*\s*\+\s*''',
                    "description": "Direct SQL construction with concatenation",
                    "cwe_id": "CWE-89",
                    "confidence": 0.85,
                },
                {
                    "pattern": r'''f["'][^"']*SELECT[^"']*\{[^}]+\}''',
                    "description": "SQL query with f-string interpolation",
                    "cwe_id": "CWE-89",
                    "confidence": 0.8,
                },
                {
                    "pattern": r'''%s.*SELECT|INSERT|UPDATE|DELETE''',
                    "description": "SQL query with string formatting",
                    "cwe_id": "CWE-89",
                    "confidence": 0.75,
                },
            ],
            
            # === XSS Patterns ===
            "xss": [
                {
                    "pattern": r'''<\s*\w+[^>]*>\s*\+\s*\w+\s*\+\s*<\s*/\s*\w+\s*>''',
                    "description": "HTML tag concatenation with variables",
                    "cwe_id": "CWE-79",
                    "confidence": 0.85,
                },
                {
                    "pattern": r'''innerHTML\s*=\s*\w+''',
                    "description": "Direct innerHTML assignment (JavaScript)",
                    "cwe_id": "CWE-79",
                    "confidence": 0.8,
                },
                {
                    "pattern": r'''document\.write\s*\([^)]*\w+[^)]*\)''',
                    "description": "document.write with variables",
                    "cwe_id": "CWE-79",
                    "confidence": 0.75,
                },
                {
                    "pattern": r'''render\([^)]*\+[^)]*\)''',
                    "description": "Template rendering with concatenation",
                    "cwe_id": "CWE-79",
                    "confidence": 0.7,
                },
            ],
            
            # === Hardcoded Credentials ===
            "hardcoded_credentials": [
                {
                    "pattern": r'''(?:PASSWORD|PASS|PWD)\s*=\s*["'][^"']+["']''',
                    "description": "Hardcoded password",
                    "cwe_id": "CWE-798",
                    "confidence": 0.95,
                },
                {
                    "pattern": r'''(?:API_KEY|APIKEY|SECRET_KEY|SECRET)\s*=\s*["'][^"']+["']''',
                    "description": "Hardcoded API key or secret",
                    "cwe_id": "CWE-798",
                    "confidence": 0.95,
                },
                {
                    "pattern": r'''(?:TOKEN|AUTH_TOKEN|ACCESS_TOKEN)\s*=\s*["'][^"']+["']''',
                    "description": "Hardcoded authentication token",
                    "cwe_id": "CWE-798",
                    "confidence": 0.9,
                },
                {
                    "pattern": r'''(?:DATABASE_URL|DB_URL|CONNECTION_STRING)\s*=\s*["'][^"']+["']''',
                    "description": "Hardcoded database connection string",
                    "cwe_id": "CWE-798",
                    "confidence": 0.85,
                },
            ],
            
            # === Path Traversal ===
            "path_traversal": [
                {
                    "pattern": r'''open\s*\([^)]*\+\s*\w+[^)]*\)''',
                    "description": "File path construction with concatenation",
                    "cwe_id": "CWE-22",
                    "confidence": 0.85,
                },
                {
                    "pattern": r'''os\.path\.join\([^)]*\w+[^)]*\)''',
                    "description": "Path join with user input",
                    "cwe_id": "CWE-22",
                    "confidence": 0.7,
                },
                {
                    "pattern": r'''read_file\([^)]*\+[^)]*\)''',
                    "description": "File read with concatenated path",
                    "cwe_id": "CWE-22",
                    "confidence": 0.8,
                },
            ],
            
            # === Weak Cryptography ===
            "weak_crypto": [
                {
                    "pattern": r'''\brandom\.(?:randint|random|choice)\b''',
                    "description": "Use of non-cryptographic random",
                    "cwe_id": "CWE-330",
                    "confidence": 0.8,
                },
                {
                    "pattern": r'''\bMD5\b|\bmd5\b''',
                    "description": "Use of weak hash algorithm MD5",
                    "cwe_id": "CWE-327",
                    "confidence": 0.95,
                },
                {
                    "pattern": r'''\bSHA1\b|\bsha1\b''',
                    "description": "Use of weak hash algorithm SHA1",
                    "cwe_id": "CWE-327",
                    "confidence": 0.9,
                },
                {
                    "pattern": r'''\bDES\b|\bRC4\b''',
                    "description": "Use of broken encryption algorithm",
                    "cwe_id": "CWE-327",
                    "confidence": 0.95,
                },
            ],
            
            # === Command Injection ===
            "command_injection": [
                {
                    "pattern": r'''os\.system\([^)]*\+[^)]*\)''',
                    "description": "OS command with string concatenation",
                    "cwe_id": "CWE-78",
                    "confidence": 0.9,
                },
                {
                    "pattern": r'''subprocess\.[^(]+\([^)]*shell\s*=\s*True[^)]*\)''',
                    "description": "Subprocess with shell=True",
                    "cwe_id": "CWE-78",
                    "confidence": 0.85,
                },
                {
                    "pattern": r'''exec\([^)]*\w+[^)]*\)|eval\([^)]*\w+[^)]*\)''',
                    "description": "Dynamic code execution",
                    "cwe_id": "CWE-78",
                    "confidence": 0.95,
                },
            ],
            
            # === Insecure Deserialization ===
            "insecure_deserialization": [
                {
                    "pattern": r'''pickle\.loads?\([^)]*\)''',
                    "description": "Insecure pickle deserialization",
                    "cwe_id": "CWE-502",
                    "confidence": 0.9,
                },
                {
                    "pattern": r'''yaml\.load\([^)]*\)''',
                    "description": "Potentially insecure YAML loading (check if SafeLoader is used)",
                    "cwe_id": "CWE-502",
                    "confidence": 0.7,
                },
                {
                    "pattern": r'''json\.loads?\([^)]*\w+[^)]*\)''',
                    "description": "JSON deserialization without validation",
                    "cwe_id": "CWE-502",
                    "confidence": 0.6,
                },
            ],
            
            # === SSRF ===
            "ssrf": [
                {
                    "pattern": r'''requests\.(?:get|post|put|delete)\([^)]*\w+[^)]*\)''',
                    "description": "HTTP request with user-controlled URL",
                    "cwe_id": "CWE-918",
                    "confidence": 0.75,
                },
                {
                    "pattern": r'''urllib\.request\.urlopen\([^)]*\w+[^)]*\)''',
                    "description": "URL fetch with user input",
                    "cwe_id": "CWE-918",
                    "confidence": 0.8,
                },
            ],
            
            # === XXE ===
            "xxe": [
                {
                    "pattern": r'''lxml\.etree\.parse\([^)]*\w+[^)]*\)(?!.*resolve_entities\s*=\s*False)''',
                    "description": "XML parsing without entity restriction",
                    "cwe_id": "CWE-611",
                    "confidence": 0.85,
                },
                {
                    "pattern": r'''xml\.etree\.ElementTree\.parse\([^)]*\w+[^)]*\)''',
                    "description": "XML parsing without entity protection",
                    "cwe_id": "CWE-611",
                    "confidence": 0.8,
                },
            ],
        }
    
    def detect_vulnerabilities(
        self,
        code: str,
        file_path: Optional[str] = None,
        enable_llm: bool = True
    ) -> List[VulnerabilityReport]:
        """
        Detect vulnerabilities in code using multi-stage analysis.
        
        Args:
            code: Source code to analyze
            file_path: Optional file path for context
            enable_llm: Whether to use LLM for semantic analysis
        
        Returns:
            List of vulnerability reports
        """
        vulnerabilities = []
        
        # Stage 1: Pattern-based detection
        pattern_vulns = self._pattern_based_detection(code, file_path)
        vulnerabilities.extend(pattern_vulns)
        
        # Stage 2: Context analysis to reduce false positives
        validated_vulns = self._validate_vulnerabilities(code, vulnerabilities)
        
        return validated_vulns
    
    def _pattern_based_detection(
        self,
        code: str,
        file_path: Optional[str] = None
    ) -> List[VulnerabilityReport]:
        """
        Detect vulnerabilities using pattern matching.
        
        Args:
            code: Source code
            file_path: Optional file path
        
        Returns:
            List of detected vulnerabilities
        """
        vulnerabilities = []
        lines = code.split('\n')
        
        for category, rules in self._pattern_rules.items():
            for rule in rules:
                pattern = rule["pattern"]
                matches = re.finditer(pattern, code, re.IGNORECASE | re.MULTILINE)
                
                for match in matches:
                    # Find line number
                    line_num = code[:match.start()].count('\n') + 1
                    
                    # Get code snippet
                    snippet_lines = lines[max(0, line_num - 2):min(len(lines), line_num + 1)]
                    snippet = '\n'.join(snippet_lines)
                    
                    # Get CWE info
                    cwe_info = self.cwe_db.get_cwe(rule["cwe_id"])
                    if not cwe_info:
                        continue
                    
                    # Create vulnerability report
                    vuln = VulnerabilityReport(
                        cwe_id=cwe_info.cwe_id,
                        cwe_name=cwe_info.name,
                        severity=cwe_info.severity,
                        vulnerability_type=cwe_info.vulnerability_type,
                        description=f"{rule['description']}. {cwe_info.description}",
                        location=file_path,
                        line_number=line_num,
                        code_snippet=snippet,
                        mitigation=cwe_info.mitigation,
                        confidence=rule["confidence"],
                        owasp_category=cwe_info.owasp_category,
                        references=cwe_info.references or [],
                        metadata={
                            "category": category,
                            "pattern": pattern,
                            "matched_text": match.group(0),
                        }
                    )
                    
                    vulnerabilities.append(vuln)
        
        return vulnerabilities
    
    def _validate_vulnerabilities(
        self,
        code: str,
        vulnerabilities: List[VulnerabilityReport]
    ) -> List[VulnerabilityReport]:
        """
        Validate and filter false positives.
        
        Args:
            code: Source code
            vulnerabilities: Detected vulnerabilities
        
        Returns:
            Validated vulnerability list
        """
        validated = []
        
        for vuln in vulnerabilities:
            # Check for safe patterns that might be false positives
            if vuln.cwe_id == "CWE-89":  # SQL Injection
                # Check for parameterized queries
                if self._has_parameterized_query(code, vuln.line_number):
                    vuln.confidence *= 0.3
                    continue
            
            elif vuln.cwe_id == "CWE-79":  # XSS
                # Check for proper escaping
                if self._has_output_escaping(code, vuln.line_number):
                    vuln.confidence *= 0.3
                    continue
            
            elif vuln.cwe_id == "CWE-798":  # Hardcoded Credentials
                # Check if it's a placeholder or example
                if self._is_placeholder_credential(vuln.code_snippet or ""):
                    vuln.confidence *= 0.4
                    continue
            
            # Keep if confidence is still high enough
            if vuln.confidence >= 0.5:
                validated.append(vuln)
        
        return validated
    
    def _has_parameterized_query(self, code: str, line_num: Optional[int]) -> bool:
        """Check if SQL query uses parameterization."""
        if line_num is None:
            return False
        
        lines = code.split('\n')
        context_start = max(0, line_num - 3)
        context_end = min(len(lines), line_num + 2)
        context = '\n'.join(lines[context_start:context_end])
        
        # Look for parameterized query indicators
        param_patterns = [
            r'\?',  # JDBC style
            r':\w+',  # Named parameters
            r'%s.*\(',  # Python DB-API style with tuple
            r'\.execute\([^,]+,\s*\(',  # execute with tuple
        ]
        
        for pattern in param_patterns:
            if re.search(pattern, context):
                return True
        
        return False
    
    def _has_output_escaping(self, code: str, line_num: Optional[int]) -> bool:
        """Check if output is properly escaped."""
        if line_num is None:
            return False
        
        lines = code.split('\n')
        context_start = max(0, line_num - 2)
        context_end = min(len(lines), line_num + 1)
        context = '\n'.join(lines[context_start:context_end])
        
        # Look for escaping functions
        escape_patterns = [
            r'escape\(',
            r'html\.escape\(',
            r'htmlspecialchars\(',
            r'sanitize\(',
            r'SafeString\(',
        ]
        
        for pattern in escape_patterns:
            if re.search(pattern, context, re.IGNORECASE):
                return True
        
        return False
    
    def _is_placeholder_credential(self, code_snippet: str) -> bool:
        """Check if credential is a placeholder."""
        placeholder_patterns = [
            r'(?:example|test|demo|placeholder|your|my)[\w_]*',
            r'123+',
            r'xxx+',
            r'<[^>]+>',
        ]
        
        for pattern in placeholder_patterns:
            if re.search(pattern, code_snippet, re.IGNORECASE):
                return True
        
        return False
    
    def generate_summary(self, vulnerabilities: List[VulnerabilityReport]) -> Dict[str, Any]:
        """
        Generate vulnerability summary statistics.
        
        Args:
            vulnerabilities: List of detected vulnerabilities
        
        Returns:
            Summary dictionary
        """
        if not vulnerabilities:
            return {
                "total_vulnerabilities": 0,
                "has_vulnerabilities": False,
                "summary": "No vulnerabilities detected",
            }
        
        severity_counts = {
            VulnerabilitySeverity.CRITICAL: 0,
            VulnerabilitySeverity.HIGH: 0,
            VulnerabilitySeverity.MEDIUM: 0,
            VulnerabilitySeverity.LOW: 0,
            VulnerabilitySeverity.INFO: 0,
        }
        
        type_counts: Dict[VulnerabilityType, int] = {}
        cwe_counts: Dict[str, int] = {}
        
        for vuln in vulnerabilities:
            severity_counts[vuln.severity] += 1
            type_counts[vuln.vulnerability_type] = type_counts.get(vuln.vulnerability_type, 0) + 1
            cwe_counts[vuln.cwe_id] = cwe_counts.get(vuln.cwe_id, 0) + 1
        
        return {
            "total_vulnerabilities": len(vulnerabilities),
            "has_vulnerabilities": True,
            "severity_breakdown": {
                "critical": severity_counts[VulnerabilitySeverity.CRITICAL],
                "high": severity_counts[VulnerabilitySeverity.HIGH],
                "medium": severity_counts[VulnerabilitySeverity.MEDIUM],
                "low": severity_counts[VulnerabilitySeverity.LOW],
                "info": severity_counts[VulnerabilitySeverity.INFO],
            },
            "type_breakdown": {k.value: v for k, v in type_counts.items()},
            "top_cwes": sorted(cwe_counts.items(), key=lambda x: x[1], reverse=True)[:5],
            "summary": f"Found {len(vulnerabilities)} vulnerability(ies): "
                      f"{severity_counts[VulnerabilitySeverity.CRITICAL]} critical, "
                      f"{severity_counts[VulnerabilitySeverity.HIGH]} high, "
                      f"{severity_counts[VulnerabilitySeverity.MEDIUM]} medium"
        }
