"""
In-memory job tracker for long-running CVSS-override runs.

The CISA Vulnrichment override against a full DB takes minutes (even
after the ZIP-snapshot speedup) — the browser's HTTP client times
out long before the work completes, leaving the user with a
'backend connection failed' modal and no way to know whether the
correction actually finished.

This module decouples the trigger from the work:

  POST /override/vulnrichment/start  → spawn thread, return job_id
  GET  /override/vulnrichment/status/{job_id} → progress + stats

State lives in a module-global dict so it survives the request but
not a backend restart. That's acceptable: a restart aborts any
in-flight job anyway, and the user just re-clicks.

Thread safety: dict assignment is atomic in CPython, and only the
worker thread mutates a given job's fields after creation. The
status reader returns whatever it sees — slightly stale reads are
fine for a 2-second poll loop.
"""
from __future__ import annotations

import logging
import threading
import uuid
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional

from sqlalchemy.orm import Session

from app.database import SessionLocal

logger = logging.getLogger(__name__)


# job_id (str) → state dict
_jobs: Dict[str, Dict[str, Any]] = {}
# Cap the in-memory ring so a misbehaving caller can't OOM us.
_MAX_JOBS = 50


def _prune_oldest() -> None:
    if len(_jobs) <= _MAX_JOBS:
        return
    # Drop the oldest finished job. Finished = has finished_at.
    finished = [
        (jid, j.get("finished_at"))
        for jid, j in _jobs.items()
        if j.get("finished_at")
    ]
    if not finished:
        return
    finished.sort(key=lambda t: t[1])
    _jobs.pop(finished[0][0], None)


def start_vulnrichment_job(
    user_id: int,
    cve_ids: Optional[List[str]] = None,
    asset_ids: Optional[List[int]] = None,
    dry_run: bool = False,
) -> str:
    """Spawn a worker thread, return its job_id immediately."""
    job_id = str(uuid.uuid4())
    _jobs[job_id] = {
        "job_id": job_id,
        "type": "vulnrichment",
        "status": "queued",       # queued | running | completed | failed
        "user_id": user_id,
        "dry_run": dry_run,
        "started_at": datetime.now().isoformat(),
        "finished_at": None,
        "stage": "starting",      # human-readable progress label
        "total": 0,
        "done": 0,
        "updated": 0,
        "checked": 0,
        "not_found": 0,
        "error": None,
        "result": None,
    }
    _prune_oldest()

    thread = threading.Thread(
        target=_run_vulnrichment_job,
        args=(job_id, cve_ids, asset_ids, dry_run),
        daemon=True,
        name=f"override-{job_id[:8]}",
    )
    thread.start()
    return job_id


def get_job(job_id: str) -> Optional[Dict[str, Any]]:
    """Return a snapshot of the job state, or None if unknown."""
    job = _jobs.get(job_id)
    if not job:
        return None
    # Return a copy so the caller can't mutate the live state.
    return dict(job)


def list_jobs(limit: int = 10) -> List[Dict[str, Any]]:
    """Most-recent-first list of jobs (for an admin overview)."""
    snapshot = list(_jobs.values())
    snapshot.sort(key=lambda j: j.get("started_at") or "", reverse=True)
    return [dict(j) for j in snapshot[:limit]]


def _run_vulnrichment_job(
    job_id: str,
    cve_ids: Optional[List[str]],
    asset_ids: Optional[List[int]],
    dry_run: bool,
) -> None:
    """Worker thread body. Owns its own DB session."""
    # Import lazily to avoid a circular import at module load.
    from app.services.vuln_override_service import (
        VulnOverrideService,
        correct_vulnerability_scores,
    )
    from app.models.vulnerability import Vulnerability, VulnerabilityStatus

    job = _jobs[job_id]
    job["status"] = "running"
    db: Session = SessionLocal()
    try:
        # Stage 1: figure out target CVE set so progress has a known total
        job["stage"] = "selecting target CVEs"
        query = db.query(Vulnerability).filter(
            Vulnerability.status == VulnerabilityStatus.open
        )
        if cve_ids:
            query = query.filter(Vulnerability.cve_id.in_(cve_ids))
        if asset_ids:
            query = query.filter(Vulnerability.asset_id.in_(asset_ids))
        vulns = query.all()
        unique = sorted({
            v.cve_id for v in vulns
            if v.cve_id and not v.cve_id.startswith("NESSUS-")
        })
        job["total"] = len(unique)

        if not unique:
            job["status"] = "completed"
            job["stage"] = "no CVEs to correct"
            job["finished_at"] = datetime.now().isoformat()
            job["result"] = {"message": "no real CVEs in scope", "updated": 0}
            return

        # Stage 2: fetch (ZIP snapshot for large batches, else per-CVE raw)
        if len(unique) > VulnOverrideService._ZIP_FALLBACK_THRESHOLD:
            job["stage"] = (
                f"downloading vulnrichment ZIP snapshot (~249 MB) for {len(unique)} CVEs"
            )
        else:
            job["stage"] = f"fetching {len(unique)} CVEs from github raw"

        # Stage 3: apply overrides via existing helper (does the heavy lift)
        # We don't get per-CVE callbacks from inside the helper, so we
        # surface "done" only after it returns. The two stages above
        # already tell the user what's currently happening.
        result = correct_vulnerability_scores(
            db,
            cve_ids=unique,
            asset_ids=asset_ids,
            dry_run=dry_run,
        )
        job["done"] = job["total"]
        job["checked"] = result.get("checked", 0)
        job["updated"] = result.get("updated", 0)
        job["not_found"] = result.get("not_found", 0)
        job["result"] = result
        job["status"] = "completed"
        job["stage"] = (
            f"done — {job['updated']} updated, "
            f"{job['not_found']} not in vulnrichment feed"
        )
    except Exception as e:
        logger.exception("override job %s failed", job_id)
        job["status"] = "failed"
        job["error"] = str(e)
        job["stage"] = f"failed: {e}"
    finally:
        job["finished_at"] = datetime.now().isoformat()
        db.close()
