#!/usr/bin/env python3
"""
prepare_data.py: build the chat SFT corpus for MetaDiffusion-600M.

Datasets:
- HuggingFaceTB/smol-smoltalk      : 460K general instruction conversations
- OpenCoder-LLM/opc-sft-stage1[lang:python] + stage2 : code instruction data
- OpenMathInstruct-2               : math (CoT solutions + answers)

Pipeline per conversation:
1. Normalize row -> messages [{role, content}]
2. Format with the Qwen chat template (<|im_start|>...<|im_end|>)
3. Tokenize; mark assistant-content tokens + rainbow padding as "response"
4. Truncate to seq_len (drop samples whose prompt alone overflows)
5. Rainbow-pad (cyclic <|r1|>..<|r7|>) so the model never sees repeated <eos>

Output: data/ids.bin (uint32 token ids; vocab is 151677, does NOT fit uint16),
data/resp.bin (uint8: 1 = assistant content / rainbow region), data/meta.json.

train.py masks ONLY the resp==1 region by default (prompt stays clean), which
matches the proven ChatDataset convention.

Usage:
    python prepare_data.py --out data --seq-len 512
    python prepare_data.py --datasets smol --max-samples 50000 --out data/smol-small
"""

import argparse
import itertools
import json
import logging
import random
import re
from pathlib import Path

import numpy as np
from transformers import AutoTokenizer

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

RAINBOW_TOKENS = [f"<|r{i}|>" for i in range(1, 8)]
NUM_NEW_TOKENS = 1 + len(RAINBOW_TOKENS)  # [MASK] + rainbow

# Per-dataset default sample caps (None = everything)
DEFAULT_CAPS = {"smol": None, "opc": 150_000, "math": 100_000, "no_robots": None}


def ensure_special_tokens(tokenizer):
    """Add [MASK] + rainbow tokens if missing.

    Works for both a fresh tokenizer (adds all 8 at the end) and the
    already-extended one saved by convert.py: [MASK] is always the first of
    the 8 appended tokens, so mask_id == len(tokenizer) - 8."""
    added = []
    if tokenizer.convert_tokens_to_ids("[MASK]") == tokenizer.unk_token_id:
        added.append("[MASK]")
    missing_rainbow = [t for t in RAINBOW_TOKENS
                       if tokenizer.convert_tokens_to_ids(t) == tokenizer.unk_token_id]
    if missing_rainbow:
        added.extend(missing_rainbow)
    if added:
        tokenizer.add_special_tokens({"additional_special_tokens": added})
    mask_id = tokenizer.convert_tokens_to_ids("[MASK]")
    expected = len(tokenizer) - NUM_NEW_TOKENS
    assert mask_id == expected, f"[MASK] at {mask_id}, expected {expected}"
    return tokenizer


def get_messages(row) -> list | None:
    """Normalize a dataset row into [{role, content}] or None."""
    if isinstance(row, dict):
        for key in ("messages", "conversations", "conversation"):
            val = row.get(key)
            if isinstance(val, list) and val:
                msgs = []
                for m in val:
                    role = str(m.get("role", "")).lower()
                    if role in ("human", "prompt", "user"):
                        role = "user"
                    elif role in ("gpt", "assistant", "response", "bot", "output"):
                        role = "assistant"
                    if role not in ("user", "assistant", "system"):
                        continue
                    content = m.get("content", m.get("value", ""))
                    if isinstance(content, list):
                        content = " ".join(str(c.get("text", c)) for c in content)
                    if content:
                        msgs.append({"role": role, "content": str(content).strip()})
                if msgs and any(m["role"] == "assistant" for m in msgs):
                    return msgs
        # Instruction/output shapes
        inst = row.get("instruction") or row.get("prompt") or row.get("question") or row.get("problem")
        out = row.get("output") or row.get("response") or row.get("answer") or row.get("solution")
        if inst and out:
            return [{"role": "user", "content": str(inst).strip()},
                    {"role": "assistant", "content": str(out).strip()}]
    return None


