"""Short-prompt (q3s) copies of q3 Clef-labelled JSONLs for the Glint student.

Keeps id, images, state, gold, teacher, teacher_logits and meta (incl. box_pairs, so box_primary_button stays
paired with primary_button) exactly as they are -- no re-labelling. Only the question texts change:
  * choice / noul: one-line question (per record language), choice criteria reduced to the option ids
    (descriptions null, the compact encoder renders ids only); noul criteria dropped.
  * box: a short element description (box_primary_button = "the primary action button", same pairing).
Sets meta.question_version = "q3s" and record["prompt_style"] = "compact" (glint/schema_compact.py).

usage: python scripts/make_short.py in.q3.clef.jsonl [more.jsonl ...]   -> writes <name>.short.jsonl next to each
       (refuses to overwrite an existing output)
"""
import json, sys
from pathlib import Path

SHORT = {
    "screen": {"en": "What is the current screen state?", "ja": "\u73fe\u5728\u306e\u753b\u9762\u306e\u72b6\u614b\u306f\uff1f"},
    "has_error": {"en": "Is an error message visible?", "ja": "\u30a8\u30e9\u30fc\u8868\u793a\u304c\u3042\u308b\u304b\uff1f"},
    "primary_button": {"en": "Which is the primary action button?", "ja": "\u4e3b\u8981\u306a\u30dc\u30bf\u30f3\u306f\u3069\u308c\u304b\uff1f"},
    "is_loading": {"en": "Is something loading?", "ja": "\u8aad\u307f\u8fbc\u307f\u4e2d\u304b\uff1f"},
    "has_modal": {"en": "Is a modal or sheet presented?", "ja": "\u30e2\u30fc\u30c0\u30eb\u3084\u30b7\u30fc\u30c8\u304c\u8868\u793a\u4e2d\u304b\uff1f"},
    "has_destructive_action": {"en": "Is a destructive action visible?", "ja": "\u7834\u58ca\u7684\u306a\u64cd\u4f5c\u304c\u3042\u308b\u304b\uff1f"},
    "is_keyboard_visible": {"en": "Is the keyboard visible?", "ja": "\u30ad\u30fc\u30dc\u30fc\u30c9\u304c\u8868\u793a\u4e2d\u304b\uff1f"},
    "is_empty": {"en": "Is it an empty state?", "ja": "\u7a7a\u306e\u72b6\u614b\u304b\uff1f"},
    "has_tab_bar": {"en": "Is a tab bar visible?", "ja": "\u30bf\u30d6\u30d0\u30fc\u304c\u3042\u308b\u304b\uff1f"},
}
BOX_SHORT = {
    "box_tab_bar": "the tab bar",
    "box_selected_tab": "the selected tab",
    "box_primary_button": "the primary action button",
    "box_modal": "the modal or sheet",
    "box_focused_text_field": "the focused text field",
    "box_keyboard": "the keyboard",
    "box_loading_indicator": "the loading indicator",
    "box_error_message": "the error message",
    "box_empty_state": "the empty-state message",
    "box_destructive_button": "the destructive button",
}


def shorten(r):
    lang = "ja" if (r.get("meta") or {}).get("lang") == "ja" else "en"
    qs = {}
    for qid, q in r["questions"].items():
        t = q["type"]
        if t == "box":
            if qid not in BOX_SHORT:
                raise SystemExit(f"{r['id']}: no short text for box question {qid}")
            qs[qid] = {"type": "box", "instructions": BOX_SHORT[qid]}
        elif t == "choice":
            qs[qid] = {"type": "choice", "instructions": SHORT[qid][lang], "criteria": {k: None for k in q["criteria"]}}
        elif t == "noul":
            qs[qid] = {"type": "noul", "instructions": SHORT[qid][lang]}
        else:
            raise SystemExit(f"{r['id']}: unsupported type {t} for {qid}")
    out = dict(r)
    out["questions"] = qs
    out["meta"] = {**r["meta"], "question_version": "q3s", "question_version_src": r["meta"].get("question_version")}
    out["prompt_style"] = "compact"
    return out


def main(paths):
    for src in map(Path, paths):
        dst = src.with_name(src.name[: -len(".jsonl")] + ".short.jsonl")
        if dst.exists():
            raise SystemExit(f"refusing to overwrite {dst}")
        n = 0
        with open(src, encoding="utf-8") as f, open(dst, "x", encoding="utf-8", newline="\n") as out:
            for line in f:
                if line.strip():
                    out.write(json.dumps(shorten(json.loads(line)), ensure_ascii=False) + "\n"); n += 1
        print(f"{n} records {src} -> {dst}")


if __name__ == "__main__":
    main(sys.argv[1:])