"""Held-out CVE/malicious-package eval against extracted DataDog samples.

Usage:
    python eval/cve_eval.py --out eval_outputs/cve_held_out
"""
from __future__ import annotations

import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path
from typing import Iterable

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from models import (
    ActionType,
    IssueView,
    OpsguardAction,
    OpsguardObservation,
)
from eval.policies import build_policy


DEFAULT_POLICIES = [
    "random", "blanket_approve_pr", "blanket_reject_pr",
    "keyword_security_triager", "paranoid_inspector", "memory_aware",
]

CATCH_ACTIONS = {
    ActionType.REJECT_PR.value,
    ActionType.CLOSE_SPAM.value,
    ActionType.FLAG_SECURITY.value,
}


def load_samples(path: Path) -> list[dict]:
    if not path.exists():
        raise FileNotFoundError(f"missing dataset: {path}; run scripts/pull_datadog.py first")
    out = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                out.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return out


def sample_to_observation(sample: dict, idx: int) -> OpsguardObservation:
    pkg = sample.get("package", "unknown")
    ver = sample.get("version", "")
    eco = sample.get("ecosystem", "unknown")
    attack = sample.get("attack_type", "unknown_malicious")
    diff = sample.get("diff_preview", "")
    files = sample.get("files", [])

    title = f"chore({eco}): bump {pkg} to {ver}" if ver else f"chore({eco}): add {pkg}"
    body = (
        f"Adds dependency `{pkg}` ({eco}).\n\n"
        f"Attack-class hint: {attack}.\n\n"
        f"This PR vendors a known-malicious sample from the DataDog dataset.\n"
    )

    looks_like = pkg
    if pkg.endswith("s"):
        looks_like = pkg[:-1]
    elif "-" in pkg:
        looks_like = pkg.replace("-", "")

    dep = {
        "package": pkg,
        "version": ver,
        "ecosystem": eco,
        "looks_like": looks_like,
    }

    iv = IssueView(
        issue_id=10_000 + idx,
        number=10_000 + idx,
        title=title,
        body=body,
        is_pr=True,
        pr_diff_preview=diff[:6000],
        pr_changed_files=files[:50],
        pr_dependency_changes=[dep],
        author_login=f"contrib_{idx % 50}",
        author_pr_count=0,
        author_account_age_days=14,
        author_first_contribution_days_ago=0,
        available_labels=["bug", "enhancement", "question", "documentation", "security"],
        comments_preview=[],
    )

    return OpsguardObservation(
        scenario_id=f"cve_heldout::{attack}",
        step=0,
        step_budget=1,
        queue_position=idx,
        queue_total=0,
        current_issue=iv,
        memory_hits=[],
        recent_actions=[],
        repo_health=1.0,
        n_attacks_landed=0,
        feedback="",
    )


def evaluate(samples: list[dict], policies: dict, out_dir: Path) -> dict:
    out_dir.mkdir(parents=True, exist_ok=True)
    per_sample_path = out_dir / "results.jsonl"

    per_policy_stats: dict[str, dict] = {
        p: {"tp": 0, "fp": 0, "fn": 0, "tn": 0, "n": 0,
            "by_attack": defaultdict(lambda: {"caught": 0, "total": 0})}
        for p in policies
    }

    with open(per_sample_path, "w", encoding="utf-8") as fout:
        for i, sample in enumerate(samples):
            obs = sample_to_observation(sample, i)
            true_malicious = sample.get("label", "malicious") == "malicious"
            attack_type = sample.get("attack_type", "unknown_malicious")
            for pname, pfn in policies.items():
                action: OpsguardAction = pfn(obs)
                at = action.action_type.value if hasattr(action.action_type, "value") else str(action.action_type)
                caught = at in CATCH_ACTIONS
                stats = per_policy_stats[pname]
                stats["n"] += 1
                stats["by_attack"][attack_type]["total"] += 1
                if true_malicious and caught:
                    stats["tp"] += 1
                    stats["by_attack"][attack_type]["caught"] += 1
                elif true_malicious and not caught:
                    stats["fn"] += 1
                elif (not true_malicious) and caught:
                    stats["fp"] += 1
                else:
                    stats["tn"] += 1

                fout.write(json.dumps({
                    "sample_idx": i,
                    "package": sample.get("package"),
                    "version": sample.get("version"),
                    "ecosystem": sample.get("ecosystem"),
                    "attack_type": attack_type,
                    "true_label": sample.get("label", "malicious"),
                    "policy": pname,
                    "predicted_action": at,
                    "caught": caught,
                    "reasoning": action.reasoning,
                    "security_verdict": (action.security_verdict.value
                                         if action.security_verdict else None),
                }) + "\n")

    summary = {"per_policy": {}}
    for pname, s in per_policy_stats.items():
        tp, fp, fn = s["tp"], s["fp"], s["fn"]
        precision = tp / (tp + fp) if (tp + fp) else 0.0
        recall = tp / (tp + fn) if (tp + fn) else 0.0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
        catch_rate = tp / max(1, tp + fn)
        summary["per_policy"][pname] = {
            "n": s["n"], "tp": tp, "fp": fp, "fn": fn, "tn": s["tn"],
            "precision": round(precision, 4),
            "recall": round(recall, 4),
            "f1": round(f1, 4),
            "catch_rate": round(catch_rate, 4),
            "by_attack": {k: dict(v) for k, v in s["by_attack"].items()},
        }

    summary["n_samples"] = len(samples)
    summary["attack_types"] = sorted({s.get("attack_type", "unknown_malicious") for s in samples})

    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    (out_dir / "summary.md").write_text(format_markdown(summary), encoding="utf-8")
    return summary


