"""漏洞评估器模块"""
from typing import Dict, List

from ..models.vulnerability import Vulnerability
from ..models.ai_domain import AIDomain


class AIVulnerabilityAssessor:
    """AI漏洞评估器，使用AI模型评估漏洞风险"""
    
    def __init__(self):
        self.ai_risk_factors = {
            # AI包特定风险因素
            "tensorflow": ["模型训练安全", "推理效率漏洞", "GPU利用漏洞"],
            "pytorch": ["张量操作安全", "自动微分漏洞", "分布式训练风险"],
            "transformers": ["模型加载安全", " tokenizer漏洞", "生成内容风险"],
            "opencv-python": ["图像处理漏洞", "视频流安全", "CVE漏洞利用风险"]
        }
        
        # 初始化轻量级风险评估模型（实际使用中可替换为更复杂的AI模型）
        self.risk_model = self._initialize_risk_model()

    def _initialize_risk_model(self) -> Dict[str, List[str]]:
        """初始化风险评估模型"""
        # 这里使用简单规则引擎模拟AI风险评估模型
        # 实际应用中可以替换为训练好的机器学习模型
        return {
            "high_risk_keywords": [
                "远程代码执行", "任意代码执行", "数据泄露", 
                "权限提升", " denial of service", "内存泄漏"
            ],
            "medium_risk_keywords": [
                "性能下降", "精度损失", "资源消耗过高", 
                "兼容性问题", "数值不稳定性"
            ]
        }

    def assess_vulnerability_risk(self, vulnerability: Vulnerability, package: str) -> float:
        """评估漏洞风险分数（0-10）"""
        score = 0.0
        desc = vulnerability.description.lower()
        
        # 基于关键词的风险评估（模拟AI模型推理）
        for kw in self.risk_model["high_risk_keywords"]:
            if kw in desc:
                score += 4.0
                
        for kw in self.risk_model["medium_risk_keywords"]:
            if kw in desc:
                score += 2.0
                
        # AI包特定风险加成
        if package in self.ai_risk_factors:
            score += 2.0  # AI包基础风险加成
            
        # 限制最高分10分
        return min(score, 10.0)

    def analyze_ai_impact(self, vulnerability: Vulnerability, package: str, domain: AIDomain) -> str:
        """分析漏洞对AI应用的具体影响"""
        impact_map = {
            AIDomain.NATURAL_LANGUAGE: "可能导致文本生成错误、模型偏见放大或输入处理漏洞",
            AIDomain.COMPUTER_VISION: "可能影响图像识别准确性、导致目标检测错误或视频处理漏洞",
            AIDomain.GENERATIVE_MODEL: "可能导致生成内容质量下降、泄露训练数据或产生有害输出",
            AIDomain.REINFORCEMENT_LEARNING: "可能影响智能体决策、奖励函数漏洞或环境交互问题"
        }
        
        base_impact = impact_map.get(domain, "可能影响AI模型的稳定性和安全性")
        
        # 根据包类型细化影响描述
        if package in ["tensorflow", "pytorch", "torch"]:
            return f"作为核心深度学习框架，此漏洞{base_impact}，可能影响模型训练和推理全过程"
        elif package in ["transformers", "diffusers"]:
            return f"作为预训练模型库，此漏洞{base_impact}，可能导致模型输出异常或安全风险"
        else:
            return base_impact

    def generate_mitigation(self, vulnerability: Vulnerability, package: str) -> str:
        """生成AI领域特定的缓解建议"""
        if vulnerability.fix_versions:
            fix建议 = f"建议升级到安全版本: {', '.join(vulnerability.fix_versions)}"
        else:
            fix建议 = "目前没有官方修复版本，建议密切关注官方更新"
            
        # AI包特定缓解措施
        if package in ["tensorflow", "pytorch"]:
            return f"{fix建议}。同时，在升级前可考虑使用沙箱环境测试模型兼容性，确保升级后模型性能不受影响。"
        elif package in ["transformers"]:
            return f"{fix建议}。对于生产环境，可临时部署输入验证机制，过滤可能触发漏洞的输入。"
        else:
            return fix建议
    