"""Enrich vulnerabilities with NVD, EPSS, and CISA KEV data."""

from __future__ import annotations

import asyncio
import json
import logging
import os
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import httpx
from rich.console import Console

from agent_bom.config import ENRICHMENT_MAX_CACHE_ENTRIES as _MAX_ENRICHMENT_CACHE_ENTRIES
from agent_bom.config import ENRICHMENT_TTL_SECONDS as _ENRICHMENT_TTL
from agent_bom.http_client import create_client, request_with_retry
from agent_bom.models import Vulnerability

console = Console(stderr=True)
_logger = logging.getLogger(__name__)

# API Endpoints
NVD_API_URL = "https://services.nvd.nist.gov/rest/json/cves/2.0"
EPSS_API_URL = "https://api.first.org/data/v1/epss"
CISA_KEV_URL = "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json"

# Cache for CISA KEV catalog (refresh daily, persisted to disk)
_kev_cache: Optional[dict] = None
_kev_cache_time: Optional[datetime] = None
_KEV_CACHE_FILE = Path.home() / ".agent-bom" / "kev_cache.json"
_KEV_CACHE_TTL_SECONDS = 86400  # 24 hours

# ─── Persistent enrichment cache (NVD + EPSS) ──────────────────────────────

_ENRICHMENT_CACHE_DIR = Path.home() / ".agent-bom"

# Module-level in-memory mirrors (loaded lazily from disk)
_nvd_file_cache: dict[str, dict] = {}
_epss_file_cache: dict[str, dict] = {}
_enrichment_cache_loaded = False


def _evict_oldest(cache: dict[str, dict], max_entries: int) -> None:
    """Evict oldest entries when cache exceeds max_entries."""
    if len(cache) <= max_entries:
        return
    # Evict entries missing _cached_at first (legacy/corrupt), then oldest
    by_age = sorted(
        cache.items(),
        key=lambda kv: int(kv[1].get("_cached_at") or 0) if isinstance(kv[1].get("_cached_at"), (int, float)) else 0,
    )
    to_remove = len(cache) - max_entries
    for key, _ in by_age[:to_remove]:
        del cache[key]


def _load_enrichment_cache() -> None:
    """Load NVD + EPSS caches from disk (once)."""
    global _nvd_file_cache, _epss_file_cache, _enrichment_cache_loaded  # noqa: PLW0603
    if _enrichment_cache_loaded:
        return
    _enrichment_cache_loaded = True
    now = time.time()
    for name, target in [("nvd_cache.json", "_nvd"), ("epss_cache.json", "_epss")]:
        path = _ENRICHMENT_CACHE_DIR / name
        if not path.exists():
            continue
        try:
            raw = json.loads(path.read_text())
            # Expire entries older than TTL
            fresh = {k: v for k, v in raw.items() if isinstance(v, dict) and (now - v.get("_cached_at", 0)) < _ENRICHMENT_TTL}
            if target == "_nvd":
                _nvd_file_cache.update(fresh)
            else:
                _epss_file_cache.update(fresh)
        except (OSError, json.JSONDecodeError, ValueError) as exc:
            _logger.warning("Failed to load enrichment cache %s: %s", name, exc)


def _save_enrichment_cache() -> None:
    """Persist NVD + EPSS caches to disk (atomic write to prevent corruption)."""
    _ENRICHMENT_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    for name, data in [("nvd_cache.json", _nvd_file_cache), ("epss_cache.json", _epss_file_cache)]:
        target = _ENRICHMENT_CACHE_DIR / name
        try:
            fd, tmp_path = tempfile.mkstemp(dir=str(_ENRICHMENT_CACHE_DIR), suffix=".tmp")
            fd_closed = False
            try:
                os.write(fd, json.dumps(data).encode("utf-8"))
                os.close(fd)
                fd_closed = True
                os.replace(tmp_path, str(target))
            except BaseException:
                if not fd_closed:
                    os.close(fd)
                try:
                    os.unlink(tmp_path)
                except OSError as cleanup_exc:
                    _logger.debug("Failed to remove temp cache file %s: %s", tmp_path, cleanup_exc)
                raise
        except OSError:
            _logger.warning("Failed to persist %s enrichment cache to disk", name)


