#!/usr/bin/env python3
"""Prototype KLEE adapter for Fold FindingConfirmationTask JSON.

This is intentionally narrow: it consumes one source task, compiles the task's
source artifact to LLVM bitcode with KLEE enabled, runs KLEE when installed, and
writes a FindingConfirmationResult-shaped JSON object. It is a test/confirmation
prototype, not a general harness generator.
"""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
from pathlib import Path


def result(
    task: dict, status: str, message: str, artifacts=None, diagnostics=None
) -> dict:
    out = {
        "version": "1",
        "task_id": task.get("task_id", ""),
        "finding_id": task.get("finding_id", ""),
        "status": status,
        "adapter": "klee",
        "message": message,
        "artifacts": artifacts or [],
        "diagnostics": diagnostics or [],
    }
    if status == "confirmed":
        # KLEE confirmation is a symbolic witness, not an ASan crash. Keep the
        # result confirmed but do not promote to concrete-witness here.
        out["promoted_confidence"] = "proved"
    return out


def klee_include_dir(klee: str) -> str:
    candidate = Path(klee).resolve().parent.parent / "include"
    return str(candidate)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--task", required=True, help="FindingConfirmationTask JSON file"
    )
    parser.add_argument("--output", required=True, help="Confirmation result JSON file")
    parser.add_argument("--work-dir", required=True, help="scratch/output directory")
    parser.add_argument("--clang", default=shutil.which("clang") or "clang")
    args = parser.parse_args()

    task_path = Path(args.task)
    out_path = Path(args.output)
    work_dir = Path(args.work_dir)
    work_dir.mkdir(parents=True, exist_ok=True)

    with task_path.open() as fh:
        task = json.load(fh)

    source = Path(task.get("target", {}).get("artifact", ""))
    if not source.exists():
        out_path.write_text(
            json.dumps(
                result(task, "unsupported", f"source not found: {source}"), indent=2
            )
            + "\n"
        )
        return 0

    klee = shutil.which("klee")
    if not klee:
        out_path.write_text(
            json.dumps(result(task, "unsupported", "klee not found"), indent=2) + "\n"
        )
        print("skip: klee not found")
        return 0

    bitcode = work_dir / (source.stem + ".bc")
    klee_out = work_dir / "klee-out"
    compile_cmd = [
        args.clang,
        "-DKLEE",
        "-I",
        klee_include_dir(klee),
        "-emit-llvm",
        "-c",
        "-g",
        "-O0",
        str(source),
        "-o",
        str(bitcode),
    ]
    compile_run = subprocess.run(compile_cmd, capture_output=True, text=True)
    if compile_run.returncode != 0:
        out_path.write_text(
            json.dumps(
                result(
                    task,
                    "unsupported",
                    "clang failed to compile KLEE bitcode",
                    diagnostics=[compile_run.stderr],
                ),
                indent=2,
            )
            + "\n"
        )
        return 0

    run_cmd = [klee, f"--output-dir={klee_out}", str(bitcode)]
    klee_run = subprocess.run(run_cmd, capture_output=True, text=True)
    tests = sorted(klee_out.glob("*.ktest")) if klee_out.exists() else []
    artifacts = [
        {"role": "bitcode", "path": str(bitcode)},
        {"role": "klee-output", "path": str(klee_out)},
    ]
    artifacts.extend({"role": "ktest", "path": str(path)} for path in tests)

    status = "confirmed" if tests else "timeout"
    message = (
        "KLEE produced witness test(s)" if tests else "KLEE produced no witness tests"
    )
    out_path.write_text(
        json.dumps(
            result(
                task,
                status,
                message,
                artifacts=artifacts,
                diagnostics=[klee_run.stderr],
            ),
            indent=2,
        )
        + "\n"
    )
    print(message)
    return 0 if klee_run.returncode == 0 or tests else 0


if __name__ == "__main__":
    sys.exit(main())