#!/usr/bin/env python3
from __future__ import annotations

import argparse
import concurrent.futures
import datetime as dt
import json
import os
import plistlib
import shlex
import shutil
import subprocess
import sys
import tempfile
import time
from collections import Counter
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parent.parent
JULIET_ROOT = REPO_ROOT / "bench" / "corpora" / "juliet" / "filtered"
GOLD_ROOT = REPO_ROOT / "bench" / "gold"
RESULTS_ROOT = REPO_ROOT / "bench" / "results"
LLVM_BIN = Path(os.environ.get("LLVM_BIN", "/opt/homebrew/opt/llvm@22/bin"))
SCAN_BUILD = LLVM_BIN / "scan-build"
CLANG = LLVM_BIN / "clang"
CLANGXX = LLVM_BIN / "clang++"


CWE_CONFIG = {
    "476": {
        "kind": "null-deref",
        "checkers": ["core.NullDereference", "core.CallAndMessage"],
        "notes": "Direct CSA support via null-dereference checkers.",
    },
    "369": {
        "kind": "div-by-zero",
        "checkers": ["core.DivideZero"],
        "notes": "Direct CSA support, but mostly for integer divide/modulo-by-zero.",
    },
    "190": {
        "kind": "int-overflow",
        "checkers": [],
        "notes": "Clang SA 22 does not ship a Juliet-relevant integer-overflow checker.",
    },
    "121": {
        "kind": "oob-access",
        "checkers": ["security.ArrayBound"],
        "notes": "Uses the current checker name in LLVM 22; older docs often refer to alpha ArrayBound variants.",
    },
}


def now_iso() -> str:
    return dt.datetime.now(dt.UTC).strftime("%Y-%m-%dT%H-%M-%SZ")


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


def load_gold(cwe: str) -> dict:
    gold_path = GOLD_ROOT / f"juliet-{cwe}.json"
    return json.loads(gold_path.read_text())


def score_file_level(findings: list[dict], gold_entries: list[dict], kind: str) -> dict:
    gold_files = {file_basename(entry["files"][0]) for entry in gold_entries if entry.get("kind") == kind}
    finding_files = {file_basename(f["file"]) for f in findings if f.get("kind") == kind}
    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 {
        "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 write_compile_db(corpus_dir: Path, db_path: Path) -> list[dict]:
    entries: list[dict] = []
    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)]
    for src in sorted(corpus_dir.rglob("*.c")):
        entries.append({
            "directory": str(corpus_dir),
            "command": [str(CLANG), "-std=c11", *inc_flags, "-c", str(src)],
            "file": str(src),
        })
    for src in sorted(corpus_dir.rglob("*.cpp")):
        entries.append({
            "directory": str(corpus_dir),
            "command": [str(CLANGXX), "-std=c++17", *inc_flags, "-c", str(src)],
            "file": str(src),
        })
    db_path.write_text(json.dumps([
        {"directory": e["directory"], "command": shlex.join(e["command"]), "file": e["file"]}
        for e in entries
    ], indent=2) + "\n")
    return entries


def parse_plist_reports(report_root: Path, enabled_checkers: set[str], kind: str) -> tuple[list[dict], Counter]:
    findings: list[dict] = []
    checker_counts: Counter = Counter()
    for plist_path in sorted(report_root.rglob("*.plist")):
        try:
            with plist_path.open("rb") as fh:
                doc = plistlib.load(fh)
        except Exception:
            continue
        files = doc.get("files", [])
        for diag in doc.get("diagnostics", []):
            checker = diag.get("check_name", "")
            if checker not in enabled_checkers:
                continue
            location = diag.get("location", {})
            file_index = location.get("file", 0)
            if not isinstance(file_index, int) or file_index >= len(files):
                continue
            findings.append({
                "file": files[file_index],
                "line": int(location.get("line", 0) or 0),
                "kind": kind,
                "checker": checker,
                "message": diag.get("description", ""),
            })
            checker_counts[checker] += 1
    return findings, checker_counts


def run_scan_build_entry(entry: dict, enabled_checkers: list[str], kind: str, reports_root: Path) -> dict:
    file_path = Path(entry["file"])
    temp_root = Path(tempfile.mkdtemp(prefix="scan-build-", dir=reports_root))
    command = [
        str(SCAN_BUILD),
        "-plist",
        "--keep-empty",
        "-disable-checker",
        "deadcode.DeadStores",
        "-disable-checker",
        "core.uninitialized.Assign",
        "-disable-checker",
        "core.uninitialized.Branch",
        "-disable-checker",
        "core.uninitialized.UndefReturn",
    ]
    for checker in enabled_checkers:
        command.extend(["-enable-checker", checker])
    command.extend([
        "-o",
        str(temp_root),
        "--use-analyzer",
        str(CLANG),
        *entry["command"],
    ])
    proc = subprocess.run(command, cwd=entry["directory"], capture_output=True, text=True)
    findings, checker_counts = parse_plist_reports(temp_root, set(enabled_checkers), kind)
    return {
        "file": str(file_path),
        "exit_code": proc.returncode,
        "stdout": proc.stdout,
        "stderr": proc.stderr,
        "findings": findings,
        "checker_counts": checker_counts,
    }


