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

import argparse
import datetime as dt
import json
import os
import re
import shlex
import subprocess
import sys
import time
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"
DEFAULT_FOLD = REPO_ROOT / "build" / "release" / "src" / "cli" / "fold"
LLVM_BIN = Path(os.environ.get("LLVM_BIN", "/opt/homebrew/opt/llvm@22/bin"))
CLANG = LLVM_BIN / "clang"
CLANGXX = LLVM_BIN / "clang++"


CWE_CONFIG = {
    "476": {"kinds": ["null-deref"], "detectors": ["null-deref"]},
    "369": {
        "kinds": ["div-by-zero"],
        "detectors": ["div-by-zero"],
        "exclude_globs": [
            "*CWE369_Divide_by_Zero__float_listenSocket_73a.cpp",
            "*CWE369_Divide_by_Zero__float_listenSocket_74a.cpp",
            "*CWE369_Divide_by_Zero__int_listen_socket_divide_73a.cpp",
            "*CWE369_Divide_by_Zero__int_listen_socket_divide_74a.cpp",
            "*CWE369_Divide_by_Zero__int_listen_socket_modulo_73a.cpp",
            "*CWE369_Divide_by_Zero__int_listen_socket_modulo_74a.cpp",
        ],
    },
    "190": {"kinds": ["int-overflow"], "detectors": ["int-overflow"]},
    "191": {"kinds": ["int-overflow"], "detectors": ["int-overflow"]},
    "194": {"kinds": ["sign-conversion"], "detectors": ["sign-conversion"]},
    "195": {"kinds": ["sign-conversion"], "detectors": ["sign-conversion"]},
    "196": {"kinds": ["sign-conversion"], "detectors": ["sign-conversion"]},
    "197": {"kinds": ["int-truncation"], "detectors": ["int-truncation"]},
    "121": {"kinds": ["oob-access"], "detectors": ["oob"]},
    "122": {"kinds": ["oob-access"], "detectors": ["oob"]},
    "124": {"kinds": ["oob-access"], "detectors": ["oob"]},
    "126": {"kinds": ["oob-access"], "detectors": ["oob"]},
    "127": {"kinds": ["oob-access"], "detectors": ["oob"]},
    "561": {
        "kinds": ["dead", "dead-statement", "unused-function"],
        "detectors": ["dead", "dead-statement", "unused-function"],
    },
}

SOLVER_RE = re.compile(
    r"^\[solver-profile\]\s+(?P<detector>\S+)\s+"
    r"queries=(?P<queries>\d+)\s+sat=(?P<sat>\d+)\s+unsat=(?P<unsat>\d+)\s+"
    r"unknown=(?P<unknown>\d+)\s+decisions=(?P<decisions>\d+)\s+"
    r"conflicts=(?P<conflicts>\d+)\s+restarts=(?P<restarts>\d+)\s+"
    r"propagations=(?P<propagations>\d+)\s+z3_time_ms=(?P<z3_time_ms>\d+)\s+"
    r"wall_ms=(?P<wall_ms>\d+)"
)
EMITTED_RE = re.compile(r"emitted\s+(?P<count>\d+)\s+findings total")


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], kinds: list[str]) -> dict:
    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 {
        "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 normalize_kind(kind: str) -> str:
    return {
        "dead-branch": "dead",
        "duplicate-class": "class-dup",
    }.get(kind, kind)


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 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] = finding
    return sorted(dedup.values(), key=lambda item: (item["file"], item["line"]))


def normalize_finding_record(finding: dict) -> list[dict]:
    if "file" in finding:
        return [{
            "file": finding.get("file", ""),
            "line": int(finding.get("line", 0) or 0),
            "kind": normalize_kind(str(finding.get("kind", "unknown"))),
            "id": finding.get("id", ""),
        }]

    kind = normalize_kind(str(finding.get("kind", "unknown")))
    fid = str(finding.get("id", ""))
    locations = finding.get("locations", [])
    out: list[dict] = []
    for loc in locations:
        out.append({
            "file": loc.get("file", ""),
            "line": int(loc.get("line", 0) or 0),
            "kind": kind,
            "id": fid,
        })
    return out


def parse_sarif(sarif_path: Path) -> list[dict]:
    if not sarif_path.exists():
        return []
    doc = json.loads(sarif_path.read_text())
    findings: list[dict] = []
    for run in doc.get("runs", []):
        for result in run.get("results", []):
            kind = normalize_kind(result.get("ruleId", "fold.unknown").removeprefix("fold."))
            fid = result.get("properties", {}).get("id", "")
            for loc in result.get("locations", []):
                physical = loc.get("physicalLocation", {})
                artifact = physical.get("artifactLocation", {})
                region = physical.get("region", {})
                findings.append({
                    "file": artifact.get("uri", ""),
                    "line": int(region.get("startLine", 0) or 0),
                    "kind": kind,
                    "id": fid,
                })
    return normalize_findings(findings)