def format_markdown(summary: dict) -> str:
    lines = [f"# CVE held-out eval ({summary['n_samples']} samples)\n"]
    lines.append("| Policy | precision | recall | F1 | catch_rate | TP | FP | FN | TN |")
    lines.append("|---|---:|---:|---:|---:|---:|---:|---:|---:|")
    for pname, s in sorted(summary["per_policy"].items()):
        lines.append(
            f"| {pname} | {s['precision']:.3f} | {s['recall']:.3f} | {s['f1']:.3f} | "
            f"{s['catch_rate']:.3f} | {s['tp']} | {s['fp']} | {s['fn']} | {s['tn']} |"
        )
    lines.append("\n## Per-attack catch rate\n")
    attacks = summary["attack_types"]
    header = "| Policy | " + " | ".join(attacks) + " |"
    sep = "|---|" + "|".join(["---:"] * len(attacks)) + "|"
    lines.append(header)
    lines.append(sep)
    for pname, s in sorted(summary["per_policy"].items()):
        row = [pname]
        for at in attacks:
            cell = s["by_attack"].get(at, {"caught": 0, "total": 0})
            tot = cell.get("total", 0)
            cau = cell.get("caught", 0)
            rate = cau / tot if tot else 0.0
            row.append(f"{cau}/{tot} ({rate:.2f})")
        lines.append("| " + " | ".join(row) + " |")
    return "\n".join(lines) + "\n"


def write_plot(summary: dict, out_path: Path) -> bool:
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except Exception as e:
        print(f"[plot] matplotlib unavailable ({e}); skipping", flush=True)
        return False
    pol = sorted(summary["per_policy"].keys())
    rates = [summary["per_policy"][p]["catch_rate"] for p in pol]
    fig, ax = plt.subplots(figsize=(max(6, len(pol) * 1.2), 4))
    ax.bar(pol, rates, color="#cc3333")
    ax.set_ylim(0, 1.05)
    ax.set_ylabel("catch rate (recall)")
    ax.set_title(f"CVE held-out catch rate (n={summary['n_samples']})")
    for i, v in enumerate(rates):
        ax.text(i, v + 0.02, f"{v:.2f}", ha="center", fontsize=9)
    plt.xticks(rotation=25, ha="right")
    plt.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=120)
    plt.close(fig)
    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="data/datadog_extracted.jsonl")
    ap.add_argument("--out", default="eval_outputs/cve_held_out")
    ap.add_argument("--policies", nargs="*", default=DEFAULT_POLICIES)
    ap.add_argument("--limit", type=int, default=0,
                    help="cap number of samples (0 = all)")
    args = ap.parse_args()

    project_root = Path(__file__).resolve().parent.parent
    data_path = (project_root / args.data).resolve()
    out_dir = (project_root / args.out).resolve()

    samples = load_samples(data_path)
    if args.limit > 0:
        samples = samples[: args.limit]
    print(f"[eval] loaded {len(samples)} samples from {data_path}", flush=True)

    policies = {n: build_policy(n) for n in args.policies}
    print(f"[eval] policies: {list(policies)}", flush=True)

    summary = evaluate(samples, policies, out_dir)
    write_plot(summary, out_dir / "catch_rate.png")

    print(f"[done] {len(samples)} samples × {len(policies)} policies → {out_dir}", flush=True)
    for p, s in sorted(summary["per_policy"].items()):
        print(f"  {p:>26}  P={s['precision']:.3f}  R={s['recall']:.3f}  "
              f"F1={s['f1']:.3f}  catch={s['catch_rate']:.3f}", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())