from typing import List, Dict
from concurrent.futures import ThreadPoolExecutor, as_completed
from models.vulnerability_intelligence import VulnerabilityIntelligence

from services.cache.cache_manager import CacheManager
from services.vulnerability_intelligence.enrichment.enrichment.cisa_kev import cisa_kev_enrich_vulnerability
from services.vulnerability_intelligence.enrichment.enrichment.trickest_cve_github import github_fetch_cve_data
from services.vulnerability_intelligence.enrichment.enrichment.trickest_cve_github_cached import github_fetch_cve_data_cached
from services.vulnerability_intelligence.enrichment.enrichment.github_poc import fetch_github_poc_data
from services.vulnerability_intelligence.enrichment.enrichment.github_poc_cached import fetch_github_poc_cached_data
from services.vulnerability_intelligence.enrichment.enrichment.nist_cached import nist_enrich_vulnerability
from services.vulnerability_intelligence.enrichment.enrichment.vulners import vulners_find_related_cve_data
from services.vulnerability_intelligence.handlers.github_poc_handler import GitHubPoCHandler
from services.vulnerability_intelligence.handlers.nist_cached_handler import NistCachedHandler
from services.vulnerability_intelligence.handlers.vulners_handler import VulnersHandler
from services.vulnerability_intelligence.handlers.github_handler import GitHubHandler
from services.vulnerability_intelligence.handlers.cisa_kev_handler import CisaKevHandler


class VulnerabilityIntelligenceEnrichmentManager:
    def __init__(self, vulnerability_intelligence_list: List[VulnerabilityIntelligence], enrichment_config: Dict, cache_manager: CacheManager):
        self.vulnerability_intelligence_list = vulnerability_intelligence_list
        self.enrichment_config = enrichment_config
        self.cache_manager = cache_manager

        self.all_enrichment_functions = {
            "vulners": vulners_find_related_cve_data,
            "trickest_cve_github": github_fetch_cve_data,
            "github_poc": fetch_github_poc_data,
            "trickest_cve_github_cached": lambda cve: github_fetch_cve_data_cached(cve, cache_manager),
            "github_poc_cached": lambda cve: fetch_github_poc_cached_data(cve, cache_manager),
            "cisa_kev": lambda cve: cisa_kev_enrich_vulnerability(cve, self.cache_manager),
            "nist_cached": lambda cve: nist_enrich_vulnerability(cve, self.cache_manager)
        }

        self.all_handlers = {
            "vulners": VulnersHandler,
            "trickest_cve_github": GitHubHandler,
            "trickest_cve_github_cached": GitHubHandler,
            "github_poc": GitHubPoCHandler,
            "github_poc_cached": GitHubPoCHandler,
            "cisa_kev": CisaKevHandler,
            "nist_cached": NistCachedHandler
        }

        self.enrichment_functions = {}
        self.handlers = {}
        self._filter_enrichment_sources()

    def _filter_enrichment_sources(self):
        enrichment_sources = self.enrichment_config.get("sources", {})

        if not enrichment_sources:
            print("[-] No enrichment sources defined in the config.")
        else:
            self.enrichment_functions = {
                name: func
                for name, func in self.all_enrichment_functions.items()
                if enrichment_sources.get(name, False)
            }
            self.handlers = {
                name: handler
                for name, handler in self.all_handlers.items()
                if name in self.enrichment_functions
            }

        print(f"\n\n[*] Enabled enrichment sources: {', '.join(list(self.enrichment_functions.keys()))}.\n")

    def enrich(self) -> List[VulnerabilityIntelligence]:
        if not self.enrichment_functions:
            print("[-] No enrichment sources are enabled. Skipping enrichment.")
            return self.vulnerability_intelligence_list

        with ThreadPoolExecutor(max_workers=10) as executor:
            future_to_vuln = {
                executor.submit(self._get_enrichment_data, vuln_intelligence): vuln_intelligence
                for vuln_intelligence in self.vulnerability_intelligence_list if vuln_intelligence.id.startswith("CVE-")
            }

            for future in as_completed(future_to_vuln):
                vuln_intelligence = future_to_vuln[future]
                try:
                    enrichment_data = future.result()
                    if enrichment_data:
                        self._apply_enrichment(vuln_intelligence, enrichment_data)
                except Exception as e:
                    print(f"[!] Error enriching data for {vuln_intelligence.id}: {e}")

        return self.vulnerability_intelligence_list

    def _get_enrichment_data(self, vuln_intelligence: VulnerabilityIntelligence) -> Dict:
        enrichment_data = {}

        for source_name, func in self.enrichment_functions.items():
            try:
                data = func(vuln_intelligence.id)
                enrichment_data[source_name] = data
            except Exception as e:
                print(f"[!] Error in {func.__name__} for CVE {vuln_intelligence.id}: {e}")

        return enrichment_data

    def _apply_enrichment(self, vuln_intelligence: VulnerabilityIntelligence, enrichment_data: Dict):
        for source, data in enrichment_data.items():
            handler_class = self.handlers.get(source)
            if handler_class:
                handler = handler_class(data)
                handler.apply(vuln_intelligence)
            else:
                print(f"[!] No handler found for source: {source}")