"""
Vulnerability data model for PyGenAI Security Framework
Comprehensive vulnerability representation with enterprise features.
"""

import time
import uuid
from dataclasses import dataclass, field
from enum import IntEnum, Enum
from typing import Dict, List, Any, Optional
from datetime import datetime, timezone


class ThreatLevel(IntEnum):
    """Threat levels with numeric values for comparison and sorting"""
    INFO = 1
    LOW = 2
    MEDIUM = 3
    HIGH = 4
    CRITICAL = 5
    
    @property
    def color_code(self) -> str:
        """Get color code for UI display"""
        colors = {
            ThreatLevel.INFO: "#17a2b8",      # Blue
            ThreatLevel.LOW: "#28a745",       # Green  
            ThreatLevel.MEDIUM: "#ffc107",    # Yellow
            ThreatLevel.HIGH: "#fd7e14",      # Orange
            ThreatLevel.CRITICAL: "#dc3545"   # Red
        }
        return colors.get(self, "#6c757d")
    
    @property
    def priority_score(self) -> int:
        """Get priority score for remediation ordering"""
        return self.value * 10


class VulnerabilityCategory(Enum):
    """Comprehensive vulnerability categories including GenAI-specific ones"""
    
    # OWASP Top 10 2023 Categories
    BROKEN_ACCESS_CONTROL = "broken_access_control"
    CRYPTOGRAPHIC_FAILURES = "cryptographic_failures"
    INJECTION = "injection"
    INSECURE_DESIGN = "insecure_design"
    SECURITY_MISCONFIGURATION = "security_misconfiguration"
    VULNERABLE_COMPONENTS = "vulnerable_components"
    IDENTIFICATION_AUTH_FAILURES = "identification_auth_failures"
    SOFTWARE_DATA_INTEGRITY = "software_data_integrity"
    LOGGING_MONITORING_FAILURES = "logging_monitoring_failures"
    SSRF = "server_side_request_forgery"
    
    # GenAI/LLM Specific Categories
    GENAI_PROMPT_INJECTION = "genai_prompt_injection"
    GENAI_DATA_LEAKAGE = "genai_data_leakage"
    GENAI_MODEL_MANIPULATION = "genai_model_manipulation"
    GENAI_BIAS_DISCRIMINATION = "genai_bias_discrimination"
    GENAI_PRIVACY_VIOLATION = "genai_privacy_violation"
    GENAI_MODEL_THEFT = "genai_model_theft"
    GENAI_TRAINING_DATA_POISONING = "genai_training_data_poisoning"
    
    # Additional Security Categories
    CODE_QUALITY = "code_quality"
    CONFIGURATION_SECURITY = "configuration_security"
    DEPENDENCY_VULNERABILITY = "dependency_vulnerability"
    SECRETS_MANAGEMENT = "secrets_management"
    
    @property
    def display_name(self) -> str:
        """Get human-readable display name"""
        return self.value.replace('_', ' ').title()
    
    @property
    def is_genai_related(self) -> bool:
        """Check if category is GenAI/LLM related"""
        return self.value.startswith('genai_')


