#!/usr/bin/env python3
"""
scripts/pack_contrastive_persona_vn.py

Packs the complete 190-pair Phase 4 Contrastive Persona + Real Visual Novel (Steins;Gate + VA-11 Hall-A)
dataset (`data/phase4_contrastive_persona_vn.jsonl`) WITHOUT truncating any user prompt or assistant reply:
- Every pair includes the FULL `<|im_start|>user\n...\n<|im_end|>\n<|im_start|>assistant\n<think>\n...\n</think>\n\n` prompt.
- Every pair includes the FULL `chosen_target` reply PLUS `<|im_end|>\n` (`[248046, 198]`).
- Preserves the 3-way stratified split (`train` 59.5% = 113, `dev` 21.1% = 40, `test` 19.5% = 37).
- Emits chunked token files (`<= 8000` tokens each for `--ctx 8192`) and a concatenated trace + metadata file
  `data/m5_val_prompts/phase4_contrastive_vn_meta.json` for exact 48-layer Strata Metal state dumping.
"""

import argparse
import json
from pathlib import Path
from tokenizers import Tokenizer


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--self-test", action="store_true")
    args = ap.parse_args()

    root = Path(__file__).resolve().parent.parent
    in_path = root / "data" / "phase4_contrastive_persona_vn.jsonl"
    out_dir = root / "data" / "m5_val_prompts"
    out_dir.mkdir(parents=True, exist_ok=True)

    tok = Tokenizer.from_pretrained("Qwen/Qwen3.8-Flash-Next")
    rows = [json.loads(ln) for ln in in_path.read_text(encoding="utf-8").splitlines() if ln.strip()]
    assert len(rows) == 190, f"Expected 190 rows, got {len(rows)}"

    all_ids = []
    chunks = [[]]
    meta_pairs = []
    max_chunk_tokens = 7800

    for r in rows:
        p_ids = tok.encode(r["prompt"], add_special_tokens=False).ids
        c_ids_body = tok.encode(r["chosen_target"], add_special_tokens=False).ids
        # Append explicit <|im_end|>\n (248046, 198) so every reply terminates cleanly
        c_ids = c_ids_body + [248046, 198]
        rej_ids = tok.encode(r["rejected_slop"], add_special_tokens=False).ids + [248046, 198]

        pair_len = len(p_ids) + len(c_ids)
        if len(chunks[-1]) + pair_len > max_chunk_tokens and chunks[-1]:
            chunks.append([])

        # Verify prompt contains <|im_start|>user and ends with </think>\n\n (248069, 271)
        assert 248045 in p_ids, f"Pair {r['id']} missing <|im_start|>!"
        assert p_ids[-2:] == [248069, 271], f"Pair {r['id']} prompt tail is {p_ids[-4:]}, expected [248069, 271] (</think>\\n\\n)"

        start_in_all = len(all_ids)
        all_ids.extend(p_ids)
        step0_pos = len(all_ids) - 1
        all_ids.extend(c_ids)

        chunks[-1].extend(p_ids)
        chunks[-1].extend(c_ids)

        meta_pairs.append({
            "id": r["id"],
            "category": r["category"],
            "split": r["split"],
            "is_voice_override": r["is_voice_override"],
            "prompt_start_pos": start_in_all,
            "step0_pos": step0_pos,
            "prompt_len": len(p_ids),
            "chosen_ids": c_ids,
            "rejected_opener_ids": rej_ids[:16],
            "rejected_closer_ids": rej_ids[-16:] if len(rej_ids) >= 16 else rej_ids,
        })

    # Append 1 trailing newline token to each chunk so --n-prompt = len(chunk) - 1 dumps all len(chunk) - 1 rows!
    chunk_files = []
    for idx, ch in enumerate(chunks):
        ch_with_tail = ch + [198]
        cpath = out_dir / f"phase4_contrastive_vn_chunk{idx}.txt"
        cpath.write_text(" ".join(str(x) for x in ch_with_tail) + "\n", encoding="utf-8")
        chunk_files.append({"chunk_index": idx, "path": str(cpath.relative_to(root)), "n_rows_dumped": len(ch)})

    trace_path = out_dir / "phase4_contrastive_vn_trace.txt"
    trace_path.write_text(" ".join(str(x) for x in (all_ids + [198])) + "\n", encoding="utf-8")

    splits = {}
    for m in meta_pairs:
        splits[m["split"]] = splits.get(m["split"], 0) + 1

    meta_obj = {
        "n_tokens": len(all_ids),
        "n_pairs": len(meta_pairs),
        "splits": splits,
        "chunks": chunk_files,
        "pairs": meta_pairs,
    }
    meta_path = out_dir / "phase4_contrastive_vn_meta.json"
    meta_path.write_text(json.dumps(meta_obj, indent=2) + "\n", encoding="utf-8")

    print(f"Packed {len(meta_pairs)} full pairs into {len(all_ids)} tokens across {len(chunks)} chunks: {splits}")
    if args.self_test:
        assert splits == {"train": 113, "dev": 40, "test": 37}, f"Unexpected split counts: {splits}"
        assert sum(c["n_rows_dumped"] for c in chunk_files) == len(all_ids)
        print("scripts/pack_contrastive_persona_vn.py self-test: PASS")


if __name__ == "__main__":
    main()