# /// script
# dependencies = [
#     "transformers>=5.14.0",
#     "datasets>=4.0",
#     "Pillow>=10.0",
#     "torch>=2.5",
#     "torchvision>=0.20",
#     "accelerate>=1.0",
#     "num2words",
# ]
# ///
"""CPU pre-flight for the Gemma 4 training collator.

Runs the EXACT collate path from train_sft.py on a few real dataset rows, using
the processor + model config (no 16GB weights), then decodes one example to
prove which tokens are supervised vs masked. Catches chat-template, token-id,
and shape bugs on cpu-basic (~$0.01) before we burn GPU minutes.

    python gemma4/validate_collator.py --dataset khalidFlex/gui-agent-smoke
"""

import argparse

import torch
from datasets import load_dataset
from transformers import AutoConfig, AutoProcessor


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="google/gemma-4-E4B-it")
    ap.add_argument("--dataset", default="khalidFlex/gui-agent-smoke")
    ap.add_argument("--n", type=int, default=4)
    args = ap.parse_args()

    print(f"[val] processor + config for {args.model}")
    processor = AutoProcessor.from_pretrained(args.model)
    config = AutoConfig.from_pretrained(args.model)

    image_token_ids = {
        tid for tid in [
            getattr(config, "image_token_id", None),
            getattr(config, "boi_token_id", None),
            getattr(config, "eoi_token_id", None),
        ] if isinstance(tid, int)
    }
    print(f"[val] image/structural token ids: {sorted(image_token_ids)}")
    processor.tokenizer.padding_side = "right"

    ds = load_dataset(args.dataset, split=f"train[:{args.n}]")
    print(f"[val] loaded {len(ds)} rows")

    def to_messages(ex):
        return [
            {"role": "user", "content": [
                {"type": "image"},
                {"type": "text", "text": f"{ex['system']}\n\n{ex['user']}"},
            ]},
            {"role": "assistant", "content": [{"type": "text", "text": ex["assistant"]}]},
        ]

    full_texts, prompt_texts, images = [], [], []
    for ex in ds:
        msgs = to_messages(ex)
        full_texts.append(processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False))
        prompt_texts.append(processor.apply_chat_template(msgs[:-1], tokenize=False, add_generation_prompt=True))
        images.append([ex["image"].convert("RGB")])

    batch = processor(text=full_texts, images=images, return_tensors="pt", padding=True)
    prompt_batch = processor(text=prompt_texts, images=images, return_tensors="pt", padding=True)
    prompt_lens = prompt_batch["attention_mask"].sum(dim=1)

    labels = batch["input_ids"].clone()
    labels[batch["attention_mask"] == 0] = -100
    for tid in image_token_ids:
        labels[labels == tid] = -100
    for i, plen in enumerate(prompt_lens):
        labels[i, : int(plen)] = -100

    print("\n[val] === batch shapes ===")
    for k, v in batch.items():
        if isinstance(v, torch.Tensor):
            print(f"  {k}: {tuple(v.shape)} {v.dtype}")

    # Prove supervision: decode only the non-masked (label != -100) tokens of row 0.
    row0 = labels[0]
    supervised = batch["input_ids"][0][row0 != -100]
    n_sup = int((row0 != -100).sum())
    n_tot = int((batch["attention_mask"][0] == 1).sum())
    print(f"\n[val] row0 supervised tokens: {n_sup}/{n_tot} real tokens")
    print(f"[val] supervised text (should be ONLY the assistant answer):")
    print("   ", repr(processor.tokenizer.decode(supervised)))

    assert n_sup > 0, "no supervised tokens — masking is wrong!"
    assert n_sup < n_tot, "everything supervised — prompt not masked!"
    decoded = processor.tokenizer.decode(supervised)
    assert "<code>" in decoded or "click" in decoded or "final_answer" in decoded, \
        "supervised span doesn't look like an action — offset is off"
    print("\n[val] ✓ collator OK — prompt+image masked, assistant supervised.")


if __name__ == "__main__":
    main()