@dataclass
class Vulnerability:
    """
    Comprehensive vulnerability representation for enterprise security scanning
    
    This class represents a security vulnerability with full metadata for
    enterprise reporting, compliance tracking, and remediation management.
    """
    
    # Core Identity
    id: str = field(default="")
    title: str = ""
    description: str = ""
    
    # Classification
    threat_level: ThreatLevel = ThreatLevel.MEDIUM
    category: VulnerabilityCategory = VulnerabilityCategory.CODE_QUALITY
    
    # Location Information
    file_path: str = ""
    line_number: int = 0
    column_number: int = 0
    end_line: int = 0
    end_column: int = 0
    
    # Code Context
    code_snippet: str = ""
    affected_code_block: str = ""
    function_name: str = ""
    class_name: str = ""
    
    # Security Analysis
    remediation: str = ""
    remediation_examples: List[str] = field(default_factory=list)
    business_impact: str = ""
    technical_impact: str = ""
    attack_vector: str = ""
    
    # Scoring and Risk Assessment
    cvss_score: float = 0.0
    cvss_vector: str = ""
    cwe_id: str = ""
    owasp_category: str = ""
    confidence: float = 1.0
    exploitability: float = 0.0
    
    # Compliance and Standards
    compliance_mappings: List[str] = field(default_factory=list)
    regulatory_impact: List[str] = field(default_factory=list)
    
    # Temporal Information
    first_detected: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    last_updated: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    detection_count: int = 1
    
    # Scanner Metadata
    scanner_name: str = "pygenai_security"
    scanner_version: str = "1.0.0"
    scan_rule_id: str = ""
    rule_description: str = ""
    
    # Enterprise Features
    assigned_to: str = ""
    priority: str = "medium"
    status: str = "open"  # open, in_progress, resolved, false_positive, accepted
    resolution_notes: str = ""
    tags: List[str] = field(default_factory=list)
    
    # Additional Metadata
    false_positive_probability: float = 0.0
    suppression_reason: str = ""
    external_references: List[str] = field(default_factory=list)
    
    def __post_init__(self):
        """Post-initialization processing"""
        if not self.id:
            self.id = self._generate_unique_id()
        
        if self.end_line == 0:
            self.end_line = self.line_number
        
        if self.end_column == 0:
            self.end_column = self.column_number
        
        self.last_updated = datetime.now(timezone.utc)
    
    def _generate_unique_id(self) -> str:
        """Generate unique vulnerability ID"""
        # Create deterministic but unique ID based on content
        content = f"{self.file_path}:{self.line_number}:{self.title}:{self.category.value}:{time.time()}"
        import hashlib
        hash_obj = hashlib.sha256(content.encode())
        return f"PYGENAI-{hash_obj.hexdigest()[:12].upper()}"
    
    @property
    def risk_score(self) -> float:
        """Calculate comprehensive risk score"""
        base_score = self.threat_level.value * 2
        cvss_weight = (self.cvss_score / 10.0) * 3
        exploitability_weight = self.exploitability * 2
        confidence_weight = self.confidence * 1.5
        
        # Adjust for GenAI vulnerabilities (emerging threat)
        genai_multiplier = 1.2 if self.category.is_genai_related else 1.0
        
        risk_score = (base_score + cvss_weight + exploitability_weight + confidence_weight) * genai_multiplier
        return min(risk_score, 25.0)  # Cap at 25
    
    @property
    def severity_label(self) -> str:
        """Get severity label for display"""
        return self.threat_level.name.title()
    
    @property
    def is_actionable(self) -> bool:
        """Check if vulnerability is actionable (not false positive or accepted)"""
        return self.status not in ['false_positive', 'accepted'] and self.confidence >= 0.5
    
    @property
    def remediation_priority(self) -> int:
        """Get remediation priority score"""
        priority_map = {
            'critical': 100,
            'high': 80,
            'medium': 60,
            'low': 40,
            'info': 20
        }
        
        base_priority = priority_map.get(self.priority.lower(), 60)
        
        # Adjust based on threat level
        threat_adjustment = self.threat_level.value * 5
        
        # Adjust based on business impact
        business_impact_adjustment = 10 if 'critical' in self.business_impact.lower() else 0
        
        return base_priority + threat_adjustment + business_impact_adjustment
    
    def update_status(self, new_status: str, notes: str = "", assigned_to: str = ""):
        """Update vulnerability status with audit trail"""
        valid_statuses = ['open', 'in_progress', 'resolved', 'false_positive', 'accepted']
        
        if new_status not in valid_statuses:
            raise ValueError(f"Invalid status: {new_status}. Must be one of {valid_statuses}")
        
        self.status = new_status
        self.resolution_notes = notes
        if assigned_to:
            self.assigned_to = assigned_to
        
        self.last_updated = datetime.now(timezone.utc)
    
    def add_tag(self, tag: str):
        """Add tag to vulnerability"""
        if tag not in self.tags:
            self.tags.append(tag)
            self.last_updated = datetime.now(timezone.utc)
    
    def remove_tag(self, tag: str):
        """Remove tag from vulnerability"""
        if tag in self.tags:
            self.tags.remove(tag)
            self.last_updated = datetime.now(timezone.utc)
    
    def mark_false_positive(self, reason: str = ""):
        """Mark vulnerability as false positive"""
        self.status = "false_positive"
        self.false_positive_probability = 1.0
        self.suppression_reason = reason
        self.last_updated = datetime.now(timezone.utc)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert vulnerability to dictionary for serialization"""
        return {
            # Core identity
            'id': self.id,
            'title': self.title,
            'description': self.description,
            
            # Classification
            'threat_level': self.threat_level.name.lower(),
            'threat_level_numeric': self.threat_level.value,
            'category': self.category.value,
            'category_display': self.category.display_name,
            'is_genai_related': self.category.is_genai_related,
            
            # Location
            'file_path': self.file_path,
            'line_number': self.line_number,
            'column_number': self.column_number,
            'end_line': self.end_line,
            'end_column': self.end_column,
            
            # Code context
            'code_snippet': self.code_snippet,
            'affected_code_block': self.affected_code_block,
            'function_name': self.function_name,
            'class_name': self.class_name,
            
            # Security analysis
            'remediation': self.remediation,
            'remediation_examples': self.remediation_examples,
            'business_impact': self.business_impact,
            'technical_impact': self.technical_impact,
            'attack_vector': self.attack_vector,
            
            # Scoring
            'cvss_score': self.cvss_score,
            'cvss_vector': self.cvss_vector,
            'cwe_id': self.cwe_id,
            'owasp_category': self.owasp_category,
            'confidence': self.confidence,
            'exploitability': self.exploitability,
            'risk_score': self.risk_score,
            
            # Compliance
            'compliance_mappings': self.compliance_mappings,
            'regulatory_impact': self.regulatory_impact,
            
            # Temporal
            'first_detected': self.first_detected.isoformat(),
            'last_updated': self.last_updated.isoformat(),
            'detection_count': self.detection_count,
            
            # Scanner metadata
            'scanner_name': self.scanner_name,
            'scanner_version': self.scanner_version,
            'scan_rule_id': self.scan_rule_id,
            'rule_description': self.rule_description,
            
            # Enterprise
            'assigned_to': self.assigned_to,
            'priority': self.priority,
            'status': self.status,
            'resolution_notes': self.resolution_notes,
            'tags': self.tags,
            'remediation_priority': self.remediation_priority,
            
            # Quality metrics
            'false_positive_probability': self.false_positive_probability,
            'suppression_reason': self.suppression_reason,
            'is_actionable': self.is_actionable,
            'external_references': self.external_references
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Vulnerability':
        """Create vulnerability from dictionary"""
        # Convert string enums back to enum objects
        if 'threat_level' in data:
            if isinstance(data['threat_level'], str):
                data['threat_level'] = ThreatLevel[data['threat_level'].upper()]
        
        if 'category' in data:
            if isinstance(data['category'], str):
                # Find matching category
                for cat in VulnerabilityCategory:
                    if cat.value == data['category']:
                        data['category'] = cat
                        break
        
        # Convert datetime strings back to datetime objects
        for field in ['first_detected', 'last_updated']:
            if field in data and isinstance(data[field], str):
                data[field] = datetime.fromisoformat(data[field])
        
        # Remove computed properties
        computed_fields = ['threat_level_numeric', 'category_display', 'is_genai_related', 
                          'risk_score', 'remediation_priority', 'is_actionable']
        for field in computed_fields:
            data.pop(field, None)
        
        return cls(**data)
    
    def __str__(self) -> str:
        """String representation for logging and debugging"""
        return f"Vulnerability(id={self.id}, title='{self.title}', threat_level={self.threat_level.name}, file='{self.file_path}:{self.line_number}')"
    
    def __repr__(self) -> str:
        """Detailed representation for debugging"""
        return (f"Vulnerability(id='{self.id}', title='{self.title}', "
                f"threat_level={self.threat_level}, category={self.category}, "
                f"file_path='{self.file_path}', line_number={self.line_number})")


class VulnerabilityCollection:
    """
    Collection class for managing multiple vulnerabilities with enterprise features
    """
    
    def __init__(self):
        self.vulnerabilities: List[Vulnerability] = []
        self._index_by_id: Dict[str, Vulnerability] = {}
        self._index_by_file: Dict[str, List[Vulnerability]] = {}
        self._index_by_category: Dict[VulnerabilityCategory, List[Vulnerability]] = {}
    
    def add(self, vulnerability: Vulnerability):
        """Add vulnerability to collection with indexing"""
        if vulnerability.id in self._index_by_id:
            # Update existing vulnerability
            self.update(vulnerability)
        else:
            # Add new vulnerability
            self.vulnerabilities.append(vulnerability)
            self._index_by_id[vulnerability.id] = vulnerability
            
            # Update file index
            if vulnerability.file_path not in self._index_by_file:
                self._index_by_file[vulnerability.file_path] = []
            self._index_by_file[vulnerability.file_path].append(vulnerability)
            
            # Update category index
            if vulnerability.category not in self._index_by_category:
                self._index_by_category[vulnerability.category] = []
            self._index_by_category[vulnerability.category].append(vulnerability)
    
    def update(self, vulnerability: Vulnerability):
        """Update existing vulnerability"""
        if vulnerability.id in self._index_by_id:
            old_vuln = self._index_by_id[vulnerability.id]
            
            # Update in main list
            index = self.vulnerabilities.index(old_vuln)
            self.vulnerabilities[index] = vulnerability
            
            # Update indexes
            self._index_by_id[vulnerability.id] = vulnerability
            
            # Rebuild file and category indexes (simple approach)
            self._rebuild_indexes()
    
    def remove(self, vulnerability_id: str):
        """Remove vulnerability by ID"""
        if vulnerability_id in self._index_by_id:
            vulnerability = self._index_by_id[vulnerability_id]
            self.vulnerabilities.remove(vulnerability)
            del self._index_by_id[vulnerability_id]
            self._rebuild_indexes()
    
    def get_by_id(self, vulnerability_id: str) -> Optional[Vulnerability]:
        """Get vulnerability by ID"""
        return self._index_by_id.get(vulnerability_id)
    
    def get_by_file(self, file_path: str) -> List[Vulnerability]:
        """Get vulnerabilities by file path"""
        return self._index_by_file.get(file_path, [])
    
    def get_by_category(self, category: VulnerabilityCategory) -> List[Vulnerability]:
        """Get vulnerabilities by category"""
        return self._index_by_category.get(category, [])
    
    def get_by_threat_level(self, threat_level: ThreatLevel) -> List[Vulnerability]:
        """Get vulnerabilities by threat level"""
        return [v for v in self.vulnerabilities if v.threat_level == threat_level]
    
    def get_actionable(self) -> List[Vulnerability]:
        """Get actionable vulnerabilities (excluding false positives and accepted)"""
        return [v for v in self.vulnerabilities if v.is_actionable]
    
    def get_genai_related(self) -> List[Vulnerability]:
        """Get GenAI-related vulnerabilities"""
        return [v for v in self.vulnerabilities if v.category.is_genai_related]
    
    def _rebuild_indexes(self):
        """Rebuild file and category indexes"""
        self._index_by_file.clear()
        self._index_by_category.clear()
        
        for vuln in self.vulnerabilities:
            # File index
            if vuln.file_path not in self._index_by_file:
                self._index_by_file[vuln.file_path] = []
            self._index_by_file[vuln.file_path].append(vuln)
            
            # Category index
            if vuln.category not in self._index_by_category:
                self._index_by_category[vuln.category] = []
            self._index_by_category[vuln.category].append(vuln)
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get comprehensive statistics"""
        if not self.vulnerabilities:
            return {
                'total_vulnerabilities': 0,
                'by_threat_level': {},
                'by_category': {},
                'by_status': {},
                'affected_files': 0,
                'genai_related_count': 0
            }
        
        stats = {
            'total_vulnerabilities': len(self.vulnerabilities),
            'by_threat_level': {},
            'by_category': {},
            'by_status': {},
            'affected_files': len(self._index_by_file),
            'genai_related_count': len(self.get_genai_related()),
            'actionable_count': len(self.get_actionable())
        }
        
        # Count by threat level
        for level in ThreatLevel:
            stats['by_threat_level'][level.name.lower()] = len(self.get_by_threat_level(level))
        
        # Count by category
        for category in VulnerabilityCategory:
            count = len(self.get_by_category(category))
            if count > 0:
                stats['by_category'][category.value] = count
        
        # Count by status
        status_counts = {}
        for vuln in self.vulnerabilities:
            status_counts[vuln.status] = status_counts.get(vuln.status, 0) + 1
        stats['by_status'] = status_counts
        
        return stats
    
    def to_list(self) -> List[Dict[str, Any]]:
        """Convert collection to list of dictionaries"""
        return [vuln.to_dict() for vuln in self.vulnerabilities]
    
    def __len__(self) -> int:
        """Get number of vulnerabilities"""
        return len(self.vulnerabilities)
    
    def __iter__(self):
        """Iterate over vulnerabilities"""
        return iter(self.vulnerabilities)
    
    def __contains__(self, vulnerability_id: str) -> bool:
        """Check if vulnerability ID exists in collection"""
        return vulnerability_id in self._index_by_id
