"""Architecture-owned repair for vLLM 0.25.1 TurboQuant hybrid cache shapes.

vLLM's hybrid cache planner allocates the correct TurboQuant page size, but
the reshape path currently represents TurboQuant with ``KVQuantMode.NONE``.
It consequently passes ``cache_dtype_str="auto"`` to the TurboQuant backend
instead of the explicit preset selected by the user.

The LOMONOSOV ZENIT model installs this narrowly scoped repair from its own
constructor.  No environment variable, external launcher, or standalone
service is required.  The repair changes shape selection only; it neither
changes tensor weights nor enables NOOSPHERE routing.
"""

from __future__ import annotations

from typing import Any


REPAIR_ID = "ZENIT_VLLM_0251_TURBOQUANT_HYBRID_SHAPE_V1"


def install_turboquant_hybrid_shape_repair(vllm_config: Any) -> bool:
    """Restore the requested TurboQuant preset when vLLM reports ``auto``."""

    cache_config = getattr(vllm_config, "cache_config", None)
    requested = str(getattr(cache_config, "cache_dtype", ""))
    if not requested.startswith("turboquant_"):
        return False

    from vllm.v1.attention.backends.turboquant_attn import (
        TurboQuantAttentionBackend,
    )

    backend = TurboQuantAttentionBackend
    backend._zenit_requested_turboquant_preset = requested
    if getattr(backend, "_zenit_model_owned_shape_repair", False):
        return True

    original = backend.get_kv_cache_shape

    def get_kv_cache_shape(
        num_blocks: int,
        block_size: int,
        num_kv_heads: int,
        head_size: int,
        cache_dtype_str: str = "turboquant_4bit_nc",
    ) -> tuple[int, ...]:
        if cache_dtype_str == "auto":
            cache_dtype_str = backend._zenit_requested_turboquant_preset
        return original(
            num_blocks,
            block_size,
            num_kv_heads,
            head_size,
            cache_dtype_str,
        )

    backend.get_kv_cache_shape = staticmethod(get_kv_cache_shape)
    backend._zenit_model_owned_shape_repair = True
    backend._zenit_model_owned_shape_repair_id = REPAIR_ID
    return True


__all__ = ["REPAIR_ID", "install_turboquant_hybrid_shape_repair"]