def parse_findings_json(json_path: Path) -> list[dict]:
    if not json_path.exists():
        return []
    doc = json.loads(json_path.read_text())

    # Runner payloads wrap the actual normalized findings under "findings".
    if isinstance(doc, dict) and isinstance(doc.get("findings"), list):
        findings: list[dict] = []
        for item in doc["findings"]:
            findings.extend(normalize_finding_record(item))
        return normalize_findings(findings)

    # Worker/native JSON can also show up as NDJSON-style envelopes flattened
    # into a list of {kind: "finding", data: {...}} items.
    if isinstance(doc, list):
        findings = []
        for item in doc:
            if not isinstance(item, dict):
                continue
            if item.get("kind") == "finding" and isinstance(item.get("data"), dict):
                findings.extend(normalize_finding_record(item["data"]))
            else:
                findings.extend(normalize_finding_record(item))
        return normalize_findings(findings)

    return []


def load_findings(cwe_out: Path) -> tuple[list[dict], str]:
    sarif_path = cwe_out / "fold.sarif"
    findings = parse_sarif(sarif_path)
    if findings:
        return findings, "sarif"

    json_path = cwe_out / "fold.findings.json"
    findings = parse_findings_json(json_path)
    if findings:
        return findings, "json"

    payload_path = cwe_out / "fold.json"
    findings = parse_findings_json(payload_path)
    if findings:
        return findings, "payload"

    return [], "none"


def parse_profiles(log_text: str, expected_detectors: list[str]) -> dict:
    totals = {
        "queries": 0,
        "sat": 0,
        "unsat": 0,
        "unknown": 0,
        "decisions": 0,
        "conflicts": 0,
        "restarts": 0,
        "propagations": 0,
        "z3_time_ms": 0,
        "wall_ms": 0,
    }
    expected = set(expected_detectors)
    for line in log_text.splitlines():
        match = SOLVER_RE.match(line.strip())
        if not match:
            continue
        if match.group("detector") not in expected:
            continue
        for key in totals:
            totals[key] += int(match.group(key))
    emitted = 0
    for match in EMITTED_RE.finditer(log_text):
        emitted = int(match.group("count"))
    return {"solver": totals, "findings_emitted_total": emitted}


def run_cwe(cwe: str, fold_bin: Path, processes: int, out_root: Path) -> dict:
    config = CWE_CONFIG[cwe]
    kinds = list(config["kinds"])
    detectors = list(config["detectors"])
    exclude_globs = config.get("exclude_globs", [])
    gold = load_gold(cwe)
    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)
    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", str(processes),
        "--profile-solver",
        "--profile-reach",
    ]
    for detector in detectors:
        command.extend(["--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
    if "subprocess_fanout: bad NDJSON line:" in combined_log:
        raise RuntimeError(
            f"parallel worker serialization dropped findings for CWE-{cwe}; rerun with --processes 1"
        )

    findings, findings_source = load_findings(cwe_out)
    if not findings and not sarif_path.exists() and proc.returncode < 0:
        raise RuntimeError(
            f"fold scan for CWE-{cwe} terminated before writing results (rc={proc.returncode})"
        )
    score = score_file_level(findings, gold["expected"], kinds)
    profile = parse_profiles(combined_log, detectors + kinds)
    payload = {
        "tool": "fold",
        "cwe": cwe,
        "kind": ",".join(kinds),
        "kinds": kinds,
        "detector": ",".join(detectors),
        "detectors": detectors,
        "version": subprocess.run([str(fold_bin), "--version"], capture_output=True, text=True).stdout.strip() or "unknown",
        "elapsed_s": elapsed,
        "exit_code": proc.returncode,
        "entries": len(entries),
        "excluded_globs": exclude_globs,
        "findings": findings,
        "findings_source": findings_source,
        "score": score,
        "profile": profile,
        "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(results: list[dict]) -> str:
    lines = [
        "# fold on Juliet",
        "",
        "| CWE | Kind | TP | FP | FN | Findings | Gold | Precision | Recall | F1 | Time | Queries | Sat | Unsat | Unknown |",
        "|-----|------|----|----|----|----------|------|-----------|--------|----|------|---------|-----|-------|---------|",
    ]
    for result in results:
        score = result["score"]
        solver = result["profile"]["solver"]
        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 | "
            f"{solver['queries']} | {solver['sat']} | {solver['unsat']} | {solver['unknown']} |"
        )
    lines.append("")
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--cwes",
        nargs="+",
        default=[
            "476", "369", "190", "191", "194", "195", "196", "197",
            "121", "122", "124", "126", "127", "561",
        ],
    )
    parser.add_argument("--fold-bin", default=str(DEFAULT_FOLD))
    parser.add_argument("--processes", type=int, default=1)
    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

    stamp = now_iso()
    out_root = Path(args.out_root).resolve() / f"{stamp}_fold-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 fold on Juliet CWE-{cwe}...", file=sys.stderr)
        results.append(run_cwe(cwe, fold_bin, args.processes, 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())