#!/usr/bin/env python3
"""Score an expanded Juliet run with per-detector precision/recall aggregates.

Reads the fold.json payloads produced by run_expanded.py and computes:
  1. Per-CWE precision/recall/F1 (file-level scoring)
  2. Per-detector aggregate precision/recall/F1 across all CWEs
  3. Comparison against the baseline (original 14 CWE coverage)

This is a standalone scorer that operates on already-produced results.

Usage:
  # Score a completed expanded run
  python3 bench/juliet/expansion/score_expanded.py \
    --run-dir bench/results/2026-05-27T12-00-00Z_fold-juliet-expanded

  # Compare against a baseline run (original 14 CWEs only)
  python3 bench/juliet/expansion/score_expanded.py \
    --run-dir bench/results/<expanded> \
    --baseline-dir bench/results/<baseline>

  # Output JSON for programmatic consumption
  python3 bench/juliet/expansion/score_expanded.py \
    --run-dir bench/results/<expanded> --json
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
HERE = Path(__file__).resolve().parent
CATALOG_PATH = HERE / "cwe_catalog.json"
GOLD_ROOT = REPO_ROOT / "bench" / "gold"
RESULTS_ROOT = REPO_ROOT / "bench" / "results"


def file_basename(path: str) -> str:
    return Path(path).name


def load_catalog() -> dict:
    return json.loads(CATALOG_PATH.read_text())


CATALOG = load_catalog()


def find_cwe_dirs(run_dir: Path) -> list[Path]:
    """Find all CWE-* subdirectories in a run directory."""
    return sorted(
        p for p in run_dir.iterdir()
        if p.is_dir() and p.name.isdigit() and len(p.name) == 3
    )


def score_per_cwe(fold_json: Path) -> dict | None:
    """Score a single CWE result against its gold standard."""
    if not fold_json.exists():
        return None

    payload = json.loads(fold_json.read_text())
    cwe = payload["cwe"]
    kinds = payload.get("kinds", [])
    findings = payload.get("findings", [])

    gold_path = GOLD_ROOT / f"juliet-{cwe}.json"
    if not gold_path.exists():
        return {
            "cwe": cwe,
            "error": f"gold file not found: {gold_path}",
        }

    gold = json.loads(gold_path.read_text())
    gold_entries = gold.get("expected", [])

    # File-level scoring
    kind_set = set(kinds)
    gold_files = {
        file_basename(entry["files"][0])
        for entry in gold_entries
        if entry.get("kind") in kind_set
    }
    finding_files = {
        file_basename(f["file"])
        for f in findings
        if f.get("kind") in kind_set
    }

    tp_files = gold_files & finding_files
    tp = len(tp_files)
    fp = len(finding_files - gold_files)
    fn = len(gold_files - finding_files)

    precision = tp / len(finding_files) if finding_files else 0.0
    recall = tp / len(gold_files) if gold_files else 0.0
    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0

    return {
        "cwe": cwe,
        "kinds": kinds,
        "detectors": payload.get("detectors", []),
        "elapsed_s": payload.get("elapsed_s", 0),
        "entries": payload.get("entries", 0),
        "true_positive": tp,
        "false_positive": fp,
        "false_negative": fn,
        "gold_total": len(gold_files),
        "finding_groups": len(finding_files),
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "matched_files": sorted(tp_files),
        "missed_files": sorted(gold_files - finding_files),
    }


def aggregate_by_detector(cwe_scores: list[dict]) -> dict[str, dict]:
    """Aggregate per-CWE scores into per-detector statistics."""
    detector_stats: dict[str, dict] = {}

    for score in cwe_scores:
        if "error" in score:
            continue
        for detector in score.get("detectors", []):
            if detector not in detector_stats:
                detector_stats[detector] = {
                    "tp": 0, "fp": 0, "fn": 0,
                    "gold_total": 0, "finding_groups": 0,
                    "cwe_count": 0, "elapsed_s": 0.0,
                    "cwe_list": [],
                }
            stats = detector_stats[detector]
            stats["tp"] += score["true_positive"]
            stats["fp"] += score["false_positive"]
            stats["fn"] += score["false_negative"]
            stats["gold_total"] += score["gold_total"]
            stats["finding_groups"] += score["finding_groups"]
            stats["cwe_count"] += 1
            stats["elapsed_s"] += score["elapsed_s"]
            stats["cwe_list"].append(score["cwe"])

    output = {}
    for detector, stats in sorted(detector_stats.items()):
        tp, fp, fn = stats["tp"], stats["fp"], stats["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
        output[detector] = {
            **stats,
            "precision": precision,
            "recall": recall,
            "f1": f1,
        }
    return output


def render_markdown(
    cwe_scores: list[dict],
    per_detector: dict[str, dict],
    baseline_comparison: dict | None = None,
) -> str:
    """Render results as Markdown."""
    lines = [
        "# fold Juliet Expanded — Scoring Report",
        "",
        f"CWEs evaluated: {len([s for s in cwe_scores if 'error' not in s])}",
        f"Errors: {len([s for s in cwe_scores if 'error' in s])}",
        "",
        "## Per-CWE Results",
        "",
        "| CWE | Detector(s) | TP | FP | FN | Gold | Precision | Recall | F1 | Time |",
        "|-----|-------------|----|----|----|------|-----------|--------|----|------|",
    ]

    for score in cwe_scores:
        if "error" in score:
            lines.append(f"| CWE-{score['cwe']} | ⚠ {score['error']} |")
            continue
        lines.append(
            f"| CWE-{score['cwe']} | {','.join(score['detectors'])} | "
            f"{score['true_positive']} | {score['false_positive']} | "
            f"{score['false_negative']} | {score['gold_total']} | "
            f"{score['precision']:.1%} | {score['recall']:.1%} | {score['f1']:.1%} | "
            f"{score['elapsed_s']:.1f}s |"
        )

    lines.extend([
        "",
        "## Per-Detector Aggregate",
        "",
        "| Detector | CWEs | TP | FP | FN | Gold | Precision | Recall | F1 | Total Time |",
        "|----------|------|----|----|----|------|-----------|--------|----|------------|",
    ])

    for detector, stats in sorted(per_detector.items()):
        lines.append(
            f"| {detector} | {stats['cwe_count']} | {stats['tp']} | {stats['fp']} | "
            f"{stats['fn']} | {stats['gold_total']} | "
            f"{stats['precision']:.1%} | {stats['recall']:.1%} | {stats['f1']:.1%} | "
            f"{stats['elapsed_s']:.1f}s |"
        )

    if baseline_comparison:
        lines.extend([
            "",
            "## Baseline Comparison",
            "",
            "Comparing expanded coverage vs. original 14-CWE baseline.",
            "",
            "| Detector | Baseline F1 | Expanded F1 | Delta |",
            "|----------|-------------|-------------|-------|",
        ])
        for detector, stats in sorted(per_detector.items()):
            base = baseline_comparison.get(detector, {})
            base_f1 = base.get("f1", 0.0)
            exp_f1 = stats["f1"]
            delta = exp_f1 - base_f1
            sign = "+" if delta >= 0 else ""
            lines.append(
                f"| {detector} | {base_f1:.1%} | {exp_f1:.1%} | {sign}{delta:.1%} |"
            )

    lines.append("")
    return "\n".join(lines)


def load_baseline(baseline_dir: Path) -> dict[str, dict]:
    """Load per-detector scores from a baseline run directory."""
    cwe_scores = []
    for cwe_dir in find_cwe_dirs(baseline_dir):
        score = score_per_cwe(cwe_dir / "fold.json")
        if score:
            cwe_scores.append(score)
    return aggregate_by_detector(cwe_scores)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--run-dir",
        required=True,
        help="Path to expanded run directory (e.g., bench/results/<stamp>_fold-juliet-expanded)",
    )
    parser.add_argument(
        "--baseline-dir",
        help="Path to baseline run directory for comparison",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output JSON instead of Markdown",
    )
    parser.add_argument(
        "--out",
        help="Write output to file instead of stdout",
    )
    args = parser.parse_args()

    run_dir = Path(args.run_dir).resolve()
    if not run_dir.is_dir():
        print(f"run-dir not found: {run_dir}", file=sys.stderr)
        return 2

    cwe_dirs = find_cwe_dirs(run_dir)
    if not cwe_dirs:
        print(f"No CWE subdirectories found in {run_dir}", file=sys.stderr)
        return 1

    cwe_scores = []
    for cwe_dir in cwe_dirs:
        fold_json = cwe_dir / "fold.json"
        score = score_per_cwe(fold_json)
        if score:
            cwe_scores.append(score)

    per_detector = aggregate_by_detector(cwe_scores)

    baseline = None
    if args.baseline_dir:
        baseline_dir = Path(args.baseline_dir).resolve()
        if baseline_dir.is_dir():
            baseline = load_baseline(baseline_dir)
        else:
            print(f"baseline-dir not found: {baseline_dir}", file=sys.stderr)

    if args.json:
        output = {
            "run_dir": str(run_dir),
            "cwe_scores": cwe_scores,
            "per_detector": per_detector,
        }
        if baseline:
            output["baseline"] = baseline
        output_text = json.dumps(output, indent=2) + "\n"
    else:
        output_text = render_markdown(cwe_scores, per_detector, baseline)

    if args.out:
        Path(args.out).write_text(output_text + "\n")
        print(f"Output written to {args.out}")
    else:
        print(output_text)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())