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

import argparse
import concurrent.futures
import datetime as dt
import json
import os
import re
import shutil
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"
COMMON_ESBMC_PATHS = [
    Path("/opt/homebrew/bin/esbmc"),
    Path("/usr/local/bin/esbmc"),
    Path("/usr/bin/esbmc"),
]


CWE_CONFIG = {
    "476": {
        "kind": "null-deref",
        "notes": "ESBMC pointer checks are filtered down to null-dereference property violations.",
    },
}


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 ensure_text(value: str | bytes | None) -> str:
    if value is None:
        return ""
    if isinstance(value, bytes):
        return value.decode(errors="replace")
    return value


def detect_esbmc() -> Path | None:
    env_bin = os.environ.get("ESBMC_BIN")
    if env_bin:
        candidate = Path(env_bin).expanduser()
        if candidate.is_file():
            return candidate.resolve()
    which = shutil.which("esbmc")
    if which:
        return Path(which).resolve()
    for candidate in COMMON_ESBMC_PATHS:
        if candidate.is_file():
            return candidate.resolve()
    return None


def print_install_instructions() -> None:
    print("ESBMC was not found in PATH or common install locations.", file=sys.stderr)
    print("Install it with:", file=sys.stderr)
    print("brew install esbmc    # macOS", file=sys.stderr)
    print("# or build from source: https://github.com/esbmc/esbmc", file=sys.stderr)


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


def read_tool_version(esbmc_bin: Path) -> str:
    proc = subprocess.run([str(esbmc_bin), "--version"], capture_output=True, text=True)
    output = (proc.stdout or proc.stderr).splitlines()
    return output[0] if output else str(esbmc_bin)


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 collect_include_dirs(corpus_dir: Path) -> list[Path]:
    include_dirs: set[Path] = set()
    for path in sorted(corpus_dir.rglob("*")):
        if path.is_dir() and path.name == "testcasesupport":
            include_dirs.add(path)
        if path.is_dir() and any(path.glob("*.h")):
            include_dirs.add(path)
    return sorted(include_dirs)


def collect_support_sources(corpus_dir: Path) -> list[Path]:
    support_dir = corpus_dir / "testcasesupport"
    if not support_dir.is_dir():
        return []
    sources: list[Path] = []
    for path in sorted(support_dir.glob("*.c")):
        if path.stem.startswith("main"):
            continue
        sources.append(path)
    for path in sorted(support_dir.glob("*.cpp")):
        if path.stem.startswith("main"):
            continue
        sources.append(path)
    return sources


def collect_testcase_sources(testcase_file: Path) -> list[Path]:
    match = re.match(r"(.+?_\d+)(?:[a-e]|_(?:bad|goodB2G|goodG2B))?$", testcase_file.stem)
    prefix = match.group(1) if match else testcase_file.stem
    siblings = sorted(testcase_file.parent.glob(f"{prefix}*.{testcase_file.suffix.lstrip('.')}"))
    return siblings or [testcase_file]


def parse_null_deref(text: str) -> dict | None:
    lines = text.splitlines()
    for index, line in enumerate(lines):
        if line.strip() != "Violated property:":
            continue
        block = "\n".join(lines[index:index + 6])
        if "NULL pointer" not in block and "null pointer" not in block:
            continue
        if "unwinding assertion" in block:
            continue
        location_match = re.search(r"file\s+(.+?)\s+line\s+(\d+)(?:\s+column\s+\d+)?\s+function\s+(.+)$", block, re.MULTILINE)
        property_line = 0
        property_file = ""
        function_name = ""
        if location_match:
            property_file = location_match.group(1).strip()
            property_line = int(location_match.group(2))
            function_name = location_match.group(3).strip()
        message = ""
        for candidate in lines[index + 1:index + 6]:
            stripped = candidate.strip()
            if "NULL pointer" in stripped or "null pointer" in stripped:
                message = stripped
                break
        return {
            "property_file": property_file,
            "property_line": property_line,
            "function": function_name,
            "message": message or "null pointer property violation",
            "block": block,
        }
    return None


