#!/usr/bin/env python
"""CPU-only compatibility check of a HF snapshot against align's Qwen3 loader (no weights are loaded).

    .venv/bin/python scripts/p2/check_qwen3_snapshot.py SNAPSHOT_DIR [--json]

  1. config.json parses through `align.model.qwen3_checkpoint.read_qwen3_config` (it REFUSES every variant align does
     not implement: sliding window, scaled RoPE, biases, tied head, non-silu, ...);
  2. the safetensors index's `weight_map` keys equal EXACTLY the key set `hf_to_align` maps (no missing, no extra);
  3. every shard exists and its header-declared tensor extents fit the file (a truncated download fails here);
  4. every tensor's shape (read from the shard HEADER, 8 bytes + JSON, no tensor data) equals the shape the config implies;
  5. dtypes are reported (bf16 expected).
Exit 0 compatible, 1 not. Memory: a few MB. Threads: 1.
"""

from __future__ import annotations

import argparse
import json
import os
import struct
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from align.model.checkpoint import INDEX_NAME  # noqa: E402
from align.model.qwen3_checkpoint import _expected_shapes, _table, read_qwen3_config  # noqa: E402


def read_header(path: str) -> dict:
    """The safetensors header dict, plus `__file_ok__`: whether the declared extents fit the file."""
    size = os.path.getsize(path)
    with open(path, "rb") as f:
        (n,) = struct.unpack("<Q", f.read(8))
        if n > 100 * 1024 * 1024 or 8 + n > size:
            raise ValueError(f"{path}: implausible safetensors header length {n}")
        h = json.loads(f.read(n))
    end = max((v["data_offsets"][1] for k, v in h.items() if k != "__metadata__"), default=0)
    h["__file_ok__"] = 8 + n + end == size
    h["__declared_bytes__"] = 8 + n + end
    h["__file_bytes__"] = size
    return h


def check(snapshot: str) -> dict:
    snap = Path(snapshot)
    rep: dict = {"snapshot": str(snap), "problems": []}
    P = rep["problems"]
    cfg = read_qwen3_config(str(snap))
    rep["config"] = {k: getattr(cfg, k) for k in ("n_layers", "hidden_size", "intermediate_size", "n_heads", "n_kv_heads", "head_dim",
                                                  "vocab_size", "rope_theta", "rms_norm_eps", "max_position_embeddings", "tie_word_embeddings")}
    want = set(_table(cfg))
    shapes = _expected_shapes(cfg)
    index = snap / INDEX_NAME
    if not index.is_file():
        P.append(f"no {INDEX_NAME}")
        return rep
    wm = json.loads(index.read_text())["weight_map"]
    rep["index"] = {"tensors": len(wm), "shards": sorted(set(wm.values()))}
    missing, extra = sorted(want - set(wm)), sorted(set(wm) - want)
    if missing:
        P.append(f"{len(missing)} key(s) hf_to_align needs are absent from the index: {missing[:5]}")
    if extra:
        P.append(f"{len(extra)} index key(s) hf_to_align cannot map: {extra[:5]}")
    dtypes: dict = {}
    seen = 0
    for shard in rep["index"]["shards"]:
        sp = snap / shard
        if not sp.is_file():
            P.append(f"shard {shard} named by the index is absent")
            continue
        h = read_header(str(sp))
        if not h["__file_ok__"]:
            P.append(f"shard {shard}: header declares {h['__declared_bytes__']} bytes, file has {h['__file_bytes__']} (truncated or padded)")
        for k, v in h.items():
            if k.startswith("__"):
                continue
            seen += 1
            dtypes[v["dtype"]] = dtypes.get(v["dtype"], 0) + 1
            if wm.get(k) != shard:
                P.append(f"{k}: in {shard} but the index says {wm.get(k)}")
            if k in shapes and tuple(v["shape"]) != shapes[k]:
                P.append(f"{k}: shape {tuple(v['shape'])} != {shapes[k]} implied by the config")
    rep["tensors_seen_in_shard_headers"], rep["dtypes"] = seen, dtypes
    if seen != len(wm):
        P.append(f"shard headers carry {seen} tensors, the index lists {len(wm)}")
    rep["compatible"] = not P
    return rep


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("snapshot")
    ap.add_argument("--json", action="store_true")
    a = ap.parse_args(argv)
    try:
        rep = check(a.snapshot)
    except Exception as e:  # noqa: BLE001 -- a refusal from the config parser is the answer
        rep = {"snapshot": a.snapshot, "compatible": False, "problems": [f"{type(e).__name__}: {e}"]}
    print(json.dumps(rep, indent=1, default=str) if a.json else
          ("COMPATIBLE" if rep.get("compatible") else "NOT COMPATIBLE") + "\n" + "\n".join(f"  - {p}" for p in rep["problems"]))
    return 0 if rep.get("compatible") else 1


if __name__ == "__main__":
    sys.exit(main())