#!/usr/bin/env python3
"""Compare Fold SAST classifications with replayed sanitizer vulnerable/fixed pairs."""

from __future__ import annotations

import argparse
import importlib.util
import json
import re
import subprocess
from pathlib import Path
from typing import Any, Mapping

from fold_eval.confirmation import summarize_confirmation_results
from fold_eval.differential import fixed_runtime_reported, vulnerable_runtime_detected
from fold_eval.metrics import aggregate_classification_metrics
from fold_eval.source_lane import run_fold_source
from fold_eval.toolchain import find_clang
from fold_eval.schema import (
    artifact_record,
    build_run_artifact,
    sha256_file,
    utc_now,
    validate_benchmark_case,
    validate_run_artifact,
    write_json,
)


REPO_ROOT = Path(__file__).resolve().parents[2]
RUNNER_COMMON_PATH = REPO_ROOT / "scripts" / "benchmarks" / "runner_common.py"
DEFAULT_SPEC = (
    REPO_ROOT / "tests" / "benchmarks" / "differential" / "sanitizer_pairs.json"
)
_spec = importlib.util.spec_from_file_location("runner_common", RUNNER_COMMON_PATH)
if _spec is None or _spec.loader is None:
    raise ImportError(f"failed to load {RUNNER_COMMON_PATH}")
common = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(common)


def sanitizer_flags(name: str) -> list[str]:
    return [f"-fsanitize={name}", f"-fno-sanitize-recover={name}"]


