"""RazorAttention послойно: короткое окно там, где нет поисковых голов.

Идея из arXiv:2407.15891 (RazorAttention) и arXiv:2410.10819 (DuoAttention):
полный кэш нужен не всем головам, а лишь поисковым; прочим достаточно локального
окна. Обычно это упирается в то, что политика задаётся на голову, а кэш
раскладывается на слой, и приходится резать один раунд внимания на два — ровно от
этого предостерегает ZigzagAttention (arXiv:2508.12407), предлагая собирать
головы одного рода в одном слое.

**Нам собирать нечего: они уже собраны.** Замер поисковых голов
(`RETRIEVAL_HEADS_DENSE.json`, критерий копирования из arXiv:2404.15574) дал 18
голов из 384, и они лежат в семи слоях полного внимания из шестнадцати. Девять
слоёв не содержат ни одной головы из верхней восемнадцатки:

    с поисковыми головами : 0, 3, 4, 11, 12, 13, 14
    без них               : 1, 2, 5, 6, 7, 8, 9, 10, 15

Арифметика: кэш KV занимает 13 ГБ и весь приходится на 16 слоёв полного внимания
(остальные 48 — Gated DeltaNet с состоянием фиксированного размера, кэша не
имеют). Это 0.81 ГБ на слой; короткое окно на девяти слоях освобождает до 7.3 ГБ,
почти четверть всей памяти модели с миллионом.

Что делает этот модуль. Перед тем как движок соберёт группы кэша, он подменяет
`FullAttentionSpec` на `SlidingWindowSpec` для перечисленных слоёв. Гибридные
спецификации vLLM 0.25.1 поддерживает штатно (`unify_hybrid_kv_cache_specs`),
новых механизмов не требуется.

**По умолчанию выключено.** Включается `ZENIT_RAZOR_WINDOW=<токенов>` или полем
профиля `razor_window`. Так сделано намеренно: пока удержание с короткими окнами
не замерено на 33 000 и 250 334, менять поведение по умолчанию нельзя — это
правило дня, градусник проверяется там, где ответ известен.
"""

from __future__ import annotations

import os
import re
from typing import Any

REPAIR_ID = "ZENIT_RAZOR_ATTENTION_V1"
ENV_WINDOW = "ZENIT_RAZOR_WINDOW"

# Замерено, не выбрано. Индексы — по слоям ПОЛНОГО ВНИМАНИЯ, от нуля, в порядке
# их следования в модели; остальные 48 слоёв рекуррентные и кэша не имеют.
RETRIEVAL_LAYERS: frozenset[int] = frozenset({0, 3, 4, 11, 12, 13, 14})

_LAYER_NUMBER = re.compile(r"\.layers\.(\d+)\.")


def _logger():
    from vllm.logger import init_logger

    return init_logger("vllm.lomonosov_zenit_altay")


def requested_window(profile: dict[str, Any] | None = None) -> int | None:
    """Сколько токенов оставить непоисковым слоям. None — не трогать ничего."""
    raw = os.environ.get(ENV_WINDOW)
    if raw is None and profile:
        raw = profile.get("razor_window")
    if raw in (None, "", "0"):
        return None
    try:
        window = int(raw)
    except (TypeError, ValueError):
        raise ValueError(
            f"{ENV_WINDOW}={raw!r} — ожидалось целое число токенов"
        ) from None
    if window <= 0:
        return None
    return window


def plan_layers(kv_cache_spec: dict[str, Any]) -> tuple[list[str], list[str]]:
    """Разложить слои полного внимания на поисковые и прочие.

    Возвращает `(оставить полными, сузить)`. Порядковый номер берётся по
    возрастанию номера слоя в имени — так же, как их нумеровал замер голов.
    """
    from vllm.v1.kv_cache_interface import FullAttentionSpec

    full = [
        name for name, spec in kv_cache_spec.items()
        if isinstance(spec, FullAttentionSpec)
    ]

    def order(name: str) -> int:
        found = _LAYER_NUMBER.search(name)
        return int(found.group(1)) if found else 1 << 30

    full.sort(key=order)
    keep, narrow = [], []
    for position, name in enumerate(full):
        (keep if position in RETRIEVAL_LAYERS else narrow).append(name)
    return keep, narrow


def _sliding_spec_class():
    """Скользящее окно, знающее про TurboQuant.

    Замерено: у модели страница 1 634 688 байт и у рекуррентных слоёв, и у
    полного внимания, потому что последнее описано классом
    `TQFullAttentionSpec` — он считает страницу как
    `block_size * num_kv_heads * tq_slot_size`, а не по сырой формуле. Стоковый
    `SlidingWindowSpec` про кодек не знает и даёт 8 454 144; эта страница
    становится максимумом, а 1 634 688 на неё не делится (5.17), и движок
    отказывается: дополнять страницу рекуррентных слоёв их бэкенд не умеет.

    Поэтому окну нужен тот же расчёт. Наследуемся от обоих: поведение окна от
    одного, размер страницы от другого.
    """
    from dataclasses import dataclass

    from vllm.v1.kv_cache_interface import SlidingWindowSpec, TQFullAttentionSpec

    @dataclass(frozen=True, kw_only=True)
    class TQSlidingWindowSpec(SlidingWindowSpec):
        tq_slot_size: int = 0

        @property
        def real_page_size_bytes(self) -> int:
            if self.tq_slot_size > 0:
                return self.block_size * self.num_kv_heads * self.tq_slot_size
            return super().real_page_size_bytes

    return TQSlidingWindowSpec


def apply(kv_cache_spec: dict[str, Any], window: int) -> dict[str, Any]:
    """Подменить спецификацию непоисковых слоёв на скользящее окно."""
    SlidingWindowSpec = _sliding_spec_class()

    keep, narrow = plan_layers(kv_cache_spec)
    if not narrow:
        return kv_cache_spec

    # Диагностика по живой спецификации, а не по классам: две догадки о причине
    # отказа уже оказались неверны, потому что строились на умолчаниях, а не на
    # том, что движок собрал на самом деле.
    if os.environ.get("ZENIT_RAZOR_DUMP"):
        seen: dict[tuple, list[str]] = {}
        for name, spec in kv_cache_spec.items():
            key = (type(spec).__name__, spec.page_size_bytes,
                   getattr(spec, "indexes_kv_by_block_stride", None),
                   str(getattr(spec, "kv_quant_mode", None)))
            seen.setdefault(key, []).append(name)
        log = _logger()
        for key, names in sorted(seen.items(), key=lambda kv: -kv[0][1]):
            log.info("%s DUMP: %s страница=%d block_stride=%s кодек=%s слоёв=%d",
                     REPAIR_ID, key[0], key[1], key[2], key[3], len(names))

    updated = dict(kv_cache_spec)
    for name in narrow:
        source = kv_cache_spec[name]
        fields = {
            "block_size": source.block_size,
            "num_kv_heads": source.num_kv_heads,
            "head_size": source.head_size,
            "dtype": source.dtype,
            "sliding_window": window,
        }
        # Поля, которых может не быть: переносим, если есть. tq_slot_size здесь
        # обязателен — без него страница считается по сырой формуле и ломает
        # согласование с рекуррентными слоями.
        for optional in ("kv_quant_mode", "page_size_padded",
                         "indexes_kv_by_block_stride", "head_size_v",
                         "tq_slot_size"):
            if hasattr(source, optional):
                fields[optional] = getattr(source, optional)
        try:
            updated[name] = SlidingWindowSpec(**fields)
        except TypeError:
            # Сигнатура разошлась — молча портить раскладку нельзя.
            raise ValueError(
                f"{REPAIR_ID}: SlidingWindowSpec не принимает поля {sorted(fields)}; "
                "раскладка кэша в этой версии vLLM изменилась, проверьте вручную"
            ) from None

    _logger().info(
        "%s: полный кэш оставлен %d слоям с поисковыми головами, %d слоя(ёв) "
        "переведены на окно %d токенов. Замер голов: 18 из 384 в семи слоях "
        "(RETRIEVAL_HEADS_DENSE.json)",
        REPAIR_ID, len(keep), len(narrow), window,
    )
    return updated


def install(profile: dict[str, Any] | None = None) -> bool:
    """Врезаться перед сборкой групп кэша. Возвращает True, если врезались."""
    window = requested_window(profile)
    if window is None:
        return False

    from vllm.v1.core import kv_cache_utils as core

    if getattr(core, "_zenit_razor_installed", False):
        return False

    # Раскладка кэша сверяет состав семей со своим ожиданием и отвергает любой
    # незнакомый — это её предохранитель, а не дефект. Мы делим шестнадцать
    # слоёв полного внимания на семь и девять, поэтому состав надо объявить,
    # иначе сверка не пройдёт, раскладка вернёт None, и движок уйдёт на путь,
    # где страницы рекуррентных слоёв согласовать некому.
    from . import vllm_kv_layout as layout

    full, recurrent = len(RETRIEVAL_LAYERS), None
    for count in layout.EXPECTED_FAMILY_COUNTS:
        if count != 16:
            recurrent = count
    if recurrent is None:
        _logger().warning(
            "%s: состав семей раскладки неожиданный (%s) — окно не применяется",
            REPAIR_ID, layout.EXPECTED_FAMILY_COUNTS)
        return False
    layout.expect_family_counts((full, 16 - full, recurrent))

    # Врезаемся в сборку ГРУПП, а не конфига: имя get_kv_cache_config в этой
    # версии отсутствует, и попытка пропатчить его даёт громкий AttributeError
    # при старте — что и случилось на первом прогоне.
    original = core.get_kv_cache_groups

    def patched(vllm_config: Any, kv_cache_spec: dict[str, Any], *args, **kwargs):
        try:
            kv_cache_spec = apply(kv_cache_spec, window)
        except ValueError:
            raise
        except Exception:  # noqa: BLE001 - раскладка не должна ронять запуск
            _logger().exception("%s: подмена спецификации не удалась", REPAIR_ID)
        return original(vllm_config, kv_cache_spec, *args, **kwargs)

    core.get_kv_cache_groups = patched
    core._zenit_razor_installed = True
    return True


__all__ = ["REPAIR_ID", "ENV_WINDOW", "RETRIEVAL_LAYERS",
           "requested_window", "plan_layers", "apply", "install"]