#!/usr/bin/env python3
"""Measure Juliet precision impact of encoder bail-out categories.

This script runs fold with `--profile-reach` on Juliet testcase files one at a
time, parses the emitted reachability report, and correlates bailout categories
with file-level TP/miss outcomes. Because `--profile-reach` is aggregate per
scan rather than per-finding, the attribution is necessarily testcase-level:
every finding from a scan is treated as affected by the bailout categories
reported for that testcase file.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shlex
import subprocess
import sys
import time
from collections import defaultdict
from pathlib import Path

from run_fold_on_juliet import (  # type: ignore
    CLANG,
    CLANGXX,
    CWE_CONFIG,
    DEFAULT_FOLD,
    GOLD_ROOT,
    JULIET_ROOT,
    LLVM_BIN,
    RESULTS_ROOT,
    file_basename,
    load_findings,
    now_iso,
)


REACH_ROW_RE = re.compile(
    r"^\s*(?P<node>\w+):\s+encoded=(?P<encoded>\d+)\s+bailed=(?P<bailed>\d+)"
    r"\s+bail_rate=(?P<rate>[0-9.]+)%"
    r"(?:\s+top_bails=(?P<top>.*))?$"
)
TOP_REASON_RE = re.compile(r"(?P<reason>.+?)\((?P<count>\d+)\)")


def latest_full_run() -> Path | None:
    candidates = sorted(RESULTS_ROOT.glob("*_fold-juliet"))
    if not candidates:
        return None
    return candidates[-1]


def load_full_scores(full_run_dir: Path) -> dict[str, dict]:
    scores: dict[str, dict] = {}
    for cwe in CWE_CONFIG:
        payload = full_run_dir / cwe / "fold.json"
        if not payload.exists():
            continue
        doc = json.loads(payload.read_text())
        scores[cwe] = doc["score"]
    return scores


def load_gold_files(cwe: str) -> list[Path]:
    gold = json.loads((GOLD_ROOT / f"juliet-{cwe}.json").read_text())
    return [JULIET_ROOT / cwe / entry["files"][0] for entry in gold["expected"]]


def build_single_file_db(src: Path, corpus_dir: Path, db_path: Path) -> None:
    include_dirs = set()
    for path in sorted(corpus_dir.rglob("*")):
        if path.is_dir() and path.name == "testcasesupport":
            include_dirs.add(str(path))
        if path.is_dir() and any(path.glob("*.h")):
            include_dirs.add(str(path))
    inc_flags = [flag for d in sorted(include_dirs) for flag in ("-I", d)]
    compiler = CLANGXX if src.suffix == ".cpp" else CLANG
    std = "-std=c++17" if src.suffix == ".cpp" else "-std=c11"
    entry = {
        "directory": str(corpus_dir),
        "command": shlex.join([str(compiler), std, *inc_flags, "-c", str(src)]),
        "file": str(src),
    }
    db_path.write_text(json.dumps([entry], indent=2) + "\n")


def category_for_reason(node: str, reason: str) -> str:
    lower = reason.lower()
    if "function call" in lower or "constructor" in lower:
        return "Call"
    if "pointer dereference" in lower or "member expr" in lower:
        return "Pointer"
    if "unsupported local type" in lower:
        return "UnsupportedType"
    if "floatingliteral" in lower:
        return "Floating"
    if "cxxnewexpr" in lower or "cxxdeleteexpr" in lower:
        return "Heap"
    if "subscript value opaque" in lower:
        return "OverApprox"
    if lower.startswith("expr kind ") or lower.startswith("stmt kind "):
        return "UnsupportedNode"
    if "assign target" in lower or "assign non-var" in lower:
        return "UnsupportedNode"
    if "unary op" in lower or "non-var declref" in lower or "encode recursion" in lower:
        return "UnsupportedNode"
    if node in {"ArraySubscriptExpr", "MemberExpr", "UnaryOperator"}:
        return "OverApprox"
    return "Other"


def parse_reach_categories(log_text: str) -> dict[str, int]:
    categories: dict[str, int] = defaultdict(int)
    in_reach = False
    for line in log_text.splitlines():
        if line.startswith("[reach-profile] reachability"):
            in_reach = True
            continue
        if not in_reach:
            continue
        if line.startswith("[") and not line.startswith("[reach-profile]"):
            break
        content = line.replace("[reach-profile]", "", 1).rstrip()
        if not content.strip():
            continue
        match = REACH_ROW_RE.match(content)
        if not match:
            continue
        node = match.group("node")
        if int(match.group("bailed")) == 0:
            continue
        top = match.group("top") or ""
        if top:
            for reason_match in TOP_REASON_RE.finditer(top):
                reason = reason_match.group("reason").strip().rstrip(",")
                count = int(reason_match.group("count"))
                categories[category_for_reason(node, reason)] += count
        else:
            categories[category_for_reason(node, node)] += int(match.group("bailed"))
    return dict(categories)


def finding_files_for_scan(cwe_out: Path, kind: str) -> set[str]:
    findings, _ = load_findings(cwe_out)
    return {file_basename(item["file"]) for item in findings if item.get("kind") == kind}


def run_file(src: Path, cwe: str, fold_bin: Path, cache_dir: Path) -> dict:
    rel_key = f"{cwe}:{src.relative_to(JULIET_ROOT / cwe)}"
    digest = hashlib.sha256(rel_key.encode("utf-8")).hexdigest()[:16]
    cache_key = f"{cwe}-{src.stem}-{digest}.json"
    cache_path = cache_dir / cache_key
    if cache_path.exists():
        return json.loads(cache_path.read_text())

    corpus_dir = JULIET_ROOT / cwe
    work_dir = cache_dir / f"{cwe}-{src.stem}-{digest}"
    work_dir.mkdir(parents=True, exist_ok=True)
    db_path = work_dir / "compile_commands.json"
    sarif_path = work_dir / "fold.sarif"
    build_single_file_db(src, corpus_dir, db_path)

    config = CWE_CONFIG[cwe]
    command = [
        str(fold_bin), "scan",
        "--compile-db", str(db_path),
        "--format", "sarif",
        "--output", str(sarif_path),
        "--min-body-stmts", "1",
        "--processes", "1",
        "--profile-reach",
        "--detectors", config["detector"],
    ]
    for pattern in config.get("exclude_globs", []):
        command.extend(["--exclude-glob", pattern])

    env = os.environ.copy()
    env["PATH"] = f"{LLVM_BIN}:{env.get('PATH', '')}"
    start = time.monotonic()
    proc = subprocess.run(command, capture_output=True, text=True, env=env)
    elapsed = time.monotonic() - start
    combined_log = proc.stdout + "\n" + proc.stderr
    categories = parse_reach_categories(combined_log)
    found = file_basename(src.name) in finding_files_for_scan(work_dir, config["kind"])

    result = {
        "file": str(src),
        "cwe": cwe,
        "kind": config["kind"],
        "elapsed_s": elapsed,
        "exit_code": proc.returncode,
        "found": found,
        "categories": categories,
    }
    cache_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    return result


def summarize(records: list[dict], full_scores: dict[str, dict]) -> tuple[str, list[dict]]:
    total_findings = sum(1 for record in records if record["found"])
    grouped: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))

    for record in records:
        categories = set(record["categories"])
        for category in categories:
            grouped[category]["functions"] += 1
            if record["found"]:
                grouped[category]["findings"] += 1
                grouped[category]["tp"] += 1
            else:
                grouped[category]["missed_gold"] += 1
            grouped[category]["gold_total"] += 1

    baseline_by_cwe = {cwe: score.get("recall", 0.0) for cwe, score in full_scores.items()}
    overall_baseline = sum(baseline_by_cwe.values()) / len(baseline_by_cwe) if baseline_by_cwe else 0.0

    rows = []
    for category, stats in sorted(grouped.items(), key=lambda item: (-item[1]["functions"], item[0])):
        gold_total = stats["gold_total"]
        tp = stats["tp"]
        findings = stats["findings"]
        category_recall = tp / gold_total if gold_total else 0.0
        tp_rate = tp / findings if findings else 0.0
        recall_loss_pp = max(0.0, (overall_baseline - category_recall) * 100.0)
        rows.append({
            "category": category,
            "functions": stats["functions"],
            "findings_pct": (100.0 * findings / total_findings) if total_findings else 0.0,
            "tp_rate": tp_rate,
            "recall_loss_pp": recall_loss_pp,
        })

    lines = [
        "Bail-out Precision Impact on Juliet",
        "=====================================",
        "Category          | Functions | % Findings | TP Rate | Recall Loss",
        "------------------|-----------|------------|---------|------------",
    ]
    for row in rows:
        lines.append(
            f"{row['category']:<18}| {row['functions']:>9} | {row['findings_pct']:>9.1f}% | "
            f"{row['tp_rate']:>6.1%} | {row['recall_loss_pp']:>9.1f}pp"
        )
    return "\n".join(lines) + "\n", rows


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cwes", nargs="+", default=["476", "369", "190", "121"])
    parser.add_argument("--fold-bin", default=str(DEFAULT_FOLD))
    parser.add_argument("--full-run-dir", default="")
    parser.add_argument("--out-root", default=str(RESULTS_ROOT))
    parser.add_argument("--limit-files", type=int, default=0, help="Debug/smoke limit across the selected Juliet files")
    args = parser.parse_args()

    fold_bin = Path(args.fold_bin).resolve()
    if not fold_bin.exists():
        print(f"fold binary not found: {fold_bin}", file=sys.stderr)
        return 2
    if not CLANG.exists() or not CLANGXX.exists():
        print(f"clang/clang++ not found under {LLVM_BIN}", file=sys.stderr)
        return 2

    full_run_dir = Path(args.full_run_dir).resolve() if args.full_run_dir else latest_full_run()
    if full_run_dir is None or not full_run_dir.is_dir():
        print("no existing full Juliet run found; pass --full-run-dir", file=sys.stderr)
        return 2
    full_scores = load_full_scores(full_run_dir)

    out_root = Path(args.out_root).resolve() / f"{now_iso()}_fold-juliet-bailouts"
    out_root.mkdir(parents=True, exist_ok=True)
    cache_dir = out_root / "per-file"
    cache_dir.mkdir(parents=True, exist_ok=True)

    files: list[tuple[str, Path]] = []
    for cwe in args.cwes:
        for src in load_gold_files(cwe):
            files.append((cwe, src))
    if args.limit_files > 0:
        files = files[:args.limit_files]

    records = []
    for index, (cwe, src) in enumerate(files, start=1):
        print(f"[{index}/{len(files)}] profiling CWE-{cwe} {src.name}", file=sys.stderr)
        records.append(run_file(src, cwe, fold_bin, cache_dir))

    table, rows = summarize(records, full_scores)
    (out_root / "summary.txt").write_text(table)
    (out_root / "summary.json").write_text(json.dumps(rows, indent=2, sort_keys=True) + "\n")
    print(table)
    print(f"results written to {out_root}")
    return 0


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