"""Check that DeepInfer's HF loader accepts a (tiny) V4.1 checkpoint, on the host CPU only.

Runs the same validation chain as DeepseekForCausalLM.load_model / model_loader.load_hf_checkpoint, but
without materializing or loading weights:
  load_hf_model_config -> prepare_hf_coverage -> V41Builder on meta -> cover_tied_parameters
  -> "every model parameter is covered" -> validate_hf_shapes (shapes + dtypes from the safetensors headers)

Usage (on the box):
    cd ~/di/DeepInfer && DEEPINFER_ACCELERATOR=neuron DEEPINFER_NEURON_DEVICE=cpu \
        ~/di/venv/bin/python examples/neuron/check_tiny_checkpoint.py ~/di/tiny_ckpt
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron")
os.environ.setdefault("DEEPINFER_NEURON_DEVICE", "cpu")
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", str(29000 + os.getpid() % 1000))


_VENDOR = ("deep_gemm", "deep_ep", "deep_select", "flash_mla", "tile_kernels", "tilelang", "flash_attn")


def _stub_missing_vendor_packages() -> None:
    """The pure-torch vendor packages are being written in parallel; building parameter containers on meta
    only needs them to be importable. Stub any that are still missing (attribute access returns a dummy)."""
    import importlib
    import importlib.abc
    import importlib.machinery
    import types

    class _Dummy:
        def __init__(self, name):
            self._name = name

        def __getattr__(self, item):
            return _Dummy(f"{self._name}.{item}")

        def __call__(self, *a, **k):
            raise RuntimeError(f"stubbed vendor function {self._name} called")

    class _StubModule(types.ModuleType):
        def __getattr__(self, item):
            if item.startswith("__"):
                raise AttributeError(item)
            return _Dummy(f"{self.__name__}.{item}")

    class _Finder(importlib.abc.MetaPathFinder, importlib.abc.Loader):
        def find_spec(self, fullname, path, target=None):
            if fullname.split(".")[0] in stubbed:
                return importlib.machinery.ModuleSpec(fullname, self, is_package=True)
            return None

        def create_module(self, spec):
            m = _StubModule(spec.name)
            m.__path__ = []
            return m

        def exec_module(self, module):
            pass

    stubbed = set()
    for name in _VENDOR:
        try:
            importlib.import_module(name)
        except Exception as exc:  # noqa: BLE001
            stubbed.add(name)
            print(f"[check] vendor package {name!r} not importable ({type(exc).__name__}: {exc}); stubbing it")
    if stubbed:
        sys.meta_path.insert(0, _Finder())


def main(hf_dir: str) -> None:
    import torch
    import torch.distributed as dist

    import deepinfer
    from deepinfer import PLATFORM  # noqa: F401  (installs transfer_to_neuron -> cpu)

    _stub_missing_vendor_packages()

    # transfer_to_neuron replaces torch.empty/zeros/... with Python wrappers. torch.utils._device caches the
    # set of factory functions it redirects inside `with torch.device("meta")` on first use, i.e. after the
    # patch, so it holds only the wrappers and the original functions they call are NOT redirected: the
    # "meta" build then allocates real tensors (and prepare_meta_fp8 asserts). Work around it here by also
    # registering the wrapped originals. (Bug to be fixed in deepinfer/neuron_ops/transfer_to_neuron.py.)
    import torch.utils._device as _tdev

    _ctors = set(_tdev._device_constructors())
    _ctors |= {f.__wrapped__ for f in list(_ctors) if hasattr(f, "__wrapped__")}
    _tdev._device_constructors = lambda: _ctors

    assert str(PLATFORM.device) == "cpu", f"refusing to run on {PLATFORM.device}; set DEEPINFER_NEURON_DEVICE=cpu"
    if not dist.is_initialized():
        dist.init_process_group("gloo", rank=0, world_size=1)
    from deepinfer.model_executor.parallel_utils.parallel_state import PG

    world = dist.group.WORLD
    for attr in ("dp", "cp", "ep", "ep_in_dp", "enp", "world", "local", "cpu_dp_group", "cpu_cp_group", "cpu_ep_in_dp"):
        setattr(PG, attr, world)
    PG.rank = 0

    from deepinfer.hugging_face import cover_tied_parameters, load_hf_model_config, model_params, prepare_hf_coverage, validate_hf_shapes
    from deepinfer.model_executor.model_loader import build_model_on_meta

    hf = Path(hf_dir).expanduser()
    model_cfg = load_hf_model_config(hf / "config.json")
    gpt = model_cfg.gpt
    # runtime fields normally filled by ModelConfig / finalize_model_config
    gpt.setdefault("runtime_max_model_len", 4096)
    gpt.setdefault("max_num_seqs", 4)
    gpt.setdefault("max_num_batched_tokens", 4096)
    gpt.moe.setdefault("n_redundant_experts", 0)
    if gpt.get("dflash_cfg") and gpt.dflash_cfg.get("moe"):
        gpt.dflash_cfg.moe.setdefault("n_redundant_experts", 0)

    coverage = prepare_hf_coverage(gpt, hf_dir=hf)
    model = build_model_on_meta(model_cfg)
    params = model_params(model)
    coverage = cover_tied_parameters(coverage, set(params), n_layer=gpt.num_hidden_layers)
    missing = sorted(set(params) - coverage.model_names())
    unused = sorted(set(coverage.model_names()) - set(params))
    print(f"model params: {len(params)}, covered names: {len(coverage.model_names())}")
    print(f"missing (model params without HF source): {len(missing)} {missing[:20]}")
    print(f"unused HF groups (no model param): {len(unused)} {unused[:20]}")
    validate_hf_shapes(params, coverage, hf_dir=hf)
    n_exp = {k: len(v) for k, v in coverage.expert_tensors.items()}
    print(f"expert groups: {len(n_exp)} (experts per group: {sorted(set(n_exp.values()))}); engram tables: {coverage.engram_tables}")
    for name in sorted(n for n in params if "engram_embedding.embedding.embedding" in n):
        p = params[name]
        print(f"  {name}: runtime rows {tuple(p.shape)} (row_start {getattr(p, 'row_start', None)}), head sizes {getattr(p, 'head_vocab_sizes', None)}")
    assert not missing, "DeepInfer would reject this checkpoint (uncovered parameters)"
    print("OK: DeepInfer HF loader accepts", hf)
    dist.destroy_process_group()


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "~/di/tiny_ckpt")