#!/usr/bin/env python3 -I
"""Translate the affected function(s) of candidate items with Charon then Aeneas (Lean backend), and check
that the generated Lean compiles against the Aeneas library.

Advisory items: Charon `--start-from-if-exists` on the functions whose bodies changed between the vulnerable and
the patched crate version (difffuncs.py), so the unit of translation is "the changed function plus its in-crate
callees", not the whole crate (whole-crate translation of real crates fails on Aeneas limitations; see TOOLCHAIN.md).
Mutant items (--mutants): the one-function crate is translated whole.

Per-item output fields:
  charon_ok   exit 0 and an .llbc was written          aeneas_ok   exit 0, no `[Error]` lines, .lean written
  fn_present  every start function appears as a `def` in the Lean output (not dropped / opaque)
  lean_ok     the generated Lean file compiles (lean --root, Aeneas lib on LEAN_PATH); only attempted if aeneas_ok
  translated  charon_ok and aeneas_ok and fn_present and lean_ok
Run after `source env.sh`. Concurrency: --jobs workers x CARGO_BUILD_JOBS=2 (default 2x2 = 4 build jobs), nice -n 15.
"""
import argparse, json, os, re, shutil, subprocess, sys, time
from concurrent.futures import ThreadPoolExecutor
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import difffuncs

TC = os.environ.get("TC", "/apps/development/rlhf/toolchains")
CHARON = os.environ.get("CHARON", f"{TC}/src/aeneas/charon/bin/charon")
AENEAS = os.environ.get("AENEAS", f"{TC}/src/aeneas/bin/aeneas")
LEAN_DIR = f"{TC}/src/aeneas/backends/lean"
ANSI = re.compile(r"\x1b\[[0-9;?]*[a-zA-Z]")
_LP = None


def lean_path():
    global _LP
    if _LP is None:
        o = subprocess.run(["lake", "env", "printenv", "LEAN_PATH"], cwd=LEAN_DIR, capture_output=True, text=True).stdout.strip()
        _LP = ":".join(dict.fromkeys(o.split(":")))
    return _LP


def start_patterns(fns):
    pats = []
    for _rel, mod, name in fns:
        base = "crate::" + (mod + "::" if mod else "")
        pats += [base + name, base + "_::" + name]
    return list(dict.fromkeys(pats))


def lean_compile(out, leanfile, timeout):
    env = dict(os.environ, LEAN_PATH=lean_path() + ":" + out)
    mod = os.path.splitext(os.path.basename(leanfile))[0]
    p = subprocess.run(["nice", "-n", "15", "lean", f"--root={out}", "-o", os.path.join(out, mod + ".olean"), leanfile],
                       cwd=LEAN_DIR, env=env, capture_output=True, text=True, timeout=timeout)
    return p.returncode == 0, (p.stdout + p.stderr)[-300:]