def math_messages(row) -> list | None:
    """OpenMathInstruct-2 shape: problem + generated_solution + expected_answer."""
    if not isinstance(row, dict):
        return None
    problem = row.get("problem") or row.get("question")
    solution = row.get("generated_solution") or row.get("solution")
    if not problem or not solution:
        return None
    answer = row.get("expected_answer")
    if answer and str(answer).strip():
        solution = f"{solution}\n\nFinal answer: {answer}"
    return [{"role": "user", "content": str(problem).strip()},
            {"role": "assistant", "content": str(solution).strip()}]


_THINK_RE = re.compile(r"<think>.*?</think>", re.S)

PLAIN_CHAT_TEMPLATE = (
    "{% for message in messages %}"
    "{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>\\n' }}"
    "{% endfor %}"
)


def is_junk_content(content):
    if "\uFFFD" in content:
        return True
    if not content:
        return False
    n_ascii = sum(1 for c in content if ord(c) < 128)
    return (len(content) - n_ascii) / len(content) > 0.25


def strip_think(messages):
    out = []
    for m in messages:
        content = m["content"]
        if m["role"] == "assistant" and isinstance(content, str):
            content = _THINK_RE.sub("", content)
            content = re.sub(r"\n{3,}", "\n\n", content).strip()
        out.append({"role": m["role"], "content": content})
    return out


def format_conversation(tokenizer, messages, seq_len, min_response_tokens):
    """Tokenize a conversation. Returns (ids, resp_flags) truncated/padded to
    seq_len, or None if the prompt alone cannot fit.

    The assistant span INCLUDES the trailing <|im_end|> token: a masked
    diffusion model can only learn to emit the terminator if it is a masked
    training target (VoidPadding format: [prompt][response][im_end][pad]*).
    """
    messages = strip_think(messages)
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
    enc = tokenizer(text, return_offsets_mapping=True, add_special_tokens=False)
    ids = enc["input_ids"]
    offsets = enc["offset_mapping"]

    resp_flags = [0] * len(ids)
    pos = 0
    for m in messages:
        if m["role"] == "assistant":
            marker = "<|im_start|>assistant\n"
            idx = text.find(marker, pos)
            if idx == -1:
                return None
            content_start = idx + len(marker)
            content_end = content_start + len(m["content"])
            # extend the span through the <|im_end|> token (the terminator is
            # a first-class masked target; this is the im_end-hardening fix)
            im_end_at = text.find("<|im_end|>", content_end)
            if im_end_at != -1:
                content_end = im_end_at + len("<|im_end|>")
            for t_i, (cs, ce) in enumerate(offsets):
                if cs >= content_start and ce <= content_end and t_i < len(resp_flags):
                    resp_flags[t_i] = 1
            pos = content_end
    if not any(resp_flags):
        return None

    # Find the first response token; prompt must fit with room for a response
    resp_start = resp_flags.index(1)
    if resp_start > seq_len - min_response_tokens:
        return None

    ids = ids[:seq_len]
    resp_flags = resp_flags[:seq_len]

    # Rainbow pad the tail (all maskable, teaches the model to stop-and-pad)
    rainbow_ids = [tokenizer.convert_tokens_to_ids(t) for t in RAINBOW_TOKENS]
    n_pad = seq_len - len(ids)
    if n_pad > 0:
        ids = ids + [rainbow_ids[i % 7] for i in range(n_pad)]
        resp_flags = resp_flags + [1] * n_pad
    return ids, resp_flags


def load_rows(dataset_name, max_samples):
    from datasets import load_dataset

    if dataset_name == "smol":
        ds = load_dataset("HuggingFaceTB/smol-smoltalk", split="train")
    elif dataset_name == "opc":
        rows = []
        for repo in ("OpenCoder-LLM/opc-sft-stage1", "OpenCoder-LLM/opc-sft-stage2"):
            loaded = None
            for config in ("lang:python", "lang:generic", None):
                try:
                    loaded = load_dataset(repo, config, split="train") if config else \
                        load_dataset(repo, split="train")
                    logger.info(f"  loaded {repo} config={config}: {len(loaded)} rows")
                    break
                except Exception:
                    continue
            if loaded is not None:
                rows.append(loaded)
        if not rows:
            raise RuntimeError("Could not load any opc-sft config")
        ds = rows[0] if len(rows) == 1 else None
    elif dataset_name == "math":
        ds = load_dataset("nvidia/OpenMathInstruct-2", split="train",
                          streaming=True)
    elif dataset_name == "no_robots":
        ds = load_dataset("HuggingFaceH4/no_robots", split="train")
    else:
        raise ValueError(f"unknown dataset: {dataset_name}")

    if ds is None:
        # opc multi-repo path: chain them
        def gen():
            for r in rows:
                yield from r
        return gen()
    return ds


def main():
    parser = argparse.ArgumentParser(description="Build MetaDiffusion-600M chat SFT corpus")
    parser.add_argument("--out", default="data", help="Output dir (ids.bin, resp.bin, meta.json)")
    parser.add_argument("--tokenizer", default="data/tokenizer",
                        help="Tokenizer dir (from convert.py) or HF id")
    parser.add_argument("--seq-len", type=int, default=512)
    parser.add_argument("--datasets", default="smol,opc,math",
                        help="Comma list of: smol, opc, math, no_robots")
    parser.add_argument("--jsonl", default=None,
                        help="Local JSONL of {\"messages\": [...]} rows")
    parser.add_argument("--max-samples", type=int, default=0,
                        help="Per-dataset cap (0 = dataset default)")
    parser.add_argument("--math-repeat", type=int, default=1,
                        help="Process the math dataset N times")
    parser.add_argument("--filter-junk", action=argparse.BooleanOptionalAction, default=True,
                        help="Drop samples whose assistant content is >25%% non-ASCII. On by default; --no-filter-junk to keep them.")
    parser.add_argument("--length-balance", action="store_true",
                        help="Duplicate samples whose response-token count is in "
                             "[--length-balance-min, --length-balance-max] by "
                             "--length-balance-mult (targets the 60-160 gen budget).")
    parser.add_argument("--length-balance-min", type=int, default=60)
    parser.add_argument("--length-balance-max", type=int, default=160)
    parser.add_argument("--length-balance-mult", type=int, default=3)
    parser.add_argument("--min-response-tokens", type=int, default=8)
    parser.add_argument("--val-fraction", type=float, default=0.05,
                        help="Hold out this fraction as ids_val.bin/resp_val.bin for early stopping")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
    tokenizer = ensure_special_tokens(tokenizer)
    tokenizer.chat_template = PLAIN_CHAT_TEMPLATE
    logger.info(f"Tokenizer: {len(tokenizer)} tokens, [MASK]={tokenizer.convert_tokens_to_ids('[MASK]')}")

    rng = random.Random(args.seed)
    all_ids = []
    all_resp = []
    n_samples = 0

    def process_rows(ds_name, rows, cap):
        nonlocal n_samples
        ds_n = 0
        ds_skipped = 0
        for row in rows:
            if cap is not None and ds_n >= cap:
                break
            msgs = math_messages(row) if ds_name == "math" else get_messages(row)
            if msgs is None:
                ds_skipped += 1
                continue
            if args.filter_junk:
                ac = msgs[-1]["content"] if msgs and msgs[-1]["role"] == "assistant" else ""
                if isinstance(ac, str) and is_junk_content(ac):
                    ds_skipped += 1
                    continue
            result = format_conversation(tokenizer, msgs, args.seq_len, args.min_response_tokens)
            if result is None:
                ds_skipped += 1
                continue
            ids, resp = result
            copies = 1
            if args.length_balance:
                n_resp = int(sum(resp))
                if args.length_balance_min <= n_resp <= args.length_balance_max:
                    copies = args.length_balance_mult
            for _ in range(copies):
                all_ids.append(np.array(ids, dtype=np.uint32))
                all_resp.append(np.array(resp, dtype=np.uint8))
                ds_n += 1
            if ds_n % 20000 == 0:
                logger.info(f"  [{ds_name}] {ds_n:,} samples")
        logger.info(f"[{ds_name}] kept {ds_n:,}, skipped {ds_skipped:,}")
        n_samples += ds_n

    for ds_name in [d.strip() for d in args.datasets.split(",") if d.strip()]:
        cap = args.max_samples if args.max_samples > 0 else DEFAULT_CAPS.get(ds_name)
        logger.info(f"[{ds_name}] loading (cap={cap})...")
        rows = load_rows(ds_name, cap)
        repeats = args.math_repeat if ds_name == "math" else 1
        if repeats > 1:
            cache_path = out / f".math_cache{'_' + str(cap) if cap else ''}.jsonl"
            if cache_path.exists():
                logger.info(f"[math] using local cache {cache_path}")
                with open(cache_path) as f:
                    rows = [json.loads(line) for line in f if line.strip()]
            else:
                rows = list(itertools.islice(rows, cap) if cap else rows)
                with open(cache_path, "w") as f:
                    for r in rows:
                        f.write(json.dumps(r) + "\n")
                logger.info(f"[math] cached {len(rows):,} rows to {cache_path}")
        for rep in range(repeats):
            if repeats > 1:
                logger.info(f"[{ds_name}] pass {rep + 1}/{repeats}")
            process_rows(ds_name, rows, cap)

    if args.jsonl:
        with open(args.jsonl) as f:
            rows = [json.loads(line) for line in f if line.strip()]
        logger.info(f"[jsonl] {len(rows)} rows from {args.jsonl}")
        cap = args.max_samples if args.max_samples > 0 else None
        process_rows("jsonl", rows, cap)

    if n_samples == 0:
        raise SystemExit("no samples kept; check --datasets / --jsonl / filters")
    logger.info(f"Shuffling {n_samples:,} samples (seed {args.seed})...")
    order = list(range(n_samples))
    rng.shuffle(order)
    n_val = int(n_samples * args.val_fraction)
    train_order = order[n_val:]
    val_order = order[:n_val]
    if not train_order:
        raise SystemExit(f"val-fraction {args.val_fraction} left zero train samples")
    ids_flat = np.concatenate([all_ids[i] for i in train_order])
    resp_flat = np.concatenate([all_resp[i] for i in train_order])

    ids_flat.tofile(out / "ids.bin")
    resp_flat.tofile(out / "resp.bin")

    if n_val > 0:
        ids_val = np.concatenate([all_ids[i] for i in val_order])
        resp_val = np.concatenate([all_resp[i] for i in val_order])
        ids_val.tofile(out / "ids_val.bin")
        resp_val.tofile(out / "resp_val.bin")
        logger.info(f"Val held out: {n_val:,} samples -> {out / 'ids_val.bin'} / {out / 'resp_val.bin'}")

    meta = {
        "seq_len": args.seq_len,
        "n_samples": n_samples - n_val,
        "n_val_samples": n_val,
        "val_held_out": True,
        "n_tokens": int(ids_flat.shape[0]),
        "mask_token_id": tokenizer.convert_tokens_to_ids("[MASK]"),
        "rainbow_token_ids": [tokenizer.convert_tokens_to_ids(t) for t in RAINBOW_TOKENS],
        "vocab_size": len(tokenizer),
        "datasets": args.datasets + (f",jsonl:{args.jsonl}" if args.jsonl else ""),
        "tokenizer_dir": str(args.tokenizer),
        "seed": args.seed,
        "filter_junk": args.filter_junk,
        "length_balance": args.length_balance,
    }
    with open(out / "meta.json", "w") as f:
        json.dump(meta, f, indent=2)
    logger.info(f"Wrote {out/'ids.bin'} ({ids_flat.nbytes/1e9:.2f} GB, {meta['n_tokens']:,} tokens), "
                f"{out/'resp.bin'}, {out/'meta.json'}")


if __name__ == "__main__":
    main()