# MTP draft-model quant workaround for vLLM 0.24.0 (mounted at PYTHONPATH). Canonical copy +
# rationale: eshpfi services/heretic2-nvfp4-quant/sitecustomize-mtp-workaround.py + the runbook.
# vLLM 0.24.0 doesn't propagate modelopt exclude_modules to the spec-decode DRAFT model -> the
# BF16 mtp head gets quantized -> shape crash. Force-skip mtp.* in is_layer_skipped to keep it BF16.
import importlib.abc
import importlib.util
import sys

TARGET = "vllm.model_executor.layers.quantization.utils.quant_utils"


class _Finder(importlib.abc.MetaPathFinder):
    def find_spec(self, name, path, target=None):
        if name != TARGET:
            return None
        sys.meta_path.remove(self)
        try:
            spec = importlib.util.find_spec(name)
        finally:
            sys.meta_path.insert(0, self)
        if not spec or not spec.loader:
            return None
        _orig_exec = spec.loader.exec_module

        def exec_module(module):
            _orig_exec(module)
            _orig_isls = module.is_layer_skipped

            def is_layer_skipped(prefix, ignored_layers, *args, **kwargs):
                pl = str(prefix)
                if pl.startswith("mtp.") or ".mtp." in pl:
                    return True
                return _orig_isls(prefix, ignored_layers, *args, **kwargs)

            module.is_layer_skipped = is_layer_skipped
            print("[mtp-workaround] is_layer_skipped force-skip for mtp.* installed", flush=True)

        spec.loader.exec_module = exec_module
        return spec


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