def normalize_findings(findings: list[dict]) -> list[dict]:
    dedup: dict[tuple[str, str], dict] = {}
    for finding in findings:
        key = (str(Path(finding["file"]).resolve()), finding["kind"])
        current = dedup.get(key)
        if current is None or finding.get("line", 0) < current.get("line", 0):
            dedup[key] = {
                "file": finding["file"],
                "line": finding.get("line", 0),
                "kind": finding["kind"],
                "id": f"{finding['kind']}:{file_basename(finding['file'])}",
                "checker": finding.get("checker", ""),
                "message": finding.get("message", ""),
            }
    return sorted(dedup.values(), key=lambda item: (item["file"], item["line"]))


def run_cwe(cwe: str, jobs: int, out_root: Path) -> dict:
    config = CWE_CONFIG[cwe]
    gold = load_gold(cwe)
    kind = config["kind"]
    corpus_dir = JULIET_ROOT / cwe
    if not corpus_dir.is_dir():
        raise FileNotFoundError(f"missing Juliet filtered corpus: {corpus_dir}")

    cwe_out = out_root / cwe
    cwe_out.mkdir(parents=True, exist_ok=True)
    db_path = cwe_out / "compile_commands.json"
    entries = write_compile_db(corpus_dir, db_path)

    if not config["checkers"]:
        normalized_findings: list[dict] = []
        checker_counts: Counter = Counter()
        elapsed = 0.0
        score = score_file_level(normalized_findings, gold["expected"], kind)
        payload = {
            "tool": "clang-sa",
            "cwe": cwe,
            "kind": kind,
            "version": subprocess.run([str(CLANG), "--version"], capture_output=True, text=True).stdout.splitlines()[0],
            "checkers": config["checkers"],
            "elapsed_s": elapsed,
            "entries": len(entries),
            "findings": normalized_findings,
            "checker_counts": dict(checker_counts),
            "score": score,
            "notes": config["notes"],
        }
        (cwe_out / "clang-sa.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
        return payload

    reports_root = cwe_out / "reports"
    reports_root.mkdir(parents=True, exist_ok=True)
    raw_findings: list[dict] = []
    checker_counts: Counter = Counter()
    start = time.monotonic()
    with concurrent.futures.ThreadPoolExecutor(max_workers=jobs) as pool:
        futures = [pool.submit(run_scan_build_entry, entry, config["checkers"], kind, reports_root) for entry in entries]
        for future in concurrent.futures.as_completed(futures):
            result = future.result()
            raw_findings.extend(result["findings"])
            checker_counts.update(result["checker_counts"])
    elapsed = time.monotonic() - start

    normalized_findings = normalize_findings(raw_findings)
    score = score_file_level(normalized_findings, gold["expected"], kind)
    payload = {
        "tool": "clang-sa",
        "cwe": cwe,
        "kind": kind,
        "version": subprocess.run([str(CLANG), "--version"], capture_output=True, text=True).stdout.splitlines()[0],
        "checkers": config["checkers"],
        "elapsed_s": elapsed,
        "entries": len(entries),
        "findings": normalized_findings,
        "checker_counts": dict(checker_counts),
        "score": score,
        "notes": config["notes"],
    }
    (cwe_out / "clang-sa.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    return payload


def render_summary(results: list[dict]) -> str:
    lines = [
        "# Clang Static Analyzer on Juliet",
        "",
        "| CWE | Kind | TP | FP | FN | Findings | Gold | Precision | Recall | F1 | Time |",
        "|-----|------|----|----|----|----------|------|-----------|--------|----|------|",
    ]
    for result in results:
        score = result["score"]
        lines.append(
            f"| CWE-{result['cwe']} | {result['kind']} | {score['true_positive']} | {score['false_positive']} | "
            f"{score['false_negative']} | {score['finding_groups']} | {score['gold_total']} | "
            f"{score['precision']:.1%} | {score['recall']:.1%} | {score['f1']:.1%} | {result['elapsed_s']:.1f}s |"
        )
    lines.append("")
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cwes", nargs="+", default=["476", "369", "190", "121"])
    parser.add_argument("--jobs", type=int, default=max(1, (os.cpu_count() or 4) // 2))
    parser.add_argument("--out-root", default=str(RESULTS_ROOT))
    args = parser.parse_args()

    if not SCAN_BUILD.exists():
        print(f"scan-build not found at {SCAN_BUILD}", 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

    stamp = now_iso()
    out_root = Path(args.out_root).resolve() / f"{stamp}_csa-juliet"
    out_root.mkdir(parents=True, exist_ok=True)

    results = []
    for cwe in args.cwes:
        if cwe not in CWE_CONFIG:
            print(f"unsupported CWE: {cwe}", file=sys.stderr)
            return 2
        print(f"running CSA on Juliet CWE-{cwe}...", file=sys.stderr)
        results.append(run_cwe(cwe, args.jobs, out_root))

    summary = render_summary(results)
    (out_root / "summary.md").write_text(summary + "\n")
    print(summary)
    print(f"results written to {out_root}")
    return 0


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