#!/usr/bin/env python3
"""Score Fold's binary lane against a CWE-grounded gold corpus.

Mirrors the source lane's bench/score_fp.py, but for the binary backend:
for each ELF in a corpus dir, run `fold scan --mode binary`, collect findings,
and score per-detector TP/FP/FN/F1 against a gold/disposition JSON.

Usage:
    python3 bench/score_binary.py --corpus tests/fixtures/binary/cwe \
        --gold bench/gold/binary-security.json \
        --fold-bin build/release/src/cli/fold
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_GOLD = ROOT / "bench" / "gold" / "binary-security.json"
DEFAULT_CORPUS = ROOT / "tests" / "fixtures" / "binary" / "cwe"
PCODE_SCAN = ROOT / "build" / "release" / "pcode_scan"

# FindingKind enum integer -> gold kind name (mirror include/fold/report/finding.hpp).
KIND_TO_NAME = {
    2: "dead",
    5: "null-deref",
    6: "int-overflow",
    7: "int-truncation",
    8: "sign-conversion",
    9: "div-by-zero",
    10: "oob-access",
    11: "taint-track",
    12: "unused-function",
    15: "uninit-local-read",
    16: "use-after-free",
    17: "use-after-move",
    24: "uninit-read",
    27: "double-free",
    29: "memory-leak",
}


def load_gold(path: Path) -> dict[str, set[str]]:
    """Return {finding_kind: {binary_basename}} from the gold JSON."""
    try:
        doc = json.loads(path.read_text())
    except (OSError, json.JSONDecodeError) as exc:
        raise SystemExit(f"gold file {path} is missing or invalid: {exc}") from exc
    out: dict[str, set[str]] = {}
    for entry in doc.get("expected", []):
        kind = entry["kind"]
        out.setdefault(kind, set()).update(entry["files"])
    return out


def run_binary_scan(binary: Path) -> list[dict]:
    """Run the dedicated P-Code scan tool and return normalized findings.

    `fold scan --mode binary` is heuristic metadata only; the P-Code security
    detectors (div-by-zero, oob-access, null-deref, ...) are exposed through the
    `pcode_scan` smoke driver, which prints `kind=<int> fn=<name>` lines.
    """
    if not PCODE_SCAN.exists():
        return []
    import re

    try:
        proc = subprocess.run(
            [str(PCODE_SCAN), str(binary)],
            capture_output=True,
            text=True,
            check=False,
            timeout=90,
        )
    except subprocess.TimeoutExpired:
        print(f"  (timeout) {binary.name}", file=sys.stderr)
        return []
    findings = []
    for line in (proc.stdout or "").splitlines():
        m = re.search(r"kind=(\d+)", line)
        if not m:
            continue
        try:
            kind_int = int(m.group(1))
        except ValueError:
            continue
        kind = KIND_TO_NAME.get(kind_int)
        if kind:
            findings.append({"kind": kind, "basename": binary.name})
    return findings


def score(
    gold: dict[str, set[str]], findings: list[dict], corpus_names: set[str]
) -> dict:
    """Score per gold kind: TP/FP/FN/F1 over the corpus binaries."""
    results = {}
    for kind, gold_files in sorted(gold.items()):
        found = {f["basename"] for f in findings if f["kind"] == kind}
        gold_in_corpus = gold_files & corpus_names
        tp = found & gold_in_corpus
        fp = found - gold_in_corpus
        fn = gold_in_corpus - found
        p = len(tp) / (len(tp) + len(fp)) if (len(tp) + len(fp)) else 0.0
        r = len(tp) / (len(tp) + len(fn)) if (len(tp) + len(fn)) else 0.0
        f1 = 2 * p * r / (p + r) if (p + r) else 0.0
        results[kind] = {
            "tp": len(tp),
            "fp": len(fp),
            "fn": len(fn),
            "precision": p,
            "recall": r,
            "f1": f1,
        }
    return results


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--corpus", type=Path, default=DEFAULT_CORPUS)
    ap.add_argument("--gold", type=Path, default=DEFAULT_GOLD)
    ap.add_argument(
        "--fold-bin", type=Path, default=ROOT / "build/release/src/cli/fold"
    )
    args = ap.parse_args()

    gold = load_gold(args.gold)
    corpus_names = {p.name for p in args.corpus.glob("*.elf")}
    all_findings: list[dict] = []
    for binary in sorted(args.corpus.glob("*.elf")):
        all_findings.extend(run_binary_scan(binary))

    results = score(gold, all_findings, corpus_names)

    print(f"{'kind':<16} {'TP':>3} {'FP':>3} {'FN':>3} {'P':>6} {'R':>6} {'F1':>6}")
    for kind, r in results.items():
        print(
            f"{kind:<16} {r['tp']:>3} {r['fp']:>3} {r['fn']:>3} "
            f"{r['precision']:>6.3f} {r['recall']:>6.3f} {r['f1']:>6.3f}"
        )
    return 0


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