"""漏洞查询服务模块"""
import requests
from enum import Enum
from typing import List, Dict, Set
from concurrent.futures import ThreadPoolExecutor, as_completed

from ..models.vulnerability import Vulnerability
from ..models.dependency import Dependency
from requests.exceptions import ConnectionError, Timeout


class VulnerabilitySource(Enum):
    OSV = "osv"
    PYPI = "pypi"
    AI_SPECIFIC = "ai-specific"  # AI特定漏洞数据库


class VulnerabilityService:
    def query(self, package: str, version: str) -> List[Vulnerability]:
        raise NotImplementedError


class OsvService(VulnerabilityService):
    BASE_URL = "https://api.osv.dev/v1/query"
    TIMEOUT = 2

    def query(self, package: str, version: str) -> List[Vulnerability]:
        try:
            response = requests.post(
                self.BASE_URL,
                json={"package": {"name": package, "ecosystem": "PyPI"}, "version": version},
                timeout=self.TIMEOUT
            )
            response.raise_for_status()
            data = response.json()

            vulns = []
            for vuln in data.get("vulns", []):
                fix_versions = {event["fixed"] for affected in vuln.get("affected", [])
                               for range_info in affected.get("ranges", [])
                               for event in range_info.get("events", [])
                               if "fixed" in event}
                
                vulns.append(Vulnerability(
                    id=vuln["id"],
                    package=package,
                    version=version,
                    description=vuln.get("details", ""),
                    fix_versions=list(fix_versions),
                    aliases=set(vuln.get("aliases", [])),
                    ai_impact=""  
                ))
            return vulns

        except (ConnectionError, Timeout):
            #print("OSV服务连接超时")
            return []
        except Exception as e:
            print(f"OSV查询失败: {str(e)}")
            return []


class AISpecificVulnerabilityService(VulnerabilityService):
    """AI特定漏洞数据库查询服务"""
    # 对接专门的AI安全数据库
    AI_VULNERABILITIES = {
        ("tensorflow", "2.5.0"): [
            Vulnerability(
                id="AISEC-TF-2023-001",
                package="tensorflow",
                version="2.5.0",
                description="TensorFlow中的恶意模型可能导致代码执行漏洞，尤其在加载不可信模型时",
                fix_versions=["2.8.0", "2.7.1"],
                aliases=set(),
                ai_impact="可能导致模型加载时执行恶意代码"
            )
        ],
        ("transformers", "4.10.0"): [
            Vulnerability(
                id="AISEC-TR-2023-002",
                package="transformers",
                version="4.10.0",
                description="在处理特定输入时可能导致模型输出有害内容或泄露训练数据",
                fix_versions=["4.18.0"],
                aliases=set(),
                ai_impact="可能影响生成式AI模型的输出安全性"
            )
        ]
    }

    def query(self, package: str, version: str) -> List[Vulnerability]:
        """查询AI特定漏洞数据库"""
        return self.AI_VULNERABILITIES.get((package, version), [])


class VulnerabilityClient:
    """AI包漏洞查询客户端"""
    def __init__(self):
        self.services = {
            VulnerabilitySource.OSV: OsvService(),
            VulnerabilitySource.AI_SPECIFIC: AISpecificVulnerabilityService()
        }

    def query_vulnerabilities(self, package: str, version: str) -> List[Vulnerability]:
        """查询所有数据源的漏洞"""
        all_vulns = []
        for service in self.services.values():
            try:
                vulns = service.query(package, version)
                all_vulns.extend(vulns)
            except Exception as e:
                print(f"查询{package}漏洞时发生错误: {e}")
        
        return self._deduplicate(all_vulns)

    def _deduplicate(self, vulns: List[Vulnerability]) -> List[Vulnerability]:
        """去重漏洞"""
        seen = set()
        unique = []
        for vuln in vulns:
            if vuln.id not in seen:
                seen.add(vuln.id)
                unique.append(vuln)
        return unique

    def query_batch(self, dependencies: Set[Dependency], max_workers: int = 5) -> Dict[Dependency, List[Vulnerability]]:
        """并行查询多个依赖的漏洞"""
        results = {}
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_dep = {
                executor.submit(self.query_vulnerabilities, dep.name, dep.version): dep
                for dep in dependencies
            }
            
            total = len(future_to_dep)
            for i, future in enumerate(as_completed(future_to_dep), 1):
                dep = future_to_dep[future]
                try:
                    vulns = future.result()
                    results[dep] = vulns
                    print(f"[{i}/{total}] 完成 {dep} 的漏洞查询")
                except Exception as e:
                    print(f"[{i}/{total}] {dep} 查询出错: {str(e)}")
                    results[dep] = []
        
        return results
    