"""
Vulnerability data model for PyGenAI Security Framework
"""

from enum import Enum
from typing import List, Dict, Any, Optional
from datetime import datetime
import uuid

class ThreatLevel(Enum):
    """Threat level enumeration"""
    INFO = 1
    LOW = 2
    MEDIUM = 3
    HIGH = 4
    CRITICAL = 5

class VulnerabilityCategory(Enum):
    """Vulnerability category enumeration"""
    INJECTION = "injection"
    BROKEN_AUTH = "broken_authentication"
    SENSITIVE_DATA = "sensitive_data_exposure"
    XXE = "xml_external_entities"
    BROKEN_ACCESS = "broken_access_control"
    SECURITY_MISCONFIG = "security_misconfiguration"
    XSS = "cross_site_scripting"
    INSECURE_DESERIALIZATION = "insecure_deserialization"
    VULNERABLE_COMPONENTS = "vulnerable_components"
    LOGGING_MONITORING = "insufficient_logging_monitoring"
    
    # GenAI specific
    GENAI_PROMPT_INJECTION = "genai_prompt_injection"
    GENAI_DATA_LEAKAGE = "genai_data_leakage"
    GENAI_MODEL_MANIPULATION = "genai_model_manipulation"
    
    # Configuration
    CONFIGURATION_SECURITY = "configuration_security"
    SECRETS_MANAGEMENT = "secrets_management"

class Vulnerability:
    """Represents a security vulnerability"""
    
    def __init__(self, title: str, description: str = "", threat_level: ThreatLevel = ThreatLevel.MEDIUM,
                 category: VulnerabilityCategory = VulnerabilityCategory.INJECTION,
                 file_path: str = "", line_number: int = 0, code_snippet: str = "",
                 remediation: str = "", confidence: float = 0.8, cvss_score: float = 5.0,
                 cwe_id: str = "", owasp_category: str = "", scanner_name: str = "",
                 **kwargs):
        
        self.id = f"PYGENAI-{uuid.uuid4().hex[:8].upper()}"
        self.title = title
        self.description = description
        self.threat_level = threat_level
        self.category = category
        self.file_path = file_path
        self.line_number = line_number
        self.code_snippet = code_snippet
        self.remediation = remediation
        self.confidence = confidence
        self.cvss_score = cvss_score
        self.cwe_id = cwe_id
        self.owasp_category = owasp_category
        self.scanner_name = scanner_name
        self.first_detected = datetime.now().isoformat()
        
        # Additional fields from kwargs
        for key, value in kwargs.items():
            setattr(self, key, value)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert vulnerability to dictionary"""
        return {
            'id': self.id,
            'title': self.title,
            'description': self.description,
            'threat_level': self.threat_level.name.lower(),
            'category': self.category.value,
            'file_path': self.file_path,
            'line_number': self.line_number,
            'code_snippet': self.code_snippet,
            'remediation': self.remediation,
            'confidence': self.confidence,
            'cvss_score': self.cvss_score,
            'cwe_id': self.cwe_id,
            'owasp_category': self.owasp_category,
            'scanner_name': self.scanner_name,
            'first_detected': self.first_detected
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Vulnerability':
        """Create vulnerability from dictionary"""
        threat_level = ThreatLevel[data.get('threat_level', 'medium').upper()]
        category = VulnerabilityCategory(data.get('category', 'injection'))
        
        return cls(
            title=data.get('title', ''),
            description=data.get('description', ''),
            threat_level=threat_level,
            category=category,
            file_path=data.get('file_path', ''),
            line_number=data.get('line_number', 0),
            code_snippet=data.get('code_snippet', ''),
            remediation=data.get('remediation', ''),
            confidence=data.get('confidence', 0.8),
            cvss_score=data.get('cvss_score', 5.0),
            cwe_id=data.get('cwe_id', ''),
            owasp_category=data.get('owasp_category', ''),
            scanner_name=data.get('scanner_name', '')
        )

class VulnerabilityCollection:
    """Collection of vulnerabilities with analysis capabilities"""
    
    def __init__(self):
        self.vulnerabilities: List[Vulnerability] = []
    
    def add(self, vulnerability: Vulnerability):
        """Add vulnerability to collection"""
        self.vulnerabilities.append(vulnerability)
    
    def __len__(self):
        return len(self.vulnerabilities)
    
    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_statistics(self) -> Dict[str, Any]:
        """Get collection statistics"""
        stats = {
            'total_vulnerabilities': len(self.vulnerabilities),
            'by_threat_level': {},
            'by_category': {}
        }
        
        for threat_level in ThreatLevel:
            count = len(self.get_by_threat_level(threat_level))
            stats['by_threat_level'][threat_level.name.lower()] = count
        
        for vuln in self.vulnerabilities:
            category = vuln.category.value
            stats['by_category'][category] = stats['by_category'].get(category, 0) + 1
        
        return stats
