#!/usr/bin/env python3
"""Smoke-test a local or Hugging Face Recon2Reason Reasoning 4B repository."""

from __future__ import annotations

import argparse

import torch
from PIL import Image
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration


EXPECTED_PARAMETERS = 4_437_815_808


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("model", help="Local directory or Hugging Face repository ID")
    parser.add_argument(
        "--device",
        default="cpu",
        help="Device map target used for loading, for example cpu or cuda:0",
    )
    parser.add_argument(
        "--generate",
        action="store_true",
        help="Also run a one-token generation test with an in-memory image",
    )
    args = parser.parse_args()

    processor = AutoProcessor.from_pretrained(args.model)
    model = Qwen3VLForConditionalGeneration.from_pretrained(
        args.model,
        dtype="auto",
        device_map={"": args.device},
        low_cpu_mem_usage=True,
    ).eval()

    parameter_count = sum(parameter.numel() for parameter in model.parameters())
    if parameter_count != EXPECTED_PARAMETERS:
        raise SystemExit(
            f"Parameter mismatch: expected {EXPECTED_PARAMETERS:,}, found {parameter_count:,}"
        )

    expected_tokens = {
        "<|feat|>": 151669,
        "<|/feat|>": 151670,
        "<|pos|>": 151671,
        "<|/pos|>": 151672,
        "<|feat_pad|>": 151673,
    }
    for token, expected_id in expected_tokens.items():
        actual_id = processor.tokenizer.convert_tokens_to_ids(token)
        if actual_id != expected_id:
            raise SystemExit(
                f"Token mismatch for {token}: expected {expected_id}, found {actual_id}"
            )

    print(f"Processor: {processor.__class__.__name__}")
    print(f"Model: {model.__class__.__name__}")
    print(f"Parameters: {parameter_count:,}")
    print(f"Dtype: {next(model.parameters()).dtype}")
    print("Special tokens: OK")

    if args.generate:
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": Image.new("RGB", (224, 224), "white")},
                    {"type": "text", "text": "What color is this image?"},
                ],
            }
        ]
        inputs = processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        ).to(args.device)
        with torch.inference_mode():
            generated_ids = model.generate(
                **inputs,
                max_new_tokens=1,
                do_sample=False,
            )
        prompt_length = inputs["input_ids"].shape[1]
        generated_text = processor.batch_decode(
            generated_ids[:, prompt_length:],
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )[0]
        print(f"One-token multimodal generation: OK ({generated_text!r})")

    print("Release smoke test: PASSED")


if __name__ == "__main__":
    main()