def probe_sanitizer(compiler: str, output_dir: Path, sanitizers: list[str]) -> str | None:
    source = output_dir / "sanitizer-probe.cpp"
    binary = output_dir / "sanitizer-probe"
    source.write_text("int main() { return 0; }\n", encoding="utf-8")
    flags = [flag for sanitizer in sanitizers for flag in sanitizer_flags(sanitizer)]
    proc = subprocess.run(
        [compiler, "-std=c++20", *flags, str(source), "-o", str(binary)],
        cwd=REPO_ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    return None if proc.returncode == 0 else proc.stdout[-4000:]


def build_runtime(
    compiler: str,
    source: Path,
    sanitizer: str,
    optimization: str,
    label: str,
    output_dir: Path,
) -> tuple[Path | None, dict[str, Any], list[str]]:
    lane_dir = output_dir / label / "runtime"
    lane_dir.mkdir(parents=True, exist_ok=True)
    binary = lane_dir / "subject"
    log = lane_dir / "compile.log"
    command = [
        compiler,
        "-std=c++20",
        optimization,
        "-g",
        *sanitizer_flags(sanitizer),
        str(source),
        "-o",
        str(binary),
    ]
    proc = subprocess.run(
        command,
        cwd=REPO_ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    log.write_text(proc.stdout, encoding="utf-8")
    failures = [] if proc.returncode == 0 else [f"sanitizer compile exited {proc.returncode}"]
    return (binary if proc.returncode == 0 else None), {
        "command": command,
        "exit_status": proc.returncode,
        "log": str(log),
    }, failures


def run_runtime_attempt(
    *,
    binary: Path,
    source: Path,
    compiler: str,
    runtime_tool: str,
    case_id: str,
    label: str,
    attempt: int,
    output_dir: Path,
) -> tuple[dict[str, Any], list[str]]:
    lane_dir = output_dir / label / "runtime"
    log_path = lane_dir / f"attempt-{attempt}.log"
    run_path = lane_dir / f"attempt-{attempt}.run.json"
    command = [str(binary)]
    started_at = utc_now()
    proc = common.run_measured(command, REPO_ROOT, timeout_seconds=10.0)
    ended_at = utc_now()
    log_path.write_text(proc.output, encoding="utf-8")
    run_artifact = build_run_artifact(
        experiment_id=f"differential:{case_id}:{label}:runtime:{attempt}",
        case_id=case_id,
        tool={
            "name": f"{runtime_tool}-subject",
            "compiler": compiler,
            "sha256": sha256_file(binary),
        },
        target={"source": str(source), "input_fingerprint": sha256_file(source)},
        command=command,
        started_at=started_at,
        ended_at=ended_at,
        exit_status=proc.returncode,
        budgets={"wall_seconds": 10.0},
        artifacts=[
            artifact_record("runtime-log", log_path, "text/plain"),
            artifact_record("instrumented-binary", binary, "application/octet-stream"),
        ],
        measurements=proc.metrics(),
    )
    failures = [
        f"run artifact: {error}" for error in validate_run_artifact(run_artifact)
    ]
    write_json(run_path, run_artifact)
    return {
        "returncode": proc.returncode,
        "output": proc.output,
        "measurements": proc.metrics(),
        "run_artifact": str(run_path),
    }, failures


def resolve_source(pair: Mapping[str, Any], field: str, spec_path: Path) -> Path:
    source = common.resolve_path(str(pair.get(field, "")), spec_path.parent)
    if not source.is_file():
        raise FileNotFoundError(source)
    return source


def evaluate_pair(
    *,
    pair: Mapping[str, Any],
    spec: Mapping[str, Any],
    spec_path: Path,
    fold: Path,
    compiler: str,
    output_dir: Path,
) -> tuple[dict[str, Any], list[str], dict[str, dict[str, int]]]:
    name = str(pair["name"])
    vulnerable_source = resolve_source(pair, "vulnerable_source", spec_path)
    fixed_source = resolve_source(pair, "fixed_source", spec_path)
    detector = str(pair["detector"])
    rule_id = str(pair["rule_id"])
    sanitizer = str(pair["sanitizer"])
    optimization = str(pair.get("optimization", spec.get("optimization", "-O1")))
    if optimization not in {"-O0", "-O1", "-O2", "-O3"}:
        raise ValueError(f"unsupported optimization level: {optimization}")
    runtime_pattern = re.compile(str(pair["runtime_regex"]), re.IGNORECASE)
    runtime_tool = str(spec.get("runtime_tool", "ubsan")).lower()
    failures: list[str] = []

    vulnerable_static, errors = run_fold_source(
        fold=fold,
        source=vulnerable_source,
        detector=detector,
        case_id=str(spec["case_id"]),
        label=f"{name}-vulnerable",
        output_dir=output_dir,
        repo_root=REPO_ROOT,
    )
    failures.extend(errors)
    fixed_static, errors = run_fold_source(
        fold=fold,
        source=fixed_source,
        detector=detector,
        case_id=str(spec["case_id"]),
        label=f"{name}-fixed",
        output_dir=output_dir,
        repo_root=REPO_ROOT,
    )
    failures.extend(errors)
    static_vulnerable = any(item["rule_id"] == rule_id for item in vulnerable_static["findings"])
    static_fixed = any(item["rule_id"] == rule_id for item in fixed_static["findings"])

    vulnerable_binary, vulnerable_build, errors = build_runtime(
        compiler,
        vulnerable_source,
        sanitizer,
        optimization,
        f"{name}-vulnerable",
        output_dir,
    )
    failures.extend(errors)
    fixed_binary, fixed_build, errors = build_runtime(
        compiler, fixed_source, sanitizer, optimization, f"{name}-fixed", output_dir
    )
    failures.extend(errors)
    if vulnerable_binary is None or fixed_binary is None:
        return {
            "name": name,
            "status": "failed",
            "static": {"vulnerable": vulnerable_static, "fixed": fixed_static},
            "build": {"vulnerable": vulnerable_build, "fixed": fixed_build},
        }, failures, {}

    vulnerable_attempts: list[dict[str, Any]] = []
    for attempt in (1, 2):
        result, errors = run_runtime_attempt(
            binary=vulnerable_binary,
            source=vulnerable_source,
            compiler=compiler,
            runtime_tool=runtime_tool,
            case_id=str(spec["case_id"]),
            label=f"{name}-vulnerable",
            attempt=attempt,
            output_dir=output_dir,
        )
        vulnerable_attempts.append(result)
        failures.extend(errors)
    fixed_attempt, errors = run_runtime_attempt(
        binary=fixed_binary,
        source=fixed_source,
        compiler=compiler,
        runtime_tool=runtime_tool,
        case_id=str(spec["case_id"]),
        label=f"{name}-fixed",
        attempt=1,
        output_dir=output_dir,
    )
    failures.extend(errors)

    vulnerable_detections = [
        vulnerable_runtime_detected(
            int(attempt["returncode"]), str(attempt["output"]), runtime_pattern
        )
        for attempt in vulnerable_attempts
    ]
    runtime_vulnerable = all(vulnerable_detections)
    # A normal program may intentionally return a non-zero application status.
    # Only the sanitizer oracle identifies the fixed revision as unsafe.
    runtime_fixed = fixed_runtime_reported(
        int(fixed_attempt["returncode"]), str(fixed_attempt["output"]), runtime_pattern
    )
    if not static_vulnerable:
        failures.append("Fold missed the vulnerable revision")
    if static_fixed:
        failures.append("Fold reported the fixed revision")
    if not runtime_vulnerable:
        failures.append(
            f"{runtime_tool.upper()} did not detect and replay the vulnerable revision"
        )
    if runtime_fixed:
        failures.append(f"{runtime_tool.upper()} reported the fixed revision")
    if static_vulnerable != runtime_vulnerable or static_fixed != runtime_fixed:
        failures.append(f"Fold and {runtime_tool.upper()} classifications disagree")

    first_ms = int(float(vulnerable_attempts[0]["measurements"]["wall_seconds"]) * 1000)
    progress: dict[str, Any] = {
        "reached": True,
        "triggered": runtime_vulnerable,
        "detected": runtime_vulnerable,
        "replayed": runtime_vulnerable and len(vulnerable_attempts) == 2,
        "time_to_reached_ms": 0,
    }
    if runtime_vulnerable:
        progress.update(
            {
                "time_to_triggered_ms": first_ms,
                "time_to_detected_ms": first_ms,
                "time_to_replayed_ms": sum(
                    int(float(item["measurements"]["wall_seconds"]) * 1000)
                    for item in vulnerable_attempts
                ),
            }
        )
    confirmation = {
        "finding_id": f"{pair['hazard']}:{name}",
        "status": "confirmed" if runtime_vulnerable else "not-confirmed",
        "progress": progress,
    }
    slices = {
        f"fold/{pair['hazard']}": {
            "true_positive": int(static_vulnerable),
            "false_positive": int(static_fixed),
            "false_negative": int(not static_vulnerable),
            "true_negative": int(not static_fixed),
        },
        f"{runtime_tool}/{pair['hazard']}": {
            "true_positive": int(runtime_vulnerable),
            "false_positive": int(runtime_fixed),
            "false_negative": int(not runtime_vulnerable),
            "true_negative": int(not runtime_fixed),
        },
    }
    return {
        "name": name,
        "hazard": pair["hazard"],
        "status": "failed" if failures else "ok",
        "classifications": {
            "fold": {"vulnerable": static_vulnerable, "fixed": static_fixed},
            runtime_tool: {"vulnerable": runtime_vulnerable, "fixed": runtime_fixed},
        },
        "static": {"vulnerable": vulnerable_static, "fixed": fixed_static},
        "runtime": {
            "vulnerable_attempts": [
                {key: value for key, value in item.items() if key != "output"}
                for item in vulnerable_attempts
            ],
            "fixed_attempt": {
                key: value for key, value in fixed_attempt.items() if key != "output"
            },
            "confirmation": confirmation,
        },
        "build": {"vulnerable": vulnerable_build, "fixed": fixed_build},
        "failures": failures,
    }, failures, slices


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--build", type=Path, default=Path("build/release"))
    parser.add_argument("--fold", type=Path)
    parser.add_argument("--compiler")
    parser.add_argument("--spec", type=Path, default=DEFAULT_SPEC)
    parser.add_argument(
        "--output-dir", type=Path, default=Path("/tmp/fold-sanitizer-differential")
    )
    args = parser.parse_args()

    build = args.build if args.build.is_absolute() else REPO_ROOT / args.build
    fold = common.find_fold_binary(REPO_ROOT, build, args.fold)
    if fold is None:
        common.print_missing_fold(args.fold or build / "src" / "cli" / "fold")
        return 2
    spec_path = args.spec.resolve()
    try:
        spec = common.load_json(spec_path)
        validation = validate_benchmark_case(spec, allow_legacy=False)
        if not validation["valid"]:
            raise ValueError("; ".join(validation["errors"]))
        pairs = spec.get("pairs", [])
        if not isinstance(pairs, list) or not pairs:
            raise ValueError("pairs must be a non-empty array")
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        print(json.dumps({"status": "failed", "reason": str(exc)}, indent=2))
        return 2

    args.output_dir.mkdir(parents=True, exist_ok=True)
    compiler = find_clang(args.compiler)
    if compiler is None:
        summary = {"status": "skipped", "reason": "no clang++ compiler found"}
        write_json(args.output_dir / "summary.json", summary)
        print(json.dumps(summary, indent=2))
        return 77
    runtime_tool = str(spec.get("runtime_tool", "ubsan")).lower()
    if runtime_tool not in {"asan", "ubsan"}:
        summary = {
            "status": "failed",
            "reason": f"unsupported runtime_tool: {runtime_tool}",
        }
        write_json(args.output_dir / "summary.json", summary)
        print(json.dumps(summary, indent=2))
        return 2
    sanitizers = sorted({str(pair.get("sanitizer", "")) for pair in pairs})
    probe_error = probe_sanitizer(compiler, args.output_dir, sanitizers)
    if probe_error is not None:
        summary = {
            "status": "skipped",
            "reason": f"compiler cannot link the required {runtime_tool.upper()} classes",
            "compiler": compiler,
            "output": probe_error,
        }
        write_json(args.output_dir / "summary.json", summary)
        print(json.dumps(summary, indent=2))
        return 77

    results: list[dict[str, Any]] = []
    failures: list[str] = []
    slices: dict[str, dict[str, int]] = {}
    confirmations: list[dict[str, Any]] = []
    for pair in pairs:
        if not isinstance(pair, dict):
            failures.append("pair is not an object")
            continue
        try:
            result, pair_failures, pair_slices = evaluate_pair(
                pair=pair,
                spec=spec,
                spec_path=spec_path,
                fold=fold,
                compiler=compiler,
                output_dir=args.output_dir,
            )
            results.append(result)
            failures.extend(f"{pair.get('name')}: {item}" for item in pair_failures)
            slices.update(pair_slices)
            confirmation = result.get("runtime", {}).get("confirmation")
            if isinstance(confirmation, dict):
                confirmations.append(confirmation)
        except (OSError, ValueError, TypeError, KeyError) as exc:
            failures.append(f"{pair.get('name', 'pair')}: {exc}")

    summary = {
        "schema_version": 1,
        "suite": f"{runtime_tool}_differential",
        "case_id": spec["case_id"],
        "runtime_tool": runtime_tool,
        "compiler": compiler,
        "status": "failed" if failures else "ok",
        "results": results,
        "classification_metrics": aggregate_classification_metrics(slices),
        "confirmation": summarize_confirmation_results(confirmations),
        "failures": failures,
    }
    write_json(args.output_dir / "summary.json", summary)
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 1 if failures else 0


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