#!/usr/bin/env python3
"""Join the hand-coded intervention rows with message metadata (read-only on
transcripts) and write the public dataset.

Inputs
  --codes    pipe-separated coding file (one row per intervention, or a NONE /
             EXCL row per message):
             msg_id|category|how|agent_reported_success|recurred|agent_did|noticed
             agent_did / noticed are PARAPHRASES written by the coder, never
             the owner's text.
  --second   second-pass codes on the consistency subsample: msg_id|CAT;CAT
  --msgs     private JSONL produced by mine_transcripts.py (never published)
  --sample   data/sample.csv

Outputs (all under --out)
  coding_pass1.csv     every sampled message: msg_id, stratum, date, project,
                       status (intervention/none/excluded), categories
  interventions.csv    one row per coded intervention (paraphrased fields only)
  consistency_check.csv  pass-1 vs pass-2 categories for the 60-message subsample

Derived flag `checks_cited`: 1 when the immediately preceding assistant text
cites passing checks (tests passing, analyze clean, green, verified, 'N/N').
It is computed by regex from the private text; only the flag is published.

Coding was done by an AI assistant (Claude) under the owner's direction.
"""
import argparse, collections, csv, json, os, re

CHECKS_RE = re.compile(r"(tests? (all )?pass|all \d+ tests|\d+/\d+|analy[sz]e (is )?clean|\bgreen\b|"
                       r"verified|passes|suite (is )?green|tsc .*clean|exit 0)", re.I)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--codes", required=True)
    ap.add_argument("--second", required=True)
    ap.add_argument("--msgs", required=True)
    ap.add_argument("--sample", required=True)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()

    msgs = {}
    by_sess = collections.defaultdict(collections.Counter)
    for line in open(a.msgs):
        r = json.loads(line)
        msgs[r["msg_id"]] = r
        if r["project_hint"] != "unknown":
            by_sess[r["session"]][r["project_hint"]] += 1
    sess_proj = {s: c.most_common(1)[0][0] for s, c in by_sess.items()}

    def project(r):
        return r["project_hint"] if r["project_hint"] != "unknown" else sess_proj.get(r["session"], "unknown")

    strata = {row["msg_id"]: row["stratum"] for row in csv.DictReader(open(a.sample))}

    rows = []
    for line in open(a.codes):
        line = line.strip()
        if not line:
            continue
        p = line.split("|")
        rows.append(p)

    per_msg = collections.OrderedDict()
    inter = []
    for p in rows:
        mid, cat = p[0], p[1]
        per_msg.setdefault(mid, [])
        if cat in ("NONE", "EXCL"):
            per_msg[mid].append(cat)
            continue
        per_msg[mid].append(cat)
        how, succ, rec = p[2], p[3], p[4]
        did = p[5] if len(p) > 5 else "-"
        saw = p[6] if len(p) > 6 else "-"
        r = msgs[mid]
        inter.append({
            "row_id": f"I{len(inter)+1:03d}", "msg_id": mid, "stratum": strata[mid],
            "date": r["ts"][:10], "project": project(r), "category": cat, "how_noticed": how,
            "agent_reported_success": succ, "checks_cited": int(bool(CHECKS_RE.search(r["prev"] or ""))),
            "recurred": rec, "agent_did_paraphrase": did if cat != "VIS" else "(visual: see sibling paper)",
            "noticed_paraphrase": saw if cat != "VIS" else "(visual: see sibling paper)",
        })

    os.makedirs(a.out, exist_ok=True)
    with open(os.path.join(a.out, "coding_pass1.csv"), "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["msg_id", "stratum", "date", "project", "status", "categories"])
        for mid in sorted(per_msg):
            cats = per_msg[mid]
            status = "excluded" if cats == ["EXCL"] else ("none" if cats == ["NONE"] else "intervention")
            r = msgs[mid]
            w.writerow([mid, strata[mid], r["ts"][:10], project(r), status,
                        ";".join(c for c in cats if c not in ("NONE", "EXCL"))])
    with open(os.path.join(a.out, "interventions.csv"), "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(inter[0].keys()))
        w.writeheader()
        w.writerows(inter)

    second = {}
    for line in open(a.second):
        line = line.strip()
        if line:
            mid, cats = line.split("|")
            second[mid] = cats
    with open(os.path.join(a.out, "consistency_check.csv"), "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["msg_id", "pass1", "pass2"])
        for mid in sorted(second):
            p1 = per_msg[mid]
            p1s = "EXCL" if p1 == ["EXCL"] else ("NONE" if p1 == ["NONE"] else ";".join(sorted(set(p1))))
            w.writerow([mid, p1s, ";".join(sorted(set(second[mid].split(";"))))])
    print(len(per_msg), "messages,", len(inter), "intervention rows")


if __name__ == "__main__":
    main()