"""vLLM entry point for the model-owned ALTAY architecture."""

from __future__ import annotations


ARCHITECTURE_ID = "LomonosovZenitAltayForConditionalGeneration"
RENDERER_MODE = "lomonosov_zenit_altay"


def _install_model_owned_renderer_default() -> None:
    """Select the model renderer automatically for this architecture only."""

    from vllm.config import ModelConfig

    if getattr(ModelConfig, "_zenit_noosphere_renderer_patch", False):
        return
    original = ModelConfig.__post_init__

    def patched(model_config, *args, **kwargs):
        original(model_config, *args, **kwargs)
        if (
            getattr(model_config, "architecture", None) == ARCHITECTURE_ID
            and model_config.tokenizer_mode == "auto"
        ):
            model_config.tokenizer_mode = RENDERER_MODE

    ModelConfig.__post_init__ = patched
    ModelConfig._zenit_noosphere_renderer_patch = True


def _register_uint3_with_zero_point(logger) -> bool:
    """Допустить несимметричные три бита: две строки, без своего ядра.

    Поздние 32 слоя MLP этого чекпоинта лежат в `pack-quantized` int3 с нулевой
    точкой. Считать их умеет уже установленный `humming_kernels`: сверка ядра с
    эталоном fp32 на процессоре даёт относительную ошибку 0.00207 при трёх битах
    с нулевой точкой — столько же, сколько без неё. Не хватает только УЧЁТА ТИПА:
    в vLLM 0.25.1 `scalar_types.uint3` не объявлен, а карта разрядностей с
    нулевой точкой знает лишь 4 и 8.

    Почему здесь, а не в обёртке вокруг запуска. vLLM v1 поднимает `EngineCore`
    отдельным процессом, который импортирует vllm заново; правка из родителя до
    него не доезжает, и загрузка падает так, будто типа нет. `register()` же
    вызывается механизмом плагинов в КАЖДОМ процессе — это единственное место,
    где правка видна и родителю, и воркеру.

    Замерено, что перевод поздних слоёв на три бита освобождает 0.90 ГиБ при
    удержании, неотличимом от отгружаемой модели (95/144 против 95/144 вместе с
    четырёхбитными эмбеддингами, p = 1.0000).
    """
    try:
        from vllm.scalar_type import ScalarType, scalar_types
        from vllm.model_executor.layers.quantization.compressed_tensors.schemes \
            import compressed_tensors_wNa16 as wna16
    except Exception:  # noqa: BLE001 - регистрация не должна ломать запуск
        return False
    if getattr(scalar_types, "uint3", None) is None:
        setattr(scalar_types, "uint3", ScalarType.uint(3, None))
    table = getattr(wna16, "WNA16_ZP_SUPPORTED_TYPES_MAP", None)
    if table is None or 3 in table:
        return bool(table and 3 in table)
    table[3] = scalar_types.uint3
    logger.info("ZENIT_W3_ASYM_V1: допущены три бита с нулевой точкой "
                "(разрядности: %s)", sorted(table))
    return True


def register() -> None:
    """Register lazily so vLLM can still use its native quantized kernels."""

    from vllm import ModelRegistry
    from vllm.renderers.registry import RENDERER_REGISTRY
    from vllm.tokenizers.registry import TokenizerRegistry

    from vllm.logger import init_logger

    from .jit_guard import install_build_parallelism_guard
    from .offline_noosphere_path import (
        REPAIR_ID as OFFLINE_PATH_ID,
        install_offline_noosphere_path,
    )
    from .positional_extension import (
        describe as describe_positional_extension,
        install as install_positional_extension,
    )
    from .serving_profiles import (
        REPAIR_ID as SERVING_PROFILES_ID,
        install_model_owned_serving_defaults,
    )
    from .tq_tiled_prefill import (
        DEFAULT_TILE_TOKENS,
        REPAIR_ID as TILED_PREFILL_ID,
        install_tiled_continuation_prefill,
    )
    from .razor_attention import (
        REPAIR_ID as RAZOR_ID,
        install as install_razor_attention,
    )
    from .vllm_kv_layout import install_altay72_compact_kv_layout

    # 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")

    _register_uint3_with_zero_point(logger)

    build_limits = install_build_parallelism_guard()
    if build_limits:
        logger.info(
            "%s: capping kernel build parallelism to %s "
            "(the FP4 CUTLASS kernels are compiled once for your GPU, and the "
            "default parallelism costs more host memory than most machines "
            "have). Set MAX_JOBS yourself to override.",
            "ZENIT_BUILD_PARALLELISM_GUARD_V1",
            ", ".join(f"{key}={value}" for key, value in build_limits.items()),
        )

    # Announce what this checkpoint changes about the engine.  A user who
    # reports a problem should be able to read the active list off their own
    # log rather than take our word for it.
    installed = {
        "compact_kv_layout": install_altay72_compact_kv_layout(),
        SERVING_PROFILES_ID: install_model_owned_serving_defaults(),
        f"{TILED_PREFILL_ID}(tile={DEFAULT_TILE_TOKENS})": (
            install_tiled_continuation_prefill()
        ),
        OFFLINE_PATH_ID: install_offline_noosphere_path(),
        # Выключено, пока не задано окно: короткий кэш на слоях без поисковых
        # голов меняет поведение, и до замера удержания включать его нельзя.
        RAZOR_ID: install_razor_attention(),
        # Reports its own state because the hook is inert unless the position
        # plan is set, and "installed" alone would not say which.
        describe_positional_extension(): install_positional_extension(),
    }
    logger.info(
        "LOMONOSOV ZENIT ALTAY runtime active: %s",
        ", ".join(
            f"{name}={'installed' if state else 'already present'}"
            for name, state in installed.items()
        ),
    )
    ModelRegistry.register_model(
        ARCHITECTURE_ID,
        "custom_generate.vllm_model:LomonosovZenitAltayForConditionalGeneration",
    )
    TokenizerRegistry.register(
        RENDERER_MODE,
        "vllm.tokenizers.hf",
        "CachedHfTokenizer",
    )
    RENDERER_REGISTRY.register(
        RENDERER_MODE,
        "custom_generate.vllm_renderer",
        "LomonosovZenitAltayRenderer",
    )
    _install_model_owned_renderer_default()


__all__ = ["ARCHITECTURE_ID", "RENDERER_MODE", "register"]