async def fetch_nvd_data(cve_id: str, client: httpx.AsyncClient, api_key: Optional[str] = None) -> Optional[dict]:
    """Fetch CVE data from NVD API (with persistent file cache).

    Args:
        cve_id: CVE identifier (e.g., "CVE-2024-1234")
        client: HTTP client
        api_key: Optional NVD API key (recommended for higher rate limits)

    Returns:
        NVD vulnerability data or None if not found
    """
    _load_enrichment_cache()

    # Check persistent cache — reject entries older than 90 days
    _nvd_max_age_secs = 90 * 86400
    _now = time.time()
    if cve_id in _nvd_file_cache:
        cached = _nvd_file_cache[cve_id]
        cached_at = cached.get("_cached_at", 0)
        if _now - cached_at < _nvd_max_age_secs:
            data = {k: v for k, v in cached.items() if k != "_cached_at"}
            return data if data else None
        # Stale — fall through to refetch

    headers = {}
    if api_key:
        headers["apiKey"] = api_key

    response = await request_with_retry(
        client,
        "GET",
        NVD_API_URL,
        params={"cveId": cve_id},
        headers=headers,
    )

    if response and response.status_code == 200:
        try:
            data = response.json()
            vulnerabilities = data.get("vulnerabilities", [])
            if vulnerabilities:
                result = vulnerabilities[0].get("cve", {})
                # Store in persistent cache
                _nvd_file_cache[cve_id] = {**result, "_cached_at": time.time()}
                _evict_oldest(_nvd_file_cache, _MAX_ENRICHMENT_CACHE_ENTRIES)
                return result
        except (ValueError, KeyError) as e:
            console.print(f"  [dim yellow]NVD parse error for {cve_id}: {e}[/dim yellow]")
    elif response and response.status_code == 403:
        _logger.warning("NVD rate limited (HTTP 403) for %s — consider using NVD_API_KEY", cve_id)
    elif response and response.status_code not in (200, 404):
        _logger.warning("NVD returned HTTP %d for %s", response.status_code, cve_id)

    return None


async def fetch_epss_scores(cve_ids: list[str], client: httpx.AsyncClient) -> dict[str, dict]:
    """Fetch EPSS scores for multiple CVEs (with persistent file cache).

    EPSS (Exploit Prediction Scoring System) provides probability that a CVE
    will be exploited in the wild within the next 30 days.

    Args:
        cve_ids: List of CVE identifiers
        client: HTTP client

    Returns:
        Dictionary mapping CVE ID to EPSS data
    """
    if not cve_ids:
        return {}

    _load_enrichment_cache()

    scores: dict[str, dict] = {}
    uncached: list[str] = []

    # Check persistent cache first — reject entries older than 30 days
    _max_age_secs = 30 * 86400  # 30 days
    now = time.time()
    for cve_id in cve_ids:
        if cve_id in _epss_file_cache:
            cached = _epss_file_cache[cve_id]
            cached_at = cached.get("_cached_at", 0)
            if now - cached_at < _max_age_secs:
                scores[cve_id] = {k: v for k, v in cached.items() if k != "_cached_at"}
            else:
                uncached.append(cve_id)  # Stale — refetch
        else:
            uncached.append(cve_id)

    if not uncached:
        return scores

    # EPSS API accepts comma-separated CVE list — paginate in batches
    epss_batch_size = 100
    for batch_start in range(0, len(uncached), epss_batch_size):
        batch = uncached[batch_start : batch_start + epss_batch_size]
        cve_param = ",".join(batch)

        response = await request_with_retry(
            client,
            "GET",
            EPSS_API_URL,
            params={"cve": cve_param},
        )

        if response and response.status_code == 200:
            try:
                data = response.json()
                for item in data.get("data", []):
                    cve = item.get("cve")
                    if cve:
                        raw_epss = item.get("epss")
                        raw_pct = item.get("percentile")
                        parsed_score = float(raw_epss) if raw_epss is not None else None
                        parsed_pct = float(raw_pct) if raw_pct is not None else None
                        # Validate EPSS ranges — reject out-of-bounds values
                        if parsed_score is not None and not (0.0 <= parsed_score <= 1.0):
                            parsed_score = None
                        if parsed_pct is not None and not (0.0 <= parsed_pct <= 100.0):
                            parsed_pct = None
                        entry = {
                            "score": parsed_score,
                            "percentile": parsed_pct,
                            "date": item.get("date"),
                        }
                        scores[cve] = entry
                        _epss_file_cache[cve] = {**entry, "_cached_at": time.time()}
            except (ValueError, KeyError) as e:
                _logger.warning("EPSS parse error for batch starting at %d: %s", batch_start, e)
        else:
            _logger.warning(
                "EPSS API request failed for %d CVEs (batch %d–%d)",
                len(batch),
                batch_start,
                batch_start + len(batch),
            )

    _evict_oldest(_epss_file_cache, _MAX_ENRICHMENT_CACHE_ENTRIES)
    return scores


