"""Let a checkpoint ship a quantised embedding table.

vLLM's compressed-tensors integration can serve a quantised
``VocabParallelEmbedding`` - it has a dequant-on-lookup method that unpacks only
the gathered rows - but the Qwen3.5 model implementation constructs the layer
without passing ``quant_config``:

    self.embed_tokens = VocabParallelEmbedding(
        self.vocab_size,
        config.hidden_size,
    )

With no quant_config the layer registers a plain ``weight`` parameter, so a
checkpoint carrying ``weight_packed`` / ``weight_scale`` fails to load with
"There is no module or parameter named 'embed_tokens.weight_packed'".

The embedding table is the largest unquantised block left in this checkpoint at
2.368 GiB of BF16; INT8 group-128 reclaims 1.166 GiB at a measured 0.65 %
relative RMSE, which is the difference between the 1,010,000-token profile
fitting on a 32 GiB card and missing by a fifth of a gigabyte.

This module rebuilds the layer with the checkpoint's own quantization config
before any weight is loaded.  It is a no-op for checkpoints whose config does
not declare a scheme for the embedding, so the same runtime serves both.
"""

from __future__ import annotations

from typing import Any

import torch

REPAIR_ID = "ZENIT_QUANTIZED_EMBEDDING_V1"


def _find_embedding(root: torch.nn.Module) -> tuple[str, Any] | None:
    """Locate the input embedding: a VocabParallelEmbedding that is not a head."""
    from vllm.model_executor.layers.vocab_parallel_embedding import (
        ParallelLMHead,
        VocabParallelEmbedding,
    )

    for name, module in root.named_modules():
        if isinstance(module, VocabParallelEmbedding) and not isinstance(
            module, ParallelLMHead
        ):
            return name, module
    return None


def _checkpoint_quantises_embedding(quant_config: Any, layer_name: str, layer: Any) -> bool:
    """Does the checkpoint declare a scheme for this layer?"""
    getter = getattr(quant_config, "get_scheme_dict", None)
    if getter is None:
        return False
    for candidate in (layer_name, f"model.{layer_name}", layer_name.split(".")[-1]):
        try:
            scheme = getter(layer, layer_name=candidate)
        except Exception:  # noqa: BLE001 - a probe must not break startup
            continue
        if scheme and scheme.get("weights") is not None:
            return True
    return False


def install_quantized_embedding(model: torch.nn.Module, vllm_config: Any) -> bool:
    """Rebuild the input embedding with the checkpoint's quant config.

    Returns True when the layer was replaced.  Safe to call for every model:
    checkpoints without a declared embedding scheme are left untouched.
    """
    from vllm.logger import init_logger
    from vllm.model_executor.layers.vocab_parallel_embedding import (
        VocabParallelEmbedding,
    )

    # vLLM attaches its handler to the "vllm" logger and sets
    # propagate=False, so a logger named after this package would have
    # its records dropped: every INFO line below would be invisible to
    # the person running the model. Naming it under "vllm." puts it
    # where the configured handler can see it.
    logger = init_logger("vllm.lomonosov_zenit_altay.embedding")
    quant_config = getattr(vllm_config, "quant_config", None)
    if quant_config is None:
        return False

    found = _find_embedding(model)
    if found is None:
        return False
    name, layer = found

    if getattr(layer, "quant_method", None) is not None and not isinstance(
        getattr(layer, "quant_method", None),
        __import__(
            "vllm.model_executor.layers.vocab_parallel_embedding",
            fromlist=["UnquantizedEmbeddingMethod"],
        ).UnquantizedEmbeddingMethod,
    ):
        return False  # already quantised by someone else

    if not _checkpoint_quantises_embedding(quant_config, name, layer):
        return False

    parent_path, _, attr = name.rpartition(".")
    parent = model.get_submodule(parent_path) if parent_path else model

    num_embeddings = getattr(layer, "org_vocab_size", None) or layer.num_embeddings
    replacement = VocabParallelEmbedding(
        num_embeddings,
        layer.embedding_dim,
        params_dtype=getattr(layer, "params_dtype", None),
        quant_config=quant_config,
        prefix=name,
    )
    replacement.to(next(layer.parameters()).device)
    setattr(parent, attr, replacement)
    logger.info(
        "%s: rebuilt %s with the checkpoint's quantization config (%s)",
        REPAIR_ID,
        name,
        type(replacement.quant_method).__name__,
    )
    return True


__all__ = ["REPAIR_ID", "install_quantized_embedding"]