#!/usr/bin/env python3
"""Export one finished patient build into the per-patient files of the benchmark.

  event_chain.jsonl     <- event_chain.jsonl (patient_id set to the exported identifier)
  evidence_chain.jsonl  <- evidence_chain.jsonl (entity null for the facts of synthetic questions)
  dialogues.jsonl       <- dialogues.jsonl (meta keeps only atom_index)
  questions.jsonl       <- questions_axis{1..8}.jsonl + checkpoints.jsonl + validation.jsonl
                           (sound questions in axis-file order with their checkpoint; axis-3 questions with the
                           same options and answer kept once; axis 8 keeps at most AXIS8_PER_MECHANISM questions per
                           mechanism, evenly spaced by evidence-to-checkpoint distance, not-yet-due deadlines first)

  python -m pipeline.export --work build/<patient_id> --out <dir> --patient-id P<nn>
"""
from __future__ import annotations

import argparse
import json
import os
import re
from collections import Counter
from functools import lru_cache
from pathlib import Path

import tiktoken

COPIES = [("event_chain.jsonl", "event_chain.jsonl"), ("evidence_chain.jsonl", "evidence_chain.jsonl"),
          ("dialogues.jsonl", "dialogues.jsonl")]
AXIS8_PER_MECHANISM = {"deadline": 10, "event_bounded": 5, "fixed_duration": 5}
AXIS8_STATE_ORDER = {"deadline": ["not_yet_due", "overdue"]}


@lru_cache(maxsize=None)
def enc():
    """The tiktoken encoding named by MEMEBENCH_TOKENIZER, loaded on first use."""
    name = os.environ.get("MEMEBENCH_TOKENIZER", "")
    if not name or name.startswith("<"):
        raise SystemExit("set MEMEBENCH_TOKENIZER (the tiktoken encoding of the exported size fields)")
    return tiktoken.get_encoding(name)


def jl(path):
    return [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]


def check_build(w):
    need = [f"questions_axis{a}.jsonl" for a in range(1, 9)] + [
        "checkpoints.jsonl", "validation.jsonl", "MANIFEST.json", "visits.jsonl"] + [s for s, _ in COPIES]
    missing = [f for f in need if not (w / f).exists()]
    if missing:
        raise SystemExit(f"{w}: build incomplete, missing {missing}")
    empty = [r["visit_number"] for r in jl(w / "dialogues.jsonl") if not r.get("turns")]
    if empty:
        raise SystemExit(f"{w}: {len(empty)} dialogues have no turns (re-run step F): {empty[:20]}")
    n_visits = len(jl(w / "visits.jsonl"))
    n_dia = len(jl(w / "dialogues.jsonl"))
    if n_dia != n_visits:
        raise SystemExit(f"{w}: {n_dia} dialogues for {n_visits} visits")


def build_questions(patient, w):
    """The sound questions of the build in axis-file order, and the (mechanism, state) of each axis-8 question."""
    dia = sorted(jl(w / "dialogues.jsonl"), key=lambda r: r["visit_number"])
    tok = {r["visit_number"]: len(enc().encode(" ".join(f"{t['speaker']}: {t['text']}" for t in r["turns"]),
                                             disallowed_special=())) for r in dia}
    date_of = {r["visit_number"]: r["date"] for r in dia}
    ck = {c["qid"]: c for c in jl(w / "checkpoints.jsonl")}
    va = {v["qid"]: v for v in jl(w / "validation.jsonl")}

    out, axis8, dropped, seen_ax3 = [], {}, [], set()
    for ax in range(1, 9):
        for q in jl(w / f"questions_axis{ax}.jsonl"):
            qid = q["qid"]
            if qid not in va or qid not in ck:
                raise SystemExit(f"{qid}: no validation or checkpoint row")
            if not va[qid]["sound"]:
                continue
            if ax == 3:
                key = json.dumps([q.get("options"), q.get("answer")], ensure_ascii=False, sort_keys=True)
                if key in seen_ax3:
                    dropped.append((qid, "axis3_duplicate"))
                    continue
                seen_ax3.add(key)
            ev, ckv = q["evidence_visits"], ck[qid]["ckpt_visit"]
            if ckv not in tok:
                raise SystemExit(f"{qid}: checkpoint v{ckv} not in dialogues")
            if ck[qid]["ckpt_date"] and str(ck[qid]["ckpt_date"]) != date_of[ckv]:
                raise SystemExit(f"{qid}: checkpoint date {ck[qid]['ckpt_date']} differs from the dialogue date")
            out_qid = f"{patient}/{qid.split('/', 1)[1]}"
            if ax == 8:
                spec = q["axis_specific"]
                axis8[out_qid] = (spec["mechanism"], spec["state"])
            out.append({
                "qid": out_qid, "patient_id": patient, "axis": q["axis"], "question": q["question"],
                "options": q.get("options"), "answer": q["answer"], "answer_type": q["answer_type"],
                "explanation": q["provenance"] if ax in (1, 2) else " ".join(q["derivation"]["steps"]),
                "checkpoint_visit": ckv, "checkpoint_date": ck[qid]["ckpt_date"],
                "evidence_visits": ev, "grounding": q["grounding"], "n_evidence_visits": len(ev),
                "evidence_context_tokens": sum(tok.get(v, 0) for v in ev),
                "context_visits": sum(1 for v in tok if v <= ckv),
                "context_tokens": sum(t for v, t in tok.items() if v <= ckv),
            })
    return out, axis8, dropped


