"""Check a chat JSONL file and split it into train.jsonl and validation.jsonl.

    python prepare_data.py example.jsonl

Every row needs id, category, group and messages, with the answer last.
Rows longer than --max-tokens are dropped. A group never spans both files.
"""
import argparse
import json
import random
from collections import Counter

from transformers import AutoTokenizer

MODEL = "Qwen/Qwen3.5-9B"
REVISION = "c202236235762e1c871ad0ccb60c8ee5ba337b9a"
FIELDS = ("id", "category", "group", "messages")


def read_rows(path):
    rows = []
    for n, line in enumerate(open(path), 1):
        if not line.strip():
            continue
        row = json.loads(line)
        missing = [k for k in FIELDS if k not in row]
        if missing:
            raise SystemExit(f"line {n}: missing {', '.join(missing)}")
        if row["messages"][-1]["role"] != "assistant":
            raise SystemExit(f"line {n}: the last message must be the answer")
        rows.append(row)
    return rows


def count_tokens(tokenizer, messages):
    return len(tokenizer.apply_chat_template(
        messages, tokenize=True, return_dict=False, enable_thinking=False))


def split_by_group(rows, share, seed):
    groups = sorted({r["group"] for r in rows})
    random.Random(seed).shuffle(groups)
    size = Counter(r["group"] for r in rows)
    held_out, n = set(), 0
    for group in groups:
        if n >= share * len(rows):
            break
        held_out.add(group)
        n += size[group]
    train = [r for r in rows if r["group"] not in held_out]
    val = [r for r in rows if r["group"] in held_out]
    return train, val


def write(path, rows):
    with open(path, "w") as f:
        f.writelines(json.dumps(r, ensure_ascii=False) + "\n" for r in rows)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("data")
    ap.add_argument("--max-tokens", type=int, default=24576)
    ap.add_argument("--validation-share", type=float, default=0.1)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(MODEL, revision=REVISION)
    rows, kept = read_rows(args.data), []
    for row in rows:
        n = count_tokens(tokenizer, row["messages"])
        if n > args.max_tokens:
            print(f"drop {row['id']}: {n} tokens")
        else:
            kept.append((n, row))

    lengths = sorted(n for n, _ in kept)
    print(f"rows={len(kept)} dropped={len(rows) - len(kept)} "
          f"median={lengths[len(lengths) // 2]} "
          f"max={lengths[-1]} total={sum(lengths)}")

    train, val = split_by_group([r for _, r in kept],
                                args.validation_share, args.seed)
    write("train.jsonl", train)
    write("validation.jsonl", val)
    for name, split in (("train", train), ("validation", val)):
        categories = dict(Counter(r["category"] for r in split))
        print(f"{name}: {len(split)} rows", categories)
    unseen = {r["category"] for r in train} - {r["category"] for r in val}
    if unseen:
        print("warning: no validation rows for", ", ".join(sorted(unseen)))


if __name__ == "__main__":
    main()