#!/usr/bin/env python3
"""Phase 0 smoke checks for Qwen/Qwen3.6-27B.

Default mode avoids downloading full weights. Use --mode load only on the GPU
host after confirming disk and VRAM are sufficient.
"""

from __future__ import annotations

import argparse
import json
import sys
from typing import Any


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="Qwen/Qwen3.6-27B")
    parser.add_argument("--mode", choices=["config", "load"], default="config")
    parser.add_argument("--dtype", choices=["auto", "bfloat16", "float16", "float32"], default="bfloat16")
    parser.add_argument("--device-map", default="auto")
    parser.add_argument("--trust-remote-code", action="store_true", default=True)
    parser.add_argument("--no-trust-remote-code", dest="trust_remote_code", action="store_false")
    return parser.parse_args()


def print_json(payload: dict[str, Any]) -> None:
    print(json.dumps(payload, indent=2, sort_keys=True, default=str))


def dtype_from_name(name: str):
    if name == "auto":
        return "auto"
    import torch

    return {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }[name]


def main() -> int:
    args = parse_args()

    try:
        import transformers
        from transformers import AutoConfig, AutoTokenizer
    except Exception as exc:
        print(f"Failed to import transformers: {exc!r}", file=sys.stderr)
        return 1

    config = AutoConfig.from_pretrained(args.model, trust_remote_code=args.trust_remote_code)
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=args.trust_remote_code)
    chat_template = getattr(tokenizer, "chat_template", "") or ""

    summary: dict[str, Any] = {
        "model": args.model,
        "transformers_version": transformers.__version__,
        "architectures": getattr(config, "architectures", None),
        "model_type": getattr(config, "model_type", None),
        "torch_dtype": str(getattr(config, "torch_dtype", None)),
        "vocab_size": getattr(config, "vocab_size", None),
        "eos_token": tokenizer.eos_token,
        "pad_token": tokenizer.pad_token,
        "has_think_template": "<think>" in chat_template and "</think>" in chat_template,
        "has_tool_call_template": "<tool_call>" in chat_template,
        "chat_template_chars": len(chat_template),
    }

    if args.mode == "config":
        print_json(summary)
        return 0

    import torch

    candidate_class_names = [
        "AutoModelForMultimodalLM",
        "AutoModelForImageTextToText",
        "AutoModelForVision2Seq",
        "AutoModelForCausalLM",
    ]

    errors: list[str] = []
    model = None
    loaded_with = None
    for class_name in candidate_class_names:
        model_cls = getattr(transformers, class_name, None)
        if model_cls is None:
            errors.append(f"{class_name}: not present in transformers {transformers.__version__}")
            continue
        try:
            model = model_cls.from_pretrained(
                args.model,
                torch_dtype=dtype_from_name(args.dtype),
                device_map=args.device_map,
                trust_remote_code=args.trust_remote_code,
            )
            loaded_with = class_name
            break
        except Exception as exc:
            errors.append(f"{class_name}: {exc!r}")

    if model is None:
        summary["load_errors"] = errors
        print_json(summary)
        return 2

    messages = [
        {"role": "system", "content": "You are validating an authorized security research training environment."},
        {"role": "user", "content": "Return one sentence confirming that the model can format a thinking response."},
    ]
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=True,
    )
    inputs = tokenizer(prompt, return_tensors="pt")
    device = next(model.parameters()).device
    inputs = {key: value.to(device) for key, value in inputs.items()}

    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=96, do_sample=False)

    decoded = tokenizer.decode(output_ids[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=False)

    summary.update(
        {
            "loaded_with": loaded_with,
            "device": str(device),
            "cuda_available": torch.cuda.is_available(),
            "cuda_device_count": torch.cuda.device_count(),
            "generated_chars": len(decoded),
            "generated_preview": decoded[:500],
        }
    )
    print_json(summary)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())