#!/usr/bin/env python3
"""Compare fold verdicts and VCs against CBMC or ESBMC.

This script is intentionally pragmatic:

- It always runs fold on a single input C file and extracts per-function verdicts.
- It runs CBMC or ESBMC with ``--show-vcc`` when either tool is installed.
- If the local toolchain does not expose raw SMT-LIB for one side, it still
  produces a useful agreement table plus an explicit capture template showing
  the command expected for a manual VC dump.

The current fold tree logs solver outcomes via ``--trace`` but does not yet ship
an always-on SMT dump flag. The harness therefore supports three extraction
paths for fold queries, in order:

1. Future ``trace_json`` records that already include ``smtlib``/``query``.
2. Pre-generated ``*.smt2`` files in ``--fold-smt-dir``.
3. A template note explaining how to supply the SMT capture.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import shlex
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_FOLD = REPO_ROOT / "build" / "release" / "src" / "cli" / "fold"
COMMON_CBMC_PATHS = [
    Path("/opt/homebrew/bin/cbmc"),
    Path("/usr/local/bin/cbmc"),
    Path("/usr/bin/cbmc"),
]
COMMON_ESBMC_PATHS = [
    Path("/opt/homebrew/bin/esbmc"),
    Path("/usr/local/bin/esbmc"),
    Path("/usr/bin/esbmc"),
]
COMMON_Z3_PATHS = [
    Path("/opt/homebrew/bin/z3"),
    Path("/usr/local/bin/z3"),
    Path("/usr/bin/z3"),
]

TRACE_JSON_RE = re.compile(r"trace_json\s+(\{.*\})")
FUNC_DEF_RE = re.compile(
    r"(?P<sig>(?:^|\n)\s*(?!if\b|for\b|while\b|switch\b)[A-Za-z_][\w\s\*]*?\s+"
    r"(?P<name>[A-Za-z_]\w*)\s*\([^;{}]*\)\s*)\{",
    re.MULTILINE,
)


@dataclass
class FunctionSpan:
    name: str
    start_line: int
    end_line: int

    def contains(self, line: int) -> bool:
        return self.start_line <= line <= self.end_line


@dataclass
class FoldFinding:
    rule: str
    function: str
    file: str
    line: int
    confidence: str
    message: str


@dataclass
class TraceEvent:
    detector: str
    file: str
    line: int
    col: int
    result: str
    detail: str
    smtlib: str = ""
    query: str = ""


@dataclass
class ToolVerdict:
    verdict: str = "UNKNOWN"
    vc_text: str = ""
    smtlib: str = ""
    source: str = ""
    note: str = ""


@dataclass
class ComparisonRow:
    function: str
    fold_verdict: str
    checker_verdict: str
    equivalent: str


def detect_binary(env_name: str, executable: str, candidates: list[Path]) -> Path | None:
    env_bin = os.environ.get(env_name)
    if env_bin:
        path = Path(env_bin).expanduser()
        if path.is_file():
            return path.resolve()
    which = shutil.which(executable)
    if which:
        return Path(which).resolve()
    for candidate in candidates:
        if candidate.is_file():
            return candidate.resolve()
    return None


def run(command: list[str], cwd: Path | None = None) -> subprocess.CompletedProcess[str]:
    return subprocess.run(command, cwd=cwd, capture_output=True, text=True)


def build_compile_db(source: Path, db_path: Path, cc: str) -> None:
    entry = [{
        "directory": str(source.parent),
        "command": shlex.join([cc, "-std=c11", "-c", str(source)]),
        "file": str(source),
    }]
    db_path.write_text(json.dumps(entry, indent=2) + "\n")


def extract_braced_block(text: str, open_brace: int) -> tuple[int, int] | 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 open_brace, idx
    return None


def discover_functions(source: Path) -> list[FunctionSpan]:
    text = source.read_text()
    spans: list[FunctionSpan] = []
    for match in FUNC_DEF_RE.finditer(text):
        name = match.group("name")
        open_brace = text.find("{", match.start("sig"))
        if open_brace < 0:
            continue
        block = extract_braced_block(text, open_brace)
        if block is None:
            continue
        start_line = text.count("\n", 0, match.start("sig")) + 1
        end_line = text.count("\n", 0, block[1]) + 1
        spans.append(FunctionSpan(name=name, start_line=start_line, end_line=end_line))
    return spans


def function_for_line(spans: list[FunctionSpan], line: int) -> str:
    for span in spans:
        if span.contains(line):
            return span.name
    return ""


def parse_fold_sarif(sarif_path: Path) -> list[FoldFinding]:
    if not sarif_path.exists():
        return []
    doc = json.loads(sarif_path.read_text())
    out: list[FoldFinding] = []
    for run_doc in doc.get("runs", []):
        for result in run_doc.get("results", []):
            locations = result.get("locations", [])
            if not locations:
                continue
            loc = locations[0]
            physical = loc.get("physicalLocation", {})
            artifact = physical.get("artifactLocation", {})
            region = physical.get("region", {})
            logical = loc.get("logicalLocations", [])
            function_name = ""
            if logical:
                function_name = logical[0].get("fullyQualifiedName", "")
            out.append(FoldFinding(
                rule=result.get("ruleId", "fold.unknown").removeprefix("fold."),
                function=function_name,
                file=artifact.get("uri", ""),
                line=int(region.get("startLine", 0) or 0),
                confidence=result.get("properties", {}).get("confidence", "unknown"),
                message=result.get("message", {}).get("text", ""),
            ))
    return out


def parse_trace_events(log_text: str, spans: list[FunctionSpan]) -> dict[str, list[TraceEvent]]:
    events: dict[str, list[TraceEvent]] = {}
    for line in log_text.splitlines():
        match = TRACE_JSON_RE.search(line)
        if not match:
            continue
        try:
            payload = json.loads(match.group(1))
        except json.JSONDecodeError:
            continue
        function_name = payload.get("function", "")
        if not function_name:
            function_name = function_for_line(spans, int(payload.get("line", 0) or 0))
        if not function_name:
            continue
        event = TraceEvent(
            detector=str(payload.get("detector", "")),
            file=str(payload.get("file", "")),
            line=int(payload.get("line", 0) or 0),
            col=int(payload.get("col", 0) or 0),
            result=str(payload.get("result", "unknown")).upper(),
            detail=str(payload.get("detail", "")),
            smtlib=str(payload.get("smtlib", "") or ""),
            query=str(payload.get("query", "") or ""),
        )
        events.setdefault(function_name, []).append(event)
    return events


def find_fold_vc(function_name: str,
                 trace_events: list[TraceEvent],
                 fold_smt_dir: Path | None) -> tuple[str, str]:
    for event in trace_events:
        candidate = event.smtlib or event.query
        if candidate:
            return candidate, "trace_json"
    if fold_smt_dir and fold_smt_dir.is_dir():
        patterns = [
            f"{function_name}.smt2",
            f"{function_name}*.smt2",
            f"*{function_name}*.smt2",
        ]
        for pattern in patterns:
            matches = sorted(fold_smt_dir.glob(pattern))
            if matches:
                return matches[0].read_text(), str(matches[0])
    return "", ""


def fold_verdicts(findings: list[FoldFinding],
                  traces: dict[str, list[TraceEvent]],
                  spans: list[FunctionSpan],
                  fold_smt_dir: Path | None) -> dict[str, ToolVerdict]:
    verdicts = {span.name: ToolVerdict() for span in spans}
    for span in spans:
        vc_text, source = find_fold_vc(span.name, traces.get(span.name, []), fold_smt_dir)
        verdicts[span.name].vc_text = vc_text
        verdicts[span.name].source = source
        verdicts[span.name].smtlib = normalize_to_smtlib(vc_text)

    for finding in findings:
        if not finding.function:
            finding.function = function_for_line(spans, finding.line)
        if not finding.function:
            continue
        verdicts.setdefault(finding.function, ToolVerdict())
        verdicts[finding.function].verdict = "BUG"

    for name, events in traces.items():
        verdict = verdicts.setdefault(name, ToolVerdict())
        if verdict.verdict == "BUG":
            continue
        if any(event.result == "SAT" for event in events):
            verdict.verdict = "BUG"
        elif any(event.result == "UNSAT" for event in events):
            verdict.verdict = "SAFE"

    return verdicts


def cbmc_command(cbmc_bin: Path, source: Path, function_name: str, unwind: int) -> list[str]:
    return [
        str(cbmc_bin),
        str(source),
        "--function", function_name,
        "--pointer-check",
        "--bounds-check",
        "--div-by-zero-check",
        "--signed-overflow-check",
        "--unwind", str(unwind),
        "--show-vcc",
    ]


def esbmc_command(esbmc_bin: Path, source: Path, function_name: str, unwind: int) -> list[str]:
    return [
        str(esbmc_bin),
        str(source),
        "--function", function_name,
        "--pointer-check",
        "--bounds-check",
        "--div-by-zero-check",
        "--overflow",
        "--unwind", str(unwind),
        "--no-unwinding-assertions",
        "--show-vcc",
    ]


def checker_verdict_from_output(text: str) -> str:
    upper = text.upper()
    if "VERIFICATION FAILED" in upper:
        return "BUG"
    if "VERIFICATION SUCCESSFUL" in upper:
        return "SAFE"
    return "UNKNOWN"


def extract_checker_vc(text: str) -> str:
    smt_match = re.search(r"(\(set-logic\b.*)", text, re.DOTALL)
    if smt_match:
        return smt_match.group(1)
    vcc_markers = [
        "VERIFICATION CONDITIONS:",
        "Verification conditions:",
        "VCCs:",
    ]
    for marker in vcc_markers:
        idx = text.find(marker)
        if idx >= 0:
            return text[idx + len(marker):].strip()
    return text.strip()


def normalize_to_smtlib(vc_text: str) -> str:
    if not vc_text:
        return ""
    stripped = vc_text.strip()
    if stripped.startswith("(set-logic") or "(assert" in stripped:
        return stripped if stripped.endswith("\n") else stripped + "\n"
    return ""


def run_checker(tool_name: str,
                tool_bin: Path,
                source: Path,
                function_name: str,
                unwind: int) -> ToolVerdict:
    command = cbmc_command(tool_bin, source, function_name, unwind) if tool_name == "cbmc" else esbmc_command(tool_bin, source, function_name, unwind)
    proc = run(command)
    text = "\n".join(part for part in (proc.stdout, proc.stderr) if part)
    vc_text = extract_checker_vc(text)
    return ToolVerdict(
        verdict=checker_verdict_from_output(text),
        vc_text=vc_text,
        smtlib=normalize_to_smtlib(vc_text),
        source=f"{tool_name} --show-vcc",
        note="",
    )


def run_z3(z3_bin: Path, smtlib: str) -> str:
    proc = subprocess.run([str(z3_bin), "-in"], input=smtlib, text=True, capture_output=True)
    text = (proc.stdout or proc.stderr).strip().splitlines()
    if not text:
        return "unknown"
    first = text[0].strip().lower()
    if first in {"sat", "unsat", "unknown"}:
        return first
    return "unknown"


def explain_mismatch(fold_verdict: str, checker_verdict: str) -> str:
    if fold_verdict == "BUG" and checker_verdict == "SAFE":
        return "NO (fold over-approximates)"
    if fold_verdict == "SAFE" and checker_verdict == "BUG":
        return "NO (fold misses bounded bug)"
    return "NO"


def print_template(source: Path, fold_bin: Path, checker_name: str) -> None:
    checker_cmd = "cbmc" if checker_name == "cbmc" else "esbmc"
    print("VC Equivalence Check")
    print("====================")
    print("No CBMC/ESBMC binary was found locally, so this run emitted a template.")
    print("")
    print("Expected commands:")
    print(f"  {fold_bin} scan --compile-db <compile_commands.json> --format sarif --output fold.sarif --trace {source}")
    print(f"  {checker_cmd} {source} --function <fn> --show-vcc <other-check-flags>")
    print("")
    print("If fold can dump SMT-LIB, pass it back in with:")
    print("  python3 bench/compare_vc_with_cbmc.py <file.c> --fold-smt-dir <dir-with-smt2>")


def render_table(rows: list[ComparisonRow], checker_label: str) -> str:
    lines = [
        "VC Equivalence Check",
        "====================",
        f"Function          | fold verdict | {checker_label} verdict | VCs equivalent?",
    ]
    yes = 0
    for row in rows:
        if row.equivalent == "YES":
            yes += 1
        lines.append(f"{row.function:<17} | {row.fold_verdict:<12} | {row.checker_verdict:<11} | {row.equivalent}")
    total = len(rows)
    rate = 0.0 if total == 0 else (100.0 * yes / total)
    lines.append("")
    lines.append(f"Agreement rate: {rate:.0f}% ({yes}/{total})")
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source", type=Path, help="Single input C file")
    parser.add_argument("--fold-bin", type=Path, default=DEFAULT_FOLD)
    parser.add_argument("--fold-smt-dir", type=Path, default=None,
                        help="Directory containing fold-emitted *.smt2 files")
    parser.add_argument("--cc", default=os.environ.get("CC", "clang"))
    parser.add_argument("--loop-unroll", type=int, default=3)
    parser.add_argument("--keep-temp", action="store_true")
    args = parser.parse_args()

    source = args.source.resolve()
    if not source.is_file():
        print(f"missing source file: {source}", file=sys.stderr)
        return 2
    if source.suffix != ".c":
        print(f"expected a C file, got: {source}", file=sys.stderr)
        return 2

    fold_bin = args.fold_bin.resolve()
    if not fold_bin.is_file():
        print(f"missing fold binary: {fold_bin}", file=sys.stderr)
        return 2

    cbmc_bin = detect_binary("CBMC_BIN", "cbmc", COMMON_CBMC_PATHS)
    esbmc_bin = detect_binary("ESBMC_BIN", "esbmc", COMMON_ESBMC_PATHS)
    z3_bin = detect_binary("Z3_BIN", "z3", COMMON_Z3_PATHS)

    checker_name = "cbmc" if cbmc_bin else "esbmc" if esbmc_bin else ""
    checker_bin = cbmc_bin or esbmc_bin
    if checker_bin is None:
        print_template(source, fold_bin, "cbmc")
        return 0

    spans = discover_functions(source)
    if not spans:
        print(f"no function definitions found in {source}", file=sys.stderr)
        return 2

    temp_dir_ctx = tempfile.TemporaryDirectory(prefix="fold-vc-compare-")
    work_dir = Path(temp_dir_ctx.name)
    try:
        db_path = work_dir / "compile_commands.json"
        sarif_path = work_dir / "fold.sarif"
        fold_log = work_dir / "fold.log"
        build_compile_db(source, db_path, args.cc)

        fold_cmd = [
            str(fold_bin), "scan",
            "--compile-db", str(db_path),
            "--format", "sarif",
            "--output", str(sarif_path),
            "--trace",
            "--min-body-stmts", "1",
            "--processes", "1",
            "--loop-unroll", str(args.loop_unroll),
            "--detectors", "dead,div-by-zero,null-deref,int-overflow,oob",
            str(source),
        ]
        fold_proc = run(fold_cmd, cwd=source.parent)
        fold_text = "\n".join(part for part in (fold_proc.stdout, fold_proc.stderr) if part)
        fold_log.write_text(fold_text)

        findings = parse_fold_sarif(sarif_path)
        traces = parse_trace_events(fold_text, spans)
        fold_by_fn = fold_verdicts(findings, traces, spans, args.fold_smt_dir)

        rows: list[ComparisonRow] = []
        for span in spans:
            checker = run_checker(checker_name, checker_bin, source, span.name, args.loop_unroll)
            fold = fold_by_fn.get(span.name, ToolVerdict())
            if fold.verdict == "UNKNOWN" or checker.verdict == "UNKNOWN":
                continue

            equivalent = ""
            if fold.verdict != checker.verdict:
                equivalent = explain_mismatch(fold.verdict, checker.verdict)
            elif z3_bin and fold.smtlib and checker.smtlib:
                fold_sat = run_z3(z3_bin, fold.smtlib)
                checker_sat = run_z3(z3_bin, checker.smtlib)
                equivalent = "YES" if fold_sat == checker_sat else f"NO ({fold_sat} vs {checker_sat})"
            elif not z3_bin:
                equivalent = "N/A (z3 not installed)"
            elif not fold.smtlib:
                equivalent = "N/A (missing fold SMT dump)"
            elif not checker.smtlib:
                equivalent = f"N/A ({checker_name} VC not SMT-LIB)"
            else:
                equivalent = "N/A"

            rows.append(ComparisonRow(
                function=span.name,
                fold_verdict=fold.verdict,
                checker_verdict=checker.verdict,
                equivalent=equivalent,
            ))

        checker_label = "CBMC" if checker_name == "cbmc" else "ESBMC"
        print(render_table(rows, checker_label))
        if rows and any(row.equivalent.startswith("N/A") for row in rows):
            print("")
            print("Notes:")
            print(f"- fold trace log: {fold_log}")
            print(f"- checker backend: {checker_name} ({checker_bin})")
            print("- For full SMT equivalence, provide fold SMT dumps via --fold-smt-dir or extend fold trace_json with an smtlib/query field.")
        return 0
    finally:
        if args.keep_temp:
            print(f"kept temp dir: {work_dir}", file=sys.stderr)
        else:
            temp_dir_ctx.cleanup()


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