async def fetch_cisa_kev_catalog(client: httpx.AsyncClient) -> dict:
    """Fetch CISA Known Exploited Vulnerabilities catalog.

    CISA maintains a catalog of CVEs known to be actively exploited.
    These represent the highest priority vulnerabilities.

    Args:
        client: HTTP client

    Returns:
        Dictionary of KEV data indexed by CVE ID
    """
    global _kev_cache, _kev_cache_time

    # Use in-memory cache if less than 24 hours old
    if _kev_cache and _kev_cache_time:
        age = datetime.now(timezone.utc) - _kev_cache_time
        if age.total_seconds() < _KEV_CACHE_TTL_SECONDS:
            return _kev_cache

    # Try loading from disk cache if in-memory cache is stale
    if not _kev_cache and _KEV_CACHE_FILE.exists():
        try:
            disk_data = json.loads(_KEV_CACHE_FILE.read_text())
            cached_at = disk_data.get("_cached_at", 0)
            if (time.time() - cached_at) < _KEV_CACHE_TTL_SECONDS:
                _kev_cache = disk_data.get("data") or {}
                _kev_cache_time = datetime.fromtimestamp(cached_at, tz=timezone.utc)
                return _kev_cache  # type: ignore[return-value]  # dict from disk cache
        except (OSError, json.JSONDecodeError, ValueError) as exc:
            _logger.warning("Failed to load KEV disk cache: %s", exc)

    response = await request_with_retry(client, "GET", CISA_KEV_URL)

    if response and response.status_code == 200:
        try:
            data = response.json()
            kev_dict = {}

            for vuln in data.get("vulnerabilities", []):
                cve_id = vuln.get("cveID")
                if cve_id:
                    kev_dict[cve_id] = {
                        "date_added": vuln.get("dateAdded"),
                        "due_date": vuln.get("dueDate"),
                        "short_description": vuln.get("shortDescription"),
                        "required_action": vuln.get("requiredAction"),
                        "vendor_project": vuln.get("vendorProject"),
                        "product": vuln.get("product"),
                    }

            _kev_cache = kev_dict
            _kev_cache_time = datetime.now(timezone.utc)

            # Persist to disk for cross-session resilience
            try:
                _KEV_CACHE_FILE.parent.mkdir(parents=True, exist_ok=True)
                _KEV_CACHE_FILE.write_text(json.dumps({"_cached_at": time.time(), "data": kev_dict}))
            except OSError as exc:
                _logger.warning("Failed to persist KEV cache to disk: %s", exc)

            return kev_dict
        except (ValueError, KeyError) as e:
            console.print(f"  [dim yellow]CISA KEV parse error: {e}[/dim yellow]")

    # Fallback: serve stale disk cache if API is down
    if _KEV_CACHE_FILE.exists():
        try:
            disk_data = json.loads(_KEV_CACHE_FILE.read_text())
            stale_data = disk_data.get("data", {})
            if stale_data:
                stale_age_hours = (time.time() - disk_data.get("_cached_at", 0)) / 3600
                _logger.warning(
                    "CISA KEV API unreachable — serving stale cache (%.0fh old, %d entries)",
                    stale_age_hours,
                    len(stale_data),
                )
                console.print("  [dim yellow]⚠ Using stale CISA KEV cache (API unreachable)[/dim yellow]")
                _kev_cache = stale_data
                # Cache in memory for 1h before retrying API — prevents
                # hitting disk on every call within this session while still
                # retrying the API on the next scan (~1h later).
                from datetime import timedelta

                _kev_cache_time = datetime.now(timezone.utc) - timedelta(seconds=_KEV_CACHE_TTL_SECONDS - 3600)
                return stale_data
        except (OSError, json.JSONDecodeError) as exc:
            _logger.warning("Failed to read stale KEV cache: %s", exc)

    return {}


def extract_cve_ids(vulnerabilities: list[Vulnerability]) -> list[str]:
    """Extract CVE IDs from vulnerability list (including aliases)."""
    cve_ids: set[str] = set()
    for vuln in vulnerabilities:
        if vuln.id.startswith("CVE-"):
            cve_ids.add(vuln.id)
        for alias in vuln.aliases:
            if alias.startswith("CVE-"):
                cve_ids.add(alias)
    return list(cve_ids)


def calculate_exploitability(epss_score: Optional[float]) -> Optional[str]:
    """Calculate exploitability level from EPSS score.

    Thresholds configurable via ``AGENT_BOM_EPSS_CRITICAL_THRESHOLD``
    and ``AGENT_BOM_EPSS_HIGH_THRESHOLD``.
    """
    if epss_score is None:
        return None

    from agent_bom.config import EPSS_CRITICAL_THRESHOLD, EPSS_HIGH_LIKELY_THRESHOLD

    if epss_score >= EPSS_CRITICAL_THRESHOLD:
        return "HIGH"
    elif epss_score >= EPSS_HIGH_LIKELY_THRESHOLD:
        return "MEDIUM"
    else:
        return "LOW"


async def enrich_vulnerabilities(
    vulnerabilities: list[Vulnerability],
    nvd_api_key: Optional[str] = None,
    enable_nvd: bool = True,
    enable_epss: bool = True,
    enable_kev: bool = True,
) -> int:
    """Enrich vulnerabilities with NVD, EPSS, and CISA KEV data.

    Args:
        vulnerabilities: List of vulnerabilities to enrich
        nvd_api_key: Optional NVD API key for higher rate limits
        enable_nvd: Fetch NVD data (CWE, dates)
        enable_epss: Fetch EPSS exploit prediction scores
        enable_kev: Check CISA Known Exploited Vulnerabilities

    Returns:
        Number of vulnerabilities enriched
    """
    if not vulnerabilities:
        return 0

    cve_ids = extract_cve_ids(vulnerabilities)
    if not cve_ids:
        return 0

    console.print(f"\n[bold blue]🔬 Enriching {len(cve_ids)} CVE(s) with external data...[/bold blue]\n")

    enriched_count = 0

    async with create_client(timeout=30.0) as client:
        # ── Parallel enrichment: EPSS + KEV + NVD run concurrently ──
        # EPSS and KEV are single-call fetches. NVD needs CWE pre-check
        # (local, no network) then batched fetch. All three are independent.
        console.print("  [cyan]→[/cyan] Fetching EPSS + KEV + NVD in parallel...")

        async def _fetch_epss() -> dict:
            if not enable_epss:
                return {}
            data = await fetch_epss_scores(cve_ids, client)
            return data or {}

        async def _fetch_kev() -> dict:
            if not enable_kev:
                return {}
            return await fetch_cisa_kev_catalog(client)

        async def _fetch_nvd() -> tuple[dict[str, dict], int, int]:
            if not enable_nvd or not cve_ids:
                return {}, 0, 0
            # Build set of CVE IDs that already have CWE data (local check)
            cves_with_cwes: set[str] = set()
            for vuln in vulnerabilities:
                if vuln.cwe_ids:
                    if vuln.id.startswith("CVE-"):
                        cves_with_cwes.add(vuln.id)
                    for alias in vuln.aliases:
                        if alias.startswith("CVE-"):
                            cves_with_cwes.add(alias)

            nvd_needed = [c for c in cve_ids if c not in cves_with_cwes]
            skipped = len(cve_ids) - len(nvd_needed)
            if not nvd_needed:
                return {}, skipped, 0

            batch_size = 50 if nvd_api_key else 5
            sleep_secs = 1.0 if nvd_api_key else 6.0
            result_data: dict[str, dict] = {}

            for batch_start in range(0, len(nvd_needed), batch_size):
                batch = nvd_needed[batch_start : batch_start + batch_size]
                tasks = [fetch_nvd_data(cve_id, client, nvd_api_key) for cve_id in batch]
                results = await asyncio.gather(*tasks, return_exceptions=True)
                for cve_id, result in zip(batch, results):
                    if result and not isinstance(result, BaseException):
                        result_data[cve_id] = result
                if batch_start + batch_size < len(nvd_needed):
                    await asyncio.sleep(sleep_secs)

            return result_data, skipped, len(nvd_needed)

        # Run all three concurrently
        epss_data, kev_data, nvd_result = await asyncio.gather(_fetch_epss(), _fetch_kev(), _fetch_nvd())
        nvd_data, nvd_skipped, nvd_total = nvd_result

        # Report results
        if epss_data:
            console.print(f"  [green]✓[/green] EPSS: {len(epss_data)} CVE(s)")
        elif enable_epss:
            console.print("  [yellow]⚠[/yellow] EPSS data unavailable")

        kev_count = sum(1 for vuln in vulnerabilities if vuln.id in kev_data) if kev_data else 0
        if kev_count:
            console.print(f"  [red]⚠[/red] KEV: {kev_count} actively exploited CVE(s)!")
        elif enable_kev:
            console.print("  [green]✓[/green] KEV: no actively exploited CVEs")

        if nvd_skipped:
            console.print(f"  [dim]NVD: skipped {nvd_skipped} CVE(s) with existing CWE data[/dim]")
        if nvd_data:
            console.print(f"  [green]✓[/green] NVD: {len(nvd_data)}/{nvd_total} CVE(s)")
        elif enable_nvd and nvd_total:
            console.print("  [yellow]⚠[/yellow] NVD data unavailable")

        # Enrich each vulnerability — match by primary ID or CVE aliases
        for vuln in vulnerabilities:
            # Collect all CVE IDs this vuln maps to (primary + aliases)
            vuln_cve_ids = []
            if vuln.id.startswith("CVE-"):
                vuln_cve_ids.append(vuln.id)
            for alias in vuln.aliases:
                if alias.startswith("CVE-"):
                    vuln_cve_ids.append(alias)

            if not vuln_cve_ids:
                continue

            # Apply EPSS data (use first matching CVE)
            vuln_was_enriched = False
            for cve in vuln_cve_ids:
                if cve in epss_data:
                    epss = epss_data[cve]
                    vuln.epss_score = epss["score"]
                    vuln.epss_percentile = epss["percentile"]
                    vuln.exploitability = calculate_exploitability(epss["score"])
                    vuln_was_enriched = True
                    break

            # Apply CISA KEV data
            for cve in vuln_cve_ids:
                if cve in kev_data:
                    kev = kev_data[cve]
                    vuln.is_kev = True
                    vuln.kev_date_added = kev["date_added"]
                    vuln.kev_due_date = kev["due_date"]
                    vuln_was_enriched = True
                    break

            # Apply NVD data
            if enable_nvd:
                for cve in vuln_cve_ids:
                    if cve in nvd_data:
                        nvd = nvd_data[cve]

                        # Extract CWE IDs (deduplicated)
                        weaknesses = nvd.get("weaknesses", [])
                        existing_cwes = set(vuln.cwe_ids)
                        for weakness in weaknesses:
                            for desc in weakness.get("description", []):
                                cwe_val = desc.get("value", "")
                                if cwe_val.startswith("CWE-") and cwe_val not in existing_cwes:
                                    vuln.cwe_ids.append(cwe_val)
                                    existing_cwes.add(cwe_val)

                        # Extract dates
                        vuln.nvd_published = nvd.get("published")
                        vuln.nvd_modified = nvd.get("lastModified")

                        # Extract NVD vulnerability review status
                        vuln.nvd_status = nvd.get("vulnStatus")

                        # Merge NVD references with existing OSV references (deduplicated)
                        nvd_refs = nvd.get("references", [])
                        existing_urls = set(vuln.references)
                        for ref in nvd_refs:
                            url = ref.get("url")
                            if url and url not in existing_urls:
                                vuln.references.append(url)
                                existing_urls.add(url)

                        # Always include canonical NVD link as first reference
                        if cve.startswith("CVE-"):
                            canonical = f"https://nvd.nist.gov/vuln/detail/{cve}"
                            if canonical not in existing_urls:
                                vuln.references.insert(0, canonical)

                        vuln_was_enriched = True
                        break

            # Count each vulnerability only once regardless of how many sources enriched it
            if vuln_was_enriched:
                enriched_count += 1

    # Persist enrichment caches to disk
    _save_enrichment_cache()

    console.print(f"\n  [bold]Enriched {enriched_count} vulnerabilities.[/bold]")
    return enriched_count


def enrich_vulnerabilities_sync(
    vulnerabilities: list[Vulnerability],
    nvd_api_key: Optional[str] = None,
    enable_nvd: bool = True,
    enable_epss: bool = True,
    enable_kev: bool = True,
) -> int:
    """Synchronous wrapper for enrich_vulnerabilities."""
    return asyncio.run(
        enrich_vulnerabilities(
            vulnerabilities,
            nvd_api_key,
            enable_nvd,
            enable_epss,
            enable_kev,
        )
    )
