#!/usr/bin/env python3 -I
"""Merge selection + translation outputs into the committed item records and an honest funnel summary.
Inputs (under /apps data dir): advisory_candidates.jsonl, advisory_translate.jsonl, mutants.jsonl, mutant_translate.jsonl
Outputs (committed, under bugfind/bench/items/): advisory_items.jsonl (kept), advisory_attempts.jsonl (all translation attempts),
mutant_items.jsonl (kept), SUMMARY.json.
Item semantics: every item is a (buggy, reference) pair. Advisory: pre-fix crate version (positive) vs first patched version
(negative control). Mutant: mutated function (positive, with a differential-execution witness) vs the original (negative control).
"""
import json, os, sys, collections
HERE = os.path.dirname(os.path.abspath(__file__))
DATA = "/apps/development/rlhf/bugfind_data/out"
OUT = os.path.join(HERE, "items"); os.makedirs(OUT, exist_ok=True)


def jl(p):
    return [json.loads(l) for l in open(p)] if os.path.exists(p) else []

adv = {r["adv_id"]: r for r in jl(f"{DATA}/advisory_candidates.jsonl")}
tr = {r["item"]: r for r in jl(f"{DATA}/advisory_translate.jsonl")}
funnel = json.load(open(f"{DATA}/advisory_funnel.json"))
pre = json.load(open(f"{DATA}/advisory_translate.jsonl.prefilter.json")) if os.path.exists(f"{DATA}/advisory_translate.jsonl.prefilter.json") else {}
keep_fields = ["adv_id", "crate", "date", "url", "title", "categories", "affected_functions", "bug_class", "pre_version", "post_version", "stage", "changed_files", "changed_lines", "unsafe_in_changed_files"]
kept, attempts = [], []
for k, t in tr.items():
    a = adv[k]
    rec = {f: a.get(f) for f in keep_fields}
    rec.update({f: t.get(f) for f in ["changed_functions", "charon_ok", "aeneas_ok", "fn_present", "lean_ok", "translated", "error", "lean_lines", "seconds"]})
    rec["item_id"] = f"adv-{k}"; rec["kind"] = "advisory"
    rec["labels"] = {"buggy": f"{a['crate']}@{a['pre_version']}", "reference": f"{a['crate']}@{a['post_version']}"}
    attempts.append(rec)
    if t.get("translated"): kept.append(rec)
with open(f"{OUT}/advisory_items.jsonl", "w") as f:
    for r in sorted(kept, key=lambda r: r["item_id"]): f.write(json.dumps(r) + "\n")
with open(f"{OUT}/advisory_attempts.jsonl", "w") as f:
    for r in sorted(attempts, key=lambda r: r["item_id"]): f.write(json.dumps(r) + "\n")

mut = jl(f"{DATA}/mutants.jsonl"); mtr = {r["item"]: r for r in jl(f"{DATA}/mutant_translate.jsonl")}
mk = []
for m in mut:
    if m["status"] != "kept_nonequivalent": continue
    t = mtr.get(m["item_id"])
    m = dict(m); m["translated"] = bool(t and t["translated"]); m["translate_error"] = (t or {}).get("error")
    if m["translated"]: mk.append(m)
with open(f"{OUT}/mutant_items.jsonl", "w") as f:
    for m in mk: f.write(json.dumps(m) + "\n")

fail = collections.Counter((r.get("error") or "ok").split(":")[0] + ":" + (r.get("error") or "").split(":")[-1].strip()[:50] for r in attempts if not r["translated"])
mstat = collections.Counter(m["status"] for m in mut)
# --- kernels (minimal-extraction pairs from RustSec fixes) ---
ker_tr = {r["item"]: r for r in jl(f"{DATA}/kernel_translate.jsonl")}
ker_fail = jl(f"{DATA}/kernel_translate.jsonl.failed_attempts")
wit = {r["item_id"]: r for r in jl(os.path.join(HERE, "ledger", "witness_ledger.jsonl"))}
kdir = os.path.join(HERE, "kernels"); kk, katt = [], []
for kid in sorted(os.listdir(kdir)):
    m = json.load(open(f"{kdir}/{kid}/meta.json")); a, b = ker_tr.get(kid + ".pre"), ker_tr.get(kid + ".post"); w = wit.get(m["item_id"], {})
    rec = dict(m, kind="kernel", pre_translated=bool(a and a["translated"]), post_translated=bool(b and b["translated"]),
               pre_lean_sha256=(a or {}).get("lean_sha256"), post_lean_sha256=(b or {}).get("lean_sha256"),
               witness_pre=w.get("pre", {}).get("outcome"), witness_post=w.get("post", {}).get("outcome"), discriminating=bool(w.get("discriminating")))
    rec["kept"] = rec["pre_translated"] and rec["post_translated"] and rec["discriminating"]
    katt.append(rec)
    if rec["kept"]: kk.append(rec)
with open(f"{OUT}/kernel_items.jsonl", "w") as f:
    for r in kk: f.write(json.dumps(r) + "\n")
with open(f"{OUT}/kernel_attempts.jsonl", "w") as f:
    for r in katt: f.write(json.dumps(r) + "\n")

S = {
  "advisories": {**{k: v for k, v in funnel.items()}, "prefilter": pre,
                 "translation_attempted": len(attempts), "translated_and_lean_compiles": len(kept),
                 "distinct_crates_kept": len({r["crate"] for r in kept}), "kept_by_bug_class": dict(collections.Counter(c for r in kept for c in r["bug_class"])),
                 "translation_failure_reasons_top": fail.most_common(12)},
  "mutants": {"generated": len(mut), "by_status": dict(mstat), "translation_attempted": len(mtr),
              "kept_nonequivalent_and_translated": len(mk), "distinct_functions": len({m["function"] for m in mk}),
              "by_operator": dict(collections.Counter(m["operator"] for m in mk))},
  "kernels": {"attempted_pairs": len(katt), "pre_and_post_lean_compile": sum(r["pre_translated"] and r["post_translated"] for r in katt),
              "witness_discriminating_on_original_rust": sum(r["discriminating"] for r in katt), "kept_pairs": len(kk),
              "distinct_crates": len({r["crate"] for r in kk}), "by_bug_class": dict(collections.Counter(r["bug_class"] for r in kk)),
              "by_fidelity": dict(collections.Counter(r["fidelity"] for r in kk)),
              "translation_failures_before_rewrite": [{"item": r["item"], "error": (r.get("error") or "")[:160]} for r in ker_fail],
              "note": "Aeneas (6c2788d) emitted ill-typed Lean for `!b` / `a || b` chains bound to a bool (Prop vs Bool, `failed to synthesize Decidable`); the post kernel of tar_pax was rewritten with explicit if/else (same semantics) and then translates."},
  "total_items_kept_now": len(kept) + len(mk) + len(kk),
  "total_real_bug_pairs": len(kept) + len(kk),
  "notes": "Mutants are clustered by source function (21 functions) and are not independent samples; the crate-clustered bootstrap treats the function as the cluster. Advisory items are clustered by crate.",
}
json.dump(S, open(f"{OUT}/SUMMARY.json", "w"), indent=1)
print(json.dumps(S, indent=1))