"""scripts/p2/check_qwen3_snapshot.py on a tiny fake HF snapshot (CPU, header-only checks).

Run: CUDA_VISIBLE_DEVICES= .venv/bin/python -m pytest tests/test_check_qwen3_snapshot.py -q
"""

from __future__ import annotations

import json
from pathlib import Path

import pytest
import torch
from safetensors.torch import save_file

from align.model.qwen3 import Qwen3Config
from align.model.qwen3_checkpoint import _expected_shapes
from test_p2_support import load_script

CK = load_script("p2_check_qwen3_snapshot", "scripts/p2/check_qwen3_snapshot.py")
TINY = dict(model_type="qwen3", architectures=["Qwen3ForCausalLM"], vocab_size=32, hidden_size=16, intermediate_size=24, num_hidden_layers=2,
            num_attention_heads=4, num_key_value_heads=2, head_dim=8, rms_norm_eps=1e-6, rope_theta=1e6, max_position_embeddings=64,
            hidden_act="silu", attention_bias=False, tie_word_embeddings=False, sliding_window=None, use_sliding_window=False)


def make(d: Path, drop=None, extra=None, bad_shape=None, cfg_over=None, truncate=False):
    d.mkdir(parents=True, exist_ok=True)
    cfg = {**TINY, **(cfg_over or {})}
    (d / "config.json").write_text(json.dumps(cfg))
    shapes = _expected_shapes(Qwen3Config.from_hf_config(TINY))
    state = {k: torch.zeros(s, dtype=torch.bfloat16) for k, s in shapes.items()}
    if drop:
        state.pop(drop)
    if extra:
        state[extra] = torch.zeros(2, dtype=torch.bfloat16)
    if bad_shape:
        state[bad_shape] = torch.zeros(3, 3, dtype=torch.bfloat16)
    keys = sorted(state)
    half = len(keys) // 2
    wm = {}
    for i, part in enumerate((keys[:half], keys[half:])):
        name = f"model-0000{i + 1}-of-00002.safetensors"
        save_file({k: state[k] for k in part}, str(d / name))
        wm.update({k: name for k in part})
    (d / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}, "weight_map": wm}))
    if truncate:
        p = d / "model-00002-of-00002.safetensors"
        p.write_bytes(p.read_bytes()[:-10])
    return d


def test_a_complete_qwen3_snapshot_is_compatible(tmp_workdir):
    rep = CK.check(str(make(Path(tmp_workdir) / "ok")))
    assert rep["compatible"] and rep["problems"] == [] and rep["dtypes"] == {"BF16": rep["index"]["tensors"]}


@pytest.mark.parametrize("kw,needle", [
    (dict(drop="model.layers.1.mlp.up_proj.weight"), "absent from the index"),
    (dict(extra="model.layers.0.self_attn.q_proj.bias"), "cannot map"),
    (dict(bad_shape="model.norm.weight"), "implied by the config"),
    (dict(truncate=True), "truncated"),
])
def test_each_defect_is_named(tmp_workdir, kw, needle):
    rep = CK.check(str(make(Path(tmp_workdir) / "bad", **kw)))
    assert not rep["compatible"] and any(needle in p for p in rep["problems"]), rep["problems"]


def test_unsupported_config_variants_are_refused(tmp_workdir):
    assert CK.main([str(make(Path(tmp_workdir) / "sw", cfg_over={"use_sliding_window": True, "sliding_window": 8}))]) == 1
    assert CK.main([str(make(Path(tmp_workdir) / "ok2"))]) == 0