"""Evaluation-only shims so stock vLLM 0.31 can load NVIDIA's Qwen3.8-Flash-Next-FP4-Local layout:
  - `<proj>.weight_scale.nvfp4_csf_fixed` + `..._exceptions` (b12x lossless CSF code) -> `<proj>.weight_scale` (E4M3);
  - NVFP4 PLE n-gram table shards (`shard_k.weight` u8 + `shard_k.weight_scale` E4M3 + one `weight_scale_2`) ->
    BF16 `shard_k.weight` (stock vLLM keeps an unquantized PLE table in pinned host memory)."""
import importlib.abc
import importlib.util
import os
import sys

TARGETS = ("vllm.model_executor.model_loader.weight_utils", "vllm.model_executor.model_loader.default_loader",
           "vllm.v1.worker.gpu.spec_decode.mtp.speculator", "vllm.model_executor.models.deepseek_v4.nvidia.model")
DRAFT_HOME = os.environ.get("Q38_DRAFT_HOME", "/root/q38prune/draft")


def _csf_decode(fixed, exceptions):
    import torch
    rows = fixed.shape[0] * 16
    cols = (fixed.shape[1] // 16 - 1) * 2
    slab = fixed.view(rows // 16, 16 * (1 + cols // 2))
    base = slab[:, :16].reshape(rows, 1).to(torch.int32)
    codes = slab[:, 16:].reshape(rows, cols // 2)
    codes = torch.stack([codes & 0xF, codes >> 4], dim=-1).reshape(rows, cols).to(torch.int32)
    out = (base + codes).view(-1)
    if exceptions.numel():
        ex = exceptions.to(torch.int64) & 0xFFFFFFFF
        out[ex & 0xFFFFFF] = (ex >> 24).to(torch.int32)
    return out.to(torch.uint8).view(rows, cols).view(torch.float8_e4m3fn)


_E2M1 = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]


def _ple_global(files):
    from safetensors import safe_open
    for f in files:
        with safe_open(f, "pt") as h:
            for n in h.keys():
                if n.endswith("ngram_embedding.weight_scale_2"):
                    return float(h.get_tensor(n).float())
    return None


def _ple_dequant(w, s, g):
    import torch
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    w, s = w.to(dev), s.to(dev).view(torch.float8_e4m3fn).float()
    lut = torch.tensor(_E2M1, device=dev)
    x = torch.stack([lut[(w & 0xF).long()], lut[(w >> 4).long()]], -1).flatten(-2)
    x = x.view(*x.shape[:-1], -1, 16) * s.unsqueeze(-1) * g
    return x.flatten(-2).to(torch.bfloat16).cpu()


def _wrap(it_fn):
    def wrapped(*a, **k):
        pending = {}
        ple = {}
        files = a[0] if a and isinstance(a[0], (list, tuple)) else k.get("hf_weights_files", [])
        g = _ple_global(files)
        for name, t in it_fn(*a, **k):
            if g is not None and ".ngram_embedding." in name and (".shard_" in name or name.endswith("weight_scale_2")):
                if name.endswith("weight_scale_2"):
                    continue
                base, leaf = name.rsplit(".", 1)
                if leaf == "weight_scale":
                    ple.setdefault(base, {})["s"] = t.clone()
                elif leaf == "weight":
                    ple.setdefault(base, {})["w"] = t.clone()
                else:
                    yield name, t
                    continue
                d = ple[base]
                if len(d) == 2:
                    del ple[base]
                    yield base + ".weight", _ple_dequant(d["w"], d["s"], g)
                continue
            for suf in (".nvfp4_csf_fixed", ".nvfp4_csf_exceptions"):
                if name.endswith(suf):
                    base = name[: -len(suf)]
                    d = pending.setdefault(base, {})
                    d[suf] = t.clone() if hasattr(t, "clone") else t
                    if len(d) == 2:
                        del pending[base]
                        yield base, _csf_decode(d[".nvfp4_csf_fixed"].cpu(), d[".nvfp4_csf_exceptions"].cpu())
                    break
            else:
                yield name, t
        assert not pending, f"unpaired CSF tensors: {list(pending)[:4]}"
        assert not ple, f"unpaired PLE shards: {list(ple)[:4]}"
    wrapped._csf = True
    return wrapped


def _patch(mod):
    for attr in ("safetensors_weights_iterator", "multi_thread_safetensors_weights_iterator"):
        fn = getattr(mod, attr, None)
        if fn is not None and not getattr(fn, "_csf", False):
            setattr(mod, attr, _wrap(fn))


class _Finder(importlib.abc.MetaPathFinder):
    def find_spec(self, fullname, path, target=None):
        if fullname not in TARGETS:
            return None
        sys.meta_path.remove(self)
        try:
            spec = importlib.util.find_spec(fullname)
        finally:
            sys.meta_path.insert(0, self)
        if spec is None or spec.loader is None:
            return spec
        orig = spec.loader.exec_module

        def exec_module(module):
            orig(module)
            if fullname.endswith("mtp.speculator"):
                if os.environ.get("Q38_DRAFT_VOCAB"):          # cheap MTP drafter (draft/q38_draft.py)
                    if DRAFT_HOME not in sys.path:
                        sys.path.insert(0, DRAFT_HOME)
                    import q38_draft
                    q38_draft.install_from_env()
            elif fullname.endswith("deepseek_v4.nvidia.model"):
                if os.environ.get("DS_STATS_DIR"):             # routed-expert usage statistics
                    sys.path.insert(0, "/root/ds4prune")
                    import ds4_stats_patch
                    ds4_stats_patch.install()
            else:
                _patch(module)
        spec.loader.exec_module = exec_module
        return spec


sys.meta_path.insert(0, _Finder())