"""Smoke test: load a quantized checkpoint through the OBLITERATUS loader.

Verifies: detection, dequantization to BF16, no quantized tensors left in
the model, config stripped of quantization_config, and a short generation.
"""

import sys

import torch

from obliteratus.models import quant_dequant as qd
from obliteratus.models.loader import load_model

repo = sys.argv[1]

det = qd.detect_quant_scheme(repo)
print(f"[detection] {det.scheme.value} (inverse={det.scale_is_inverse}, group={det.group_size})")
assert det.scheme not in (qd.QuantScheme.NONE, qd.QuantScheme.UNSUPPORTED), det

handle = load_model(
    repo,
    task="causal_lm",
    device="auto",
    dtype="bfloat16",
    trust_remote_code=True,
)
model = handle.model
print(f"[load] ok — scheme tag: {getattr(model, '_obliteratus_dequantized_scheme', None)}")

bad = []
n_params = 0
for name, p in model.named_parameters():
    n_params += 1
    if p.dtype in qd.FP8_DTYPES or p.dtype == torch.uint8:
        bad.append((name, str(p.dtype)))
    if torch.isnan(p.data).any().item() if p.data.is_floating_point() else False:
        bad.append((name, "NaN"))
print(f"[check] {n_params} params, quantized/NaN leftovers: {bad[:10] or 'NONE'}")
assert not bad, bad

assert getattr(model.config, "quantization_config", None) is None, "quantization_config survived"

tok = handle.tokenizer
prompt = "The capital of France is"
try:
    inputs = tok(prompt, return_tensors="pt").to(model.device)
except Exception:
    # Omni processors: fall back to bare tokenizer encode
    ids = tok.encode(prompt, return_tensors="pt").to(model.device)
    inputs = {"input_ids": ids}
with torch.no_grad():
    out = model.generate(**inputs, max_new_tokens=16, do_sample=False)
text = tok.decode(out[0][-16:] if out.dim() > 1 else out[-16:], skip_special_tokens=True)
print(f"[generate] {text!r}")
print("SMOKE_LOAD_OK")