"""Random tiny-model CPU test of custom generate/prefill. NOT a 4B accuracy test."""
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from grounding_demo.native_runtime import inspect_contract, user_messages


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model_path", type=Path, help="Reviewed native code/config/tokenizer directory")
    args = parser.parse_args()
    root = str(args.model_path.resolve())
    inspect_contract(Path(root), require_weights=False)
    import torch
    from PIL import Image
    from transformers import AutoConfig, AutoProcessor
    from transformers.dynamic_module_utils import get_class_from_dynamic_module

    c = AutoConfig.from_pretrained(root, trust_remote_code=True)
    for key, value in dict(hidden_size=32, intermediate_size=64, num_hidden_layers=1,
                           num_attention_heads=4, num_key_value_heads=2, head_dim=8).items():
        setattr(c.text_config, key, value)
    for key, value in dict(hidden_size=32, qkv_hidden_size=32, intermediate_size=64,
                           num_hidden_layers=1, num_attention_heads=4, out_hidden_size=32,
                           projector_hidden_size=64, init_pos_emb_height=2,
                           init_pos_emb_width=2).items():
        setattr(c.vision_config, key, value)
    for config in [c, c.text_config, c.vision_config]:
        config._attn_implementation = "sdpa"
    model_class = get_class_from_dynamic_module(c.auto_map["AutoModelForImageTextToText"], root)
    # Full float32 only for this CPU/random-weight structural test; production is BF16.
    model = model_class(c).eval().float()
    processor = AutoProcessor.from_pretrained(root, trust_remote_code=True)
    rendered = processor.apply_chat_template(user_messages("Find the red square."),
                                             tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[rendered], images=[Image.new("RGB", (28, 28), "red")],
                       return_tensors="pt")
    calls = []
    hook = model.model.visual.register_forward_pre_hook(lambda *args: calls.append("vision"))
    try:
        with torch.inference_mode():
            output = model.generate(**inputs, max_new_tokens=3, do_sample=False,
                                    use_cache=True, eos_token_id=[151645, 151643])
    finally:
        hook.remove()
    generated = int(output.shape[-1] - inputs["input_ids"].shape[-1])
    if len(calls) != 1 or generated < 1:
        raise RuntimeError("Vision must execute exactly once during cached generate prefill.")
    print(json.dumps({"scope": "RANDOM_TINY_MODEL_CPU_STRUCTURE_ONLY_NOT_4B_ACCURACY",
                      "vision_calls": len(calls), "new_tokens": generated,
                      "parameters": sum(p.numel() for p in model.parameters())}, indent=2))


if __name__ == "__main__":
    main()