def one(it, work, timeout):
    t0 = time.time()
    name = it.get("adv_id") or it["item_id"]
    d = it["pre_dir"]
    out = os.path.join(work, name); shutil.rmtree(out, ignore_errors=True); os.makedirs(out)
    res = dict(item=name, crate=it.get("crate"), charon_ok=False, aeneas_ok=False, fn_present=False, lean_ok=False, translated=False)
    cmd = ["nice", "-n", "15", CHARON, "cargo", "--preset=aeneas", "--dest-file", os.path.join(out, "crate.llbc")]
    names = []
    if "post_dir" in it:
        fns = difffuncs.changed_functions(d, it["post_dir"])
        res["changed_functions"] = [f"{m}::{n}" if m else n for _r, m, n in fns]
        if not fns:
            res["error"] = "no_changed_function_identified"; res["seconds"] = round(time.time() - t0, 1); return res
        for p in start_patterns(fns): cmd += ["--start-from-if-exists", p]
        names = sorted({n for _r, _m, n in fns})
    env = dict(os.environ, CARGO_TARGET_DIR=os.path.join(out, "target"), CARGO_BUILD_JOBS="2")
    res["mode"] = "start_from_changed_fns" if names else "whole_crate"
    try:
        p = subprocess.run(cmd, cwd=d, env=env, capture_output=True, text=True, timeout=timeout)
        log = ANSI.sub("", p.stdout + p.stderr)
        res["charon_ok"] = p.returncode == 0 and os.path.exists(os.path.join(out, "crate.llbc"))
        if not res["charon_ok"]:
            errs = [l for l in log.splitlines() if re.match(r"\s*(error|Error)", l)]
            res["error"] = ("charon: " + (errs[0] if errs else log[-300:]))[:300]
    except subprocess.TimeoutExpired:
        res["error"] = "charon timeout"
    if res["charon_ok"]:
        try:
            p = subprocess.run(["nice", "-n", "15", AENEAS, "-backend", "lean", "-no-progress-bar", os.path.join(out, "crate.llbc"), "-dest", out],
                               capture_output=True, text=True, timeout=timeout)
            log = ANSI.sub("", p.stdout + p.stderr)
            leans = [f for f in os.listdir(out) if f.endswith(".lean")]
            nerr = len(re.findall(r"\[\s*Error\s*\]", log))
            res["n_aeneas_errors"] = nerr
            res["aeneas_ok"] = p.returncode == 0 and nerr == 0 and bool(leans)
            if not res["aeneas_ok"]:
                errs = [l for l in log.splitlines() if "Error" in l]
                res["error"] = ("aeneas: " + (errs[0] if errs else log[-200:]))[:300]
            elif leans:
                src = open(os.path.join(out, leans[0]), encoding="utf8").read()
                res["fn_present"] = all(re.search(rf"^(?:@\[[^\]]*\]\s*)*def\s+(?:[\w.]*\.)?{re.escape(n)}(?:_loop)?\b", src, re.M) for n in names) if names else True
                res["lean_file"] = leans[0]; res["lean_lines"] = src.count("\n")
                if names and not res["fn_present"] and res["mode"] == "start_from_changed_fns":
                    # patterns matched nothing (trait-impl methods, cfg'd-out code): retry from the whole crate
                    res["mode"] = "whole_crate_fallback"
                    cmd2 = ["nice", "-n", "15", CHARON, "cargo", "--preset=aeneas", "--dest-file", os.path.join(out, "crate.llbc")]
                    subprocess.run(cmd2, cwd=d, env=env, capture_output=True, text=True, timeout=timeout)
                    p = subprocess.run(["nice", "-n", "15", AENEAS, "-backend", "lean", "-no-progress-bar", os.path.join(out, "crate.llbc"), "-dest", out],
                                       capture_output=True, text=True, timeout=timeout)
                    log = ANSI.sub("", p.stdout + p.stderr)
                    res["n_aeneas_errors"] = len(re.findall(r"\[\s*Error\s*\]", log))
                    leans = [f for f in os.listdir(out) if f.endswith(".lean")]
                    if leans:
                        src = open(os.path.join(out, leans[0]), encoding="utf8").read()
                        res["fn_present"] = all(re.search(rf"^(?:@\[[^\]]*\]\s*)*def\s+(?:[\w.]*\.)?{re.escape(n)}(?:_loop)?\b", src, re.M) for n in names)
                        res["lean_file"] = leans[0]; res["lean_lines"] = src.count("\n")
                        # failures elsewhere in the crate are tolerated here: Aeneas drops failing decls, and what we need is the target fn
                        res["aeneas_ok"] = True
                if res["fn_present"]:
                    ok, msg = lean_compile(out, os.path.join(out, leans[0]), timeout)
                    res["lean_ok"] = ok
                    if not ok: res["error"] = "lean: " + msg
                else:
                    res["error"] = "start function not in generated Lean"
        except subprocess.TimeoutExpired:
            res["error"] = "aeneas/lean timeout"
    res["translated"] = res["charon_ok"] and res["aeneas_ok"] and res["fn_present"] and res["lean_ok"]
    shutil.rmtree(os.path.join(out, "target"), ignore_errors=True)  # keep llbc + lean, drop build products
    res["seconds"] = round(time.time() - t0, 1)
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("items"); ap.add_argument("--out", required=True)
    ap.add_argument("--work", default="/apps/development/rlhf/bugfind_data/work/translate")
    ap.add_argument("--jobs", type=int, default=2); ap.add_argument("--timeout", type=int, default=600)
    ap.add_argument("--mutants", action="store_true")
    ap.add_argument("--stages", default="diff_safe,strict_safe"); ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--max-changed-fns", type=int, default=3, help="advisory items touching more functions are not 'targeted fixes'")
    a = ap.parse_args()
    for p in (a.out, a.work):
        if not os.path.abspath(p).startswith("/apps/"): sys.exit("paths must be under /apps")
    items = [json.loads(l) for l in open(a.items)]
    if a.mutants:  # materialise each kept mutant as a one-file crate (`lib.rs` = the mutated function)
        items = [i for i in items if i.get("status") == "kept_nonequivalent"]
        for i in items:
            d = os.path.join(a.work, "_crates", i["item_id"]); os.makedirs(d + "/src", exist_ok=True)
            open(d + "/Cargo.toml", "w").write(f'[package]\nname = "m_{i["function"]}"\nversion = "0.1.0"\nedition = "2021"\n\n[workspace]\n')
            open(d + "/src/lib.rs", "w").write(i["mutant_source"] + "\n")
            i["pre_dir"], i["stage"] = d, "diff_safe"
    items = [i for i in items if "pre_dir" in i and i.get("stage") in a.stages.split(",")]
    pre_skip = []
    if not a.mutants:
        keep = []
        for i in items:
            n = len(difffuncs.changed_functions(i["pre_dir"], i["post_dir"]))
            i["n_changed_fns"] = n
            (keep if 1 <= n <= a.max_changed_fns else pre_skip).append(i)
        items = keep
        json.dump({"skipped_not_targeted_or_no_fn": len(pre_skip), "to_attempt": len(items)}, open(a.out + ".prefilter.json", "w"))
    items.sort(key=lambda i: (i.get("stage") != "strict_safe", i.get("n_changed_fns", 0)))
    if a.limit: items = items[: a.limit]
    done = {}
    if os.path.exists(a.out):
        for l in open(a.out): r = json.loads(l); done[r["item"]] = r
    todo = [i for i in items if (i.get("adv_id") or i["item_id"]) not in done]
    os.makedirs(a.work, exist_ok=True)
    with open(a.out, "a") as f, ThreadPoolExecutor(a.jobs) as ex:
        for r in ex.map(lambda i: one(i, a.work, a.timeout), todo):
            f.write(json.dumps(r) + "\n"); f.flush()
            print(r["item"], "T" if r["translated"] else "-", r.get("error", "")[:80], r["seconds"], flush=True)

if __name__ == "__main__":
    main()