# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from __future__ import annotations

from typing import Union

import torch

import vllm.envs as envs
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON

logger = init_logger(__name__)
_DISABLE_SNAPKV_TRITON: bool = False


def snapkv_query_aware_token_scores(
    *,
    query: torch.Tensor,  # [N_q, Hq, D]
    key: torch.Tensor,  # [N_k, Hkv, D]
    cu_seqlens_q: torch.Tensor,  # [B+1]
    cu_seqlens_k: torch.Tensor,  # [B+1]
    window: Union[int, torch.Tensor],
    sm_scale: float,
) -> torch.Tensor:
    """Compute token-shared SnapKV scores for packed, varlen q/k inputs.

    Returns a [N_k] float32 tensor, reduced across TP ranks so every rank makes
    an identical Top-K selection.
    """
    global _DISABLE_SNAPKV_TRITON

    device = query.device
    if query.ndim != 3 or key.ndim != 3:
        raise ValueError("query and key must be 3D tensors.")

    _, Hq, D = query.shape
    N_k, Hkv, Dk = key.shape
    if D != Dk:
        raise ValueError("query and key must have the same head size.")
    if Hq % Hkv != 0:
        raise ValueError("Query heads must be a multiple of KV heads.")
    if cu_seqlens_q.numel() != cu_seqlens_k.numel():
        raise ValueError("cu_seqlens_q and cu_seqlens_k must match.")

    # NOTE: Triton SnapKV scoring on ROCm is experimental. It is enabled by
    # default (uses a ROCm-safe kernel); set
    # VLLM_KV_COMPRESSION_SNAPKV_USE_TRITON_ROCM=0 to force the PyTorch
    # reference implementation.
    if (HAS_TRITON and not _DISABLE_SNAPKV_TRITON and device.type == "cuda"
            and (not current_platform.is_rocm()
                 or envs.VLLM_KV_COMPRESSION_SNAPKV_USE_TRITON_ROCM)
            and query.stride(-1) == 1 and key.stride(-1) == 1):
        try:
            from vllm.v1.kv_compression.snapkv_triton import (
                query_aware_key_scores)

            w = int(window) if isinstance(window, int) else window
            scores_per_head = query_aware_key_scores(
                q=query,
                k=key,
                cu_seqlens_q=cu_seqlens_q,
                cu_seqlens_k=cu_seqlens_k,
                w=w,
                sm_scale=float(sm_scale),
                pool=False,
                protect_last=False,
                normalize=False,
            )
            token_scores = scores_per_head.sum(dim=1)

            from vllm.distributed.parallel_state import get_tp_group

            return get_tp_group().all_reduce(token_scores)
        except Exception as e:
            _DISABLE_SNAPKV_TRITON = True
            logger.warning(
                "Triton SnapKV scoring failed; falling back to PyTorch. "
                "Error: %s", e)

    group = Hq // Hkv

    # Read boundaries and (optional) per-request window sizes on host.
    qsl = cu_seqlens_q.tolist()
    ksl = cu_seqlens_k.tolist()
    B = len(qsl) - 1
    if len(ksl) - 1 != B:
        raise ValueError("cu_seqlens_q and cu_seqlens_k must match.")

    wsl = None
    if not isinstance(window, int):
        if int(window.numel()) != B:
            raise ValueError("window must be a scalar int or have shape [B].")
        wsl = window.to(device="cpu", dtype=torch.int64).tolist()

    scores = torch.zeros((N_k, ), device=device, dtype=torch.float32)
    for b in range(B):
        qs = int(qsl[b])
        qe = int(qsl[b + 1])
        ks = int(ksl[b])
        ke = int(ksl[b + 1])
        q_len = qe - qs
        k_len = ke - ks
        if q_len <= 0 or k_len <= 0:
            continue

        win_b = int(window) if wsl is None else int(wsl[b])
        if win_b <= 0:
            continue
        win = min(win_b, q_len, k_len)

        k_eff_end = ke - win
        if k_eff_end <= ks:
            continue

        q_win = query[qe - win:qe]  # [win, Hq, D]
        # Aggregate query heads to KV heads (token-shared selection).
        q_win = q_win.reshape(win, Hkv, group, D).mean(dim=2)  # [win, Hkv, D]
        k_eff = key[ks:k_eff_end]  # [K_eff, Hkv, D]

        qh = q_win.permute(1, 0, 2).to(torch.float32)  # [Hkv, win, D]
        kh = k_eff.permute(1, 0, 2).to(torch.float32)  # [Hkv, K_eff, D]
        logits = torch.matmul(qh, kh.transpose(1, 2)) * float(sm_scale)
        probs = torch.softmax(logits, dim=-1)
        scores[ks:k_eff_end] = probs.sum(dim=1).sum(dim=0)

    from vllm.distributed.parallel_state import get_tp_group

    return get_tp_group().all_reduce(scores)


def snapkv_like_token_scores(
    *,
    query: torch.Tensor,  # [T, Hq, D]
    key: torch.Tensor,  # [T, Hkv, D]
    query_start_loc: torch.Tensor,  # [B+1]
    window: Union[int, torch.Tensor],
    sm_scale: float,
) -> torch.Tensor:
    """SnapKV-like token scores when q/k share the same packed layout."""
    return snapkv_query_aware_token_scores(
        query=query,
        key=key,
        cu_seqlens_q=query_start_loc,
        cu_seqlens_k=query_start_loc,
        window=window,
        sm_scale=sm_scale,
    )