def esbmc_command(esbmc_bin: Path, testcase_file: Path, corpus_dir: Path) -> list[str]:
    sources = collect_testcase_sources(testcase_file)
    support_sources = collect_support_sources(corpus_dir)
    include_dirs = collect_include_dirs(corpus_dir)
    command = [
        str(esbmc_bin),
        "--overflow",
        "--memory-leak-check",
        "--div-by-zero-check",
        "--pointer-check",
        "--bounds-check",
        "--unwind",
        "3",
        "--no-unwinding-assertions",
        "--timeout",
        "30",
        "-DINCLUDEMAIN",
        "-DOMITGOOD",
    ]
    if testcase_file.suffix == ".c":
        command.append("-std=c11")
    else:
        command.append("-std=c++17")
    for include_dir in include_dirs:
        command.extend(["-I", str(include_dir)])
    command.extend(str(path) for path in sources)
    command.extend(str(path) for path in support_sources)
    return command


def run_testcase(esbmc_bin: Path, testcase_file: Path, corpus_dir: Path, kind: str, cases_dir: Path) -> dict:
    command = esbmc_command(esbmc_bin, testcase_file, corpus_dir)
    started = time.monotonic()
    timed_out = False
    try:
        proc = subprocess.run(command, cwd=corpus_dir, capture_output=True, text=True, timeout=45)
        exit_code = proc.returncode
        stdout = proc.stdout
        stderr = proc.stderr
    except subprocess.TimeoutExpired as exc:
        timed_out = True
        exit_code = 124
        stdout = ensure_text(exc.stdout)
        stderr = ensure_text(exc.stderr)
    elapsed = time.monotonic() - started
    combined = "\n".join(part for part in (stdout, stderr) if part)
    violation = parse_null_deref(combined)
    finding = None
    if violation is not None:
        finding = {
            "file": str(testcase_file),
            "line": violation["property_line"],
            "kind": kind,
            "id": f"{kind}:{file_basename(str(testcase_file))}",
            "property_file": violation["property_file"],
            "function": violation["function"],
            "message": violation["message"],
        }
    payload = {
        "testcase": str(testcase_file),
        "command": command,
        "elapsed_s": elapsed,
        "exit_code": exit_code,
        "timed_out": timed_out,
        "verification_failed": "VERIFICATION FAILED" in combined,
        "verification_successful": "VERIFICATION SUCCESSFUL" in combined,
        "finding": finding,
        "stdout": stdout,
        "stderr": stderr,
    }
    case_name = testcase_file.name.replace(testcase_file.suffix, "")
    (cases_dir / f"{case_name}.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    return payload


def run_cwe(cwe: str, esbmc_bin: Path, 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
    cases_dir = cwe_out / "cases"
    cases_dir.mkdir(parents=True, exist_ok=True)
    testcase_files = [corpus_dir / entry["files"][0] for entry in gold["expected"] if entry.get("kind") == kind]

    raw_results: list[dict] = []
    start = time.monotonic()
    with concurrent.futures.ThreadPoolExecutor(max_workers=jobs) as pool:
        futures = [pool.submit(run_testcase, esbmc_bin, testcase_file, corpus_dir, kind, cases_dir) for testcase_file in testcase_files]
        for future in concurrent.futures.as_completed(futures):
            raw_results.append(future.result())
    elapsed = time.monotonic() - start

    findings = sorted(
        [result["finding"] for result in raw_results if result["finding"] is not None],
        key=lambda finding: (finding["file"], finding["line"]),
    )
    score = score_file_level(findings, gold["expected"], kind)
    payload = {
        "tool": "esbmc",
        "cwe": cwe,
        "kind": kind,
        "version": read_tool_version(esbmc_bin),
        "command_flags": [
            "--overflow",
            "--memory-leak-check",
            "--div-by-zero-check",
            "--pointer-check",
            "--bounds-check",
            "--unwind 3",
            "--no-unwinding-assertions",
            "--timeout 30",
            "-DINCLUDEMAIN",
            "-DOMITGOOD",
        ],
        "elapsed_s": elapsed,
        "entries": len(testcase_files),
        "findings": findings,
        "score": score,
        "notes": config["notes"],
        "timeouts": sorted(file_basename(result["testcase"]) for result in raw_results if result["timed_out"]),
        "verification_failures": sorted(file_basename(result["testcase"]) for result in raw_results if result["verification_failed"]),
    }
    (cwe_out / "esbmc.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    return payload


def render_summary(results: list[dict]) -> str:
    lines = [
        "# ESBMC 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"])
    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()

    esbmc_bin = detect_esbmc()
    if esbmc_bin is None:
        print_install_instructions()
        return 2

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