def spread(qs, k):
    """k questions evenly spaced over the list sorted by the evidence-to-checkpoint gap."""
    if k <= 0:
        return []
    if len(qs) <= k:
        return qs
    s = sorted(qs, key=lambda q: q["checkpoint_visit"] - q["evidence_visits"][0])
    idx = sorted(set(round(i * (len(s) - 1) / (k - 1)) for i in range(k))) if k > 1 else [0]
    return [s[i] for i in idx]


def select_axis8(rows, axis8):
    """At most AXIS8_PER_MECHANISM questions per mechanism, evenly spaced by evidence-to-checkpoint distance; for
    deadlines, not-yet-due questions first, then overdue ones."""
    keep = set()
    for mech, cap in AXIS8_PER_MECHANISM.items():
        taken = []
        for state in AXIS8_STATE_ORDER.get(mech, [None]):
            pool = [r for r in rows if r["qid"] in axis8 and axis8[r["qid"]][0] == mech
                    and state in (None, axis8[r["qid"]][1])]
            taken += spread(pool, cap - len(taken))
        keep.update(r["qid"] for r in taken)
    return [r for r in rows if r["qid"] not in axis8 or r["qid"] in keep]


def export(w, dest, patient):
    """Write the per-patient files of the build in w into the new directory dest, as patient `patient` (P<nn>)."""
    if not re.fullmatch(r"P\d{2}", str(patient or ""), re.ASCII):
        raise SystemExit(f"patient id {patient!r} is not of the form P<nn>")
    check_build(w)
    rows, axis8, dropped = build_questions(patient, w)
    unknown = sorted({m for m, _ in axis8.values()} - set(AXIS8_PER_MECHANISM))
    if unknown:
        raise SystemExit(f"axis-8 mechanisms without a cap: {unknown}")
    kept = select_axis8(rows, axis8)
    dest.mkdir(parents=True)
    for src, dst in COPIES:
        records = jl(w / src)
        for r in records:
            if src == "dialogues.jsonl":
                r["meta"] = {"atom_index": (r.get("meta") or {}).get("atom_index", {})}
            if src == "event_chain.jsonl" and "patient_id" in r:
                r["patient_id"] = patient
            if src == "evidence_chain.jsonl" and str(r.get("batch_id", "")).startswith("synthetic-"):
                r["entity"] = None
        with open(dest / dst, "w", encoding="utf-8") as fh:
            for r in records:
                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    with open(dest / "questions.jsonl", "w", encoding="utf-8") as fh:
        for r in kept:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    return rows, kept, axis8, dropped


def main():
    ap = argparse.ArgumentParser(description="Export one finished patient build.")
    ap.add_argument("--work", required=True, help="the patient's build directory (build/<patient_id>)")
    ap.add_argument("--out", required=True, help="output directory (must not exist)")
    ap.add_argument("--patient-id", required=True,
                    help="identifier written as patient_id and as the qid prefix, of the form P<nn>")
    a = ap.parse_args()
    dest = Path(a.out)
    if dest.exists():
        raise SystemExit(f"{dest} already exists; refusing to overwrite")
    rows, kept, axis8, dropped = export(Path(a.work), dest, a.patient_id)
    per_axis = Counter(r["axis"][0] for r in kept)
    print(f"-> {dest}")
    print(f"sound questions: {len(rows)}  axis-3 duplicates dropped: {len(dropped)}")
    print(f"axis 8: {len(axis8)} sound, kept per mechanism "
          f"{dict(Counter(axis8[r['qid']][0] for r in kept if r['qid'] in axis8))}")
    print(f"questions.jsonl: {len(kept)}  per axis {dict(sorted(per_axis.items()))}")


if __name__ == "__main__":
    main()