#!/usr/bin/env python3
"""Measure fold performance on a filtered Juliet completeness class.

Completeness class:
- loop-free
- integer-arithmetic-only
- no heap
- no extern / user function calls
- no pointer casts

The filter is intentionally source-level and conservative: it inspects the
Juliet testcase's `_bad` function body and rejects any testcase that contains a
disqualifying construct. The measurement still scores at Juliet's existing
file-level granularity so results remain directly comparable with the existing
`bench/run_fold_on_juliet.py` output.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shlex
import subprocess
import sys
import time
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,
    load_gold,
    now_iso,
    parse_profiles,
    score_file_level,
)


COMMENT_RE = re.compile(r"//.*?$|/\*.*?\*/", re.DOTALL | re.MULTILINE)
STRING_RE = re.compile(r'"(?:\\.|[^"\\])*"|\'(?:\\.|[^\'\\])*\'')
BAD_FUNC_RE = re.compile(
    r"\b[A-Za-z_]\w*?_bad\s*\([^;{}]*\)\s*\{",
    re.MULTILINE,
)
CALL_RE = re.compile(r"\b([A-Za-z_]\w*)\s*\(")
FLOAT_LITERAL_RE = re.compile(
    r"(?<![\w.])(?:\d+\.\d*|\.\d+)(?:[eE][+-]?\d+)?[fFlL]?|"
    r"(?<![\w.])\d+[eE][+-]?\d+[fFlL]?"
)
POINTER_CAST_RE = re.compile(
    r"\(\s*(?:const\s+|volatile\s+|struct\s+|enum\s+|unsigned\s+|signed\s+|long\s+|short\s+)*"
    r"[A-Za-z_]\w*(?:\s+[A-Za-z_]\w*)*\s*\*+[^)]*\)"
)

CONTROL_WORDS = {
    "if",
    "for",
    "while",
    "switch",
    "return",
    "sizeof",
}
TYPE_WORDS = {
    "bool",
    "char",
    "int",
    "long",
    "short",
    "signed",
    "unsigned",
    "size_t",
    "ssize_t",
    "uintptr_t",
    "intptr_t",
    "void",
    "wchar_t",
}
HEAP_TOKENS = {"malloc", "calloc", "realloc", "free", "new", "delete"}


def strip_comments_and_strings(text: str) -> str:
    text = COMMENT_RE.sub("", text)
    return STRING_RE.sub("\"\"", text)


def extract_braced_block(text: str, open_brace: int) -> str | None:
    depth = 0
    for idx in range(open_brace, len(text)):
        ch = text[idx]
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return text[open_brace + 1:idx]
    return None


def extract_bad_function_body(source_text: str) -> str | None:
    sanitized = strip_comments_and_strings(source_text)
    match = BAD_FUNC_RE.search(sanitized)
    if not match:
        return None
    brace_index = sanitized.find("{", match.start())
    if brace_index < 0:
        return None
    return extract_braced_block(sanitized, brace_index)


def classify_body(body: str) -> tuple[bool, list[str]]:
    reasons: list[str] = []

    if re.search(r"\b(?:for|while|do)\b", body):
        reasons.append("loop")
    if re.search(r"\b(?:float|double|long\s+double)\b", body):
        reasons.append("float-type")
    if FLOAT_LITERAL_RE.search(body):
        reasons.append("float-literal")
    if POINTER_CAST_RE.search(body):
        reasons.append("pointer-cast")
    if any(re.search(rf"\b{token}\b", body) for token in HEAP_TOKENS):
        reasons.append("heap")

    calls: set[str] = set()
    for match in CALL_RE.finditer(body):
        callee = match.group(1)
        if callee in CONTROL_WORDS or callee in TYPE_WORDS:
            continue
        calls.add(callee)
    if calls:
        reasons.append("calls:" + ",".join(sorted(calls)))

    return not reasons, reasons


def build_compile_db(files: list[Path], corpus_dir: Path, db_path: Path) -> int:
    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)]

    entries = []
    for src in sorted(files):
        compiler = CLANGXX if src.suffix == ".cpp" else CLANG
        std = "-std=c++17" if src.suffix == ".cpp" else "-std=c11"
        entries.append({
            "directory": str(corpus_dir),
            "command": shlex.join([str(compiler), std, *inc_flags, "-c", str(src)]),
            "file": str(src),
        })
    db_path.write_text(json.dumps(entries, indent=2) + "\n")
    return len(entries)


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 select_gold_entries(cwe: str) -> tuple[list[dict], dict[str, list[str]]]:
    gold = load_gold(cwe)
    selected: list[dict] = []
    skipped: dict[str, list[str]] = {}
    corpus_dir = JULIET_ROOT / cwe

    for entry in gold["expected"]:
        rel = Path(entry["files"][0])
        src = corpus_dir / rel
        if not src.exists():
            skipped[str(rel)] = ["missing-source"]
            continue
        body = extract_bad_function_body(src.read_text())
        if body is None:
            skipped[str(rel)] = ["missing-bad-body"]
            continue
        keep, reasons = classify_body(body)
        if keep:
            selected.append(entry)
        else:
            skipped[str(rel)] = reasons
    return selected, skipped


def run_subset(cwe: str, selected_entries: list[dict], fold_bin: Path, out_root: Path) -> dict:
    config = CWE_CONFIG[cwe]
    detector = config["detector"]
    kind = config["kind"]
    exclude_globs = config.get("exclude_globs", [])
    corpus_dir = JULIET_ROOT / cwe
    cwe_out = out_root / cwe
    cwe_out.mkdir(parents=True, exist_ok=True)

    selected_files = [corpus_dir / entry["files"][0] for entry in selected_entries]
    db_path = cwe_out / "compile_commands.json"
    file_count = build_compile_db(selected_files, corpus_dir, db_path)
    sarif_path = cwe_out / "fold.sarif"
    stdout_path = cwe_out / "fold.stdout.log"
    stderr_path = cwe_out / "fold.stderr.log"

    command = [
        str(fold_bin), "scan",
        "--compile-db", str(db_path),
        "--format", "sarif",
        "--output", str(sarif_path),
        "--min-body-stmts", "1",
        "--processes", "1",
        "--profile-solver",
        "--profile-reach",
        "--detectors", detector,
    ]
    for pattern in 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
    stdout_path.write_text(proc.stdout)
    stderr_path.write_text(proc.stderr)

    combined_log = proc.stdout + "\n" + proc.stderr
    findings, findings_source = load_findings(cwe_out)
    score = score_file_level(findings, selected_entries, kind)
    profile = parse_profiles(combined_log, kind)
    payload = {
        "tool": "fold",
        "cwe": cwe,
        "kind": kind,
        "detector": detector,
        "elapsed_s": elapsed,
        "exit_code": proc.returncode,
        "entries": file_count,
        "findings": findings,
        "findings_source": findings_source,
        "score": score,
        "profile": profile,
        "selected_gold": len(selected_entries),
        "artifacts": {
            "compile_db": str(db_path),
            "sarif": str(sarif_path),
            "stdout_log": str(stdout_path),
            "stderr_log": str(stderr_path),
        },
    }
    (cwe_out / "fold.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    return payload


def render_summary(cwes: list[str], results: list[dict], full_scores: dict[str, dict], selection_manifest: dict[str, dict]) -> str:
    by_cwe = {result["cwe"]: result for result in results}
    lines = [
        "Completeness Class Results (loop-free, no-heap, no-calls, int-only)",
        "===================================================================",
    ]
    for cwe in cwes:
        result = by_cwe.get(cwe)
        if result is None:
            full = full_scores.get(cwe, {})
            lines.append(
                f"CWE-{cwe} ({CWE_CONFIG[cwe]['kind']}): no completeness-class Juliet cases matched "
                f"the current source-level filter (full-set recall {full.get('recall', 0.0):.1%})"
            )
            continue
        score = result["score"]
        full = full_scores.get(result["cwe"], {})
        full_recall = full.get("recall", 0.0)
        lines.append(
            f"CWE-{result['cwe']} ({result['kind']}): "
            f"{score['true_positive']}/{score['gold_total']} TP "
            f"({score['recall']:.1%} recall, {score['precision']:.1%} precision, {score['f1']:.1%} F1) "
            f"on {result['selected_gold']} completeness-class cases "
            f"(vs {full_recall:.1%} recall on full set)"
        )
    return "\n".join(lines) + "\n"


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="",
        help="Existing bench/results/*_fold-juliet directory used for full-set comparison",
    )
    parser.add_argument("--out-root", default=str(RESULTS_ROOT))
    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-completeness"
    out_root.mkdir(parents=True, exist_ok=True)

    results = []
    selection_manifest: dict[str, dict] = {}
    for cwe in args.cwes:
        if cwe not in CWE_CONFIG:
            print(f"unsupported CWE: {cwe}", file=sys.stderr)
            return 2
        selected_entries, skipped = select_gold_entries(cwe)
        selection_manifest[cwe] = {
            "selected": len(selected_entries),
            "skipped": len(skipped),
            "skipped_reasons": skipped,
        }
        if not selected_entries:
            continue
        print(f"running completeness subset for CWE-{cwe} ({len(selected_entries)} files)...", file=sys.stderr)
        result = run_subset(cwe, selected_entries, fold_bin, out_root)
        results.append(result)

    (out_root / "selection.json").write_text(json.dumps(selection_manifest, indent=2, sort_keys=True) + "\n")
    summary = render_summary(args.cwes, results, full_scores, selection_manifest)
    (out_root / "summary.txt").write_text(summary)
    print(summary)
    print(f"results written to {out_root}")
    return 0


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