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

from __future__ import annotations

from typing import Any, Optional, Protocol

import torch

import vllm.envs as envs
from vllm.forward_context import get_forward_context
from vllm.platforms import current_platform
from vllm.v1.kv_compression.compaction_step import compute_compact_dst_slots_for_step
from vllm.v1.kv_compression.forward_context import (
    get_kv_compression_compact_slots,
    get_kv_compression_prompt_payload,
    set_kv_compression_compact_slots,
    set_kv_compression_prompt_payload,
)
from vllm.v1.kv_compression.prompt_end import compute_prompt_end_indices
from vllm.v1.kv_compression.slot_mapping import kv_compaction_dst_rewrite_mapping


class _ReshapeAndCacheFn(Protocol):

    def __call__(
        self,
        key: torch.Tensor,
        value: torch.Tensor,
        key_cache: torch.Tensor,
        value_cache: torch.Tensor,
        slot_mapping: torch.Tensor,
        kv_cache_dtype: str,
        k_scale: torch.Tensor,
        v_scale: torch.Tensor,
    ) -> None: ...


def maybe_compute_prompt_end_payload_flash_attn(
    *,
    kv_sharing_target_layer_name: Optional[str],
    query: torch.Tensor,
    num_actual_tokens: int,
    key_cache: torch.Tensor,
    cache_block_size: int,
    attn_metadata: Any,
    sm_scale: float,
) -> None:
    """Compute and stash prompt-end Top-K indices for chunked-prefill scheme 3.

    The payload is cached in the forward context and later consumed by the
    model runner to perform one-shot prompt KV compaction before the first
    decode step.
    """
    if not envs.VLLM_ENABLE_KV_COMPRESSION or kv_sharing_target_layer_name is not None:
        return

    prompt_end = getattr(attn_metadata, "kv_compression_prompt_end", None)
    prompt_lens = getattr(attn_metadata, "kv_compression_prompt_lens", None)
    topk_keep = getattr(attn_metadata, "kv_compression_prompt_topk_keep", None)
    if prompt_end is None or prompt_lens is None or topk_keep is None:
        return
    B = int(prompt_end.numel())
    if B <= 0:
        return

    forward_context = get_forward_context()
    if get_kv_compression_prompt_payload(forward_context) is not None:
        return

    payload = compute_prompt_end_indices(
        query=query[:num_actual_tokens],
        key_cache=key_cache,
        block_size=cache_block_size,
        query_start_loc=attn_metadata.query_start_loc[:B + 1],
        block_table=attn_metadata.block_table[:B],
        prompt_end=prompt_end,
        prompt_lens=prompt_lens,
        topk_keep=topk_keep,
        topk_keep_max=getattr(attn_metadata, "kv_compression_prompt_topk_keep_max",
                              None),
        sm_scale=sm_scale,
    )
    if payload is not None:
        set_kv_compression_prompt_payload(forward_context, payload)


def maybe_compact_kv_cache_flash_attn(
    *,
    kv_sharing_target_layer_name: Optional[str],
    layer: Any,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    num_actual_tokens: int,
    cache_block_size: int,
    attn_metadata: Any,
    sm_scale: float,
    kv_cache_dtype: str,
    reshape_and_cache: _ReshapeAndCacheFn,
) -> None:
    """Optional per-step KV compaction for scheme 1/2 token-shared selection."""
    if not envs.VLLM_ENABLE_KV_COMPRESSION or kv_sharing_target_layer_name is not None:
        return

    must_keep = getattr(attn_metadata, "kv_compression_must_keep", None)
    topk_budget = getattr(attn_metadata, "kv_compression_topk_budget", None)
    if must_keep is None or topk_budget is None:
        return
    B = int(topk_budget.numel())
    if B <= 0:
        return

    forward_context = get_forward_context()
    per_layer_topk = envs.VLLM_KV_COMPRESSION_TOPK_PER_LAYER
    dst = get_kv_compression_compact_slots(
        forward_context,
        per_layer_topk=per_layer_topk,
        layer=layer,
    )

    if dst is None:
        topk_budget_max = int(
            getattr(attn_metadata, "kv_compression_topk_budget_max", 0) or 0)
        dst = compute_compact_dst_slots_for_step(
            query=query[:num_actual_tokens],
            key=key[:num_actual_tokens],
            query_start_loc=attn_metadata.query_start_loc[:B + 1],
            seq_lens=attn_metadata.seq_lens[:B],
            block_table=attn_metadata.block_table[:B],
            block_size=cache_block_size,
            must_keep=must_keep[:num_actual_tokens],
            topk_budget=topk_budget,
            topk_budget_max=topk_budget_max,
            max_query_len=attn_metadata.max_query_len,
            sm_scale=sm_scale,
        )
        set_kv_compression_compact_slots(
            forward_context,
            per_layer_topk=per_layer_topk,
            layer=layer,
            dst=dst,
        )

    if dst is None:
        return

    src = attn_metadata.slot_mapping[:num_actual_tokens]
    dst_rewrite = kv_compaction_dst_rewrite_mapping(dst_slots=dst,
                                                    src_slots=src)

    if not current_platform.is_rocm():
        reshape_and_cache(
            key,
            value,
            key_cache,
            value_cache,
            dst_rewrite,
            kv_cache_dtype,
            layer._k_scale,
            layer._v_scale,
        )
        return

    # ROCm: optionally prefer the optimized reshape-and-cache kernel.
    if (envs.VLLM_USE_OPT_RESHAPE_AND_CACHE and key.dtype == value.dtype
            and key.dtype == torch.float16):
        from lightop import reshape_and_cache_cuda

        reshape_and_cache_cuda(
            key,
            value,
            key_cache,
            value_cache,
            dst_rewrite,
            kv_cache_dtype,
            layer._k_scale,
            layer._v_scale,
        )
    else:
        reshape_and_cache(
            key,
            value,
            key_cache,
            value_cache,
            dst_rewrite,
            kv_cache_dtype,
            layer._k_scale,
            layer._v_scale,
        )
