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

from __future__ import annotations

from typing import Optional

import torch

import vllm.envs as envs
from vllm.v1.kv_compression.kv_cache_view import paged_k_cache_view_for_triton_gather
from vllm.v1.kv_compression.snapkv_score import snapkv_query_aware_token_scores
from vllm.v1.kv_compression.topk_select import (_packed_varlen_coords,
                                                _topk_keep_mask_and_local_rank)


def _prompt_end_topk_keep_indices(
    *,
    token_scores: torch.Tensor,  # [T] float32
    prompt_lens: torch.Tensor,  # [B] int32
    topk_keep: torch.Tensor,  # [B] int32 (candidates only)
    protected_prefix: int,
    protected_suffix: int,
    keep_last_token: bool,
    topk_keep_max: Optional[int] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Select kept prompt indices (ascending) for one-shot compaction.

    Returns:
      idx_sorted: [B, K_max] int32, per-request kept token indices (0..L-1)
      keep_len: [B] int32, number of kept tokens per request
    """
    device = token_scores.device
    B = int(prompt_lens.numel())
    if B == 0:
        empty = torch.empty((0, 0), device=device, dtype=torch.int32)
        return empty, torch.empty((0, ), device=device, dtype=torch.int32)

    prompt_lens_i64 = prompt_lens.to(torch.long)
    cu = torch.zeros((B + 1, ), device=device, dtype=torch.long)
    cu[1:] = torch.cumsum(prompt_lens_i64, dim=0)
    T = int(token_scores.numel())
    if T == 0:
        empty = torch.empty((B, 0), device=device, dtype=torch.int32)
        return empty, torch.zeros((B, ), device=device, dtype=torch.int32)

    starts, _, lengths, req_ids, pos_in_req = _packed_varlen_coords(
        cu_seqlens=cu,
        total_tokens=T,
    )

    # Must-keep mask (protected prefix/suffix + optional last prompt token).
    prefix_len = torch.clamp(prompt_lens_i64,
                             min=0).clamp_max(max(protected_prefix, 0))
    suffix = torch.clamp(prompt_lens_i64,
                         min=0).clamp_max(max(protected_suffix, 0))
    suffix_start = (prompt_lens_i64 - suffix).clamp_min(0)
    prefix_len_t = prefix_len.index_select(0, req_ids)
    suffix_start_t = suffix_start.index_select(0, req_ids)
    must_keep = (pos_in_req < prefix_len_t) | (pos_in_req >= suffix_start_t)
    if keep_last_token:
        last = (prompt_lens_i64 - 1).clamp_min(0)
        last_t = last.index_select(0, req_ids)
        must_keep |= pos_in_req == last_t

    keep_mask, local_rank, keep_len = _topk_keep_mask_and_local_rank(
        token_scores=token_scores,
        must_keep=must_keep,
        topk_budget=topk_keep,
        starts=starts,
        lengths=lengths,
        req_ids=req_ids,
        pos_in_req=pos_in_req,
        max_len=int(prompt_lens_i64.max().item()),
        topk_budget_max=topk_keep_max,
    )

    keep_max_len = int(keep_len.max().item()) if B > 0 else 0
    if keep_max_len <= 0:
        empty = torch.empty((B, 0), device=device, dtype=torch.int32)
        return empty, keep_len

    idx_sorted = torch.zeros((B, keep_max_len), device=device, dtype=torch.int32)
    lin_out = (req_ids * keep_max_len + local_rank).masked_select(keep_mask)
    vals = pos_in_req.to(torch.int32).masked_select(keep_mask)
    idx_sorted.view(-1).scatter_(0, lin_out, vals)
    return idx_sorted, keep_len


def compute_prompt_end_indices(
    *,
    query: torch.Tensor,  # [T, Hq, D] scheduled tokens for this step
    key_cache: torch.Tensor,  # layer KV cache view (platform-dependent)
    block_size: int,
    query_start_loc: torch.Tensor,  # [B+1] int32
    block_table: torch.Tensor,  # [B, max_blocks] int32
    prompt_end: torch.Tensor,  # [B] bool
    prompt_lens: torch.Tensor,  # [B] int32
    topk_keep: torch.Tensor,  # [B] int32
    topk_keep_max: Optional[int],
    sm_scale: float,
) -> Optional[dict[str, torch.Tensor]]:
    """Compute one-shot prompt compaction indices on the last prefill chunk."""
    device = query.device
    if prompt_end.numel() == 0:
        return None
    sel = torch.nonzero(prompt_end, as_tuple=False).flatten()
    if int(sel.numel()) == 0:
        return None

    window = int(envs.VLLM_KV_COMPRESSION_SNAPKV_WINDOW)
    keep_last = bool(envs.VLLM_KV_COMPRESSION_KEEP_LAST_TOKEN)
    protected_prefix = int(envs.VLLM_KV_COMPRESSION_PROTECTED_PREFIX)
    protected_suffix = int(envs.VLLM_KV_COMPRESSION_PROTECTED_SUFFIX)

    # Build packed Q window (last `window` queries per selected request).
    sel_list = sel.to(device="cpu", dtype=torch.int64).tolist()
    qsl = query_start_loc.to(device="cpu", dtype=torch.int64).tolist()
    q_chunks = []
    cu_q = [0]
    w_list = []
    for b in sel_list:
        s = int(qsl[b])
        e = int(qsl[b + 1])
        q_len = max(0, e - s)
        win = min(window, q_len)
        w_list.append(int(win))
        if win > 0:
            q_chunks.append(query[e - win:e])
        cu_q.append(cu_q[-1] + int(win))

    if cu_q[-1] <= 0:
        return None
    q_packed = torch.cat(q_chunks, dim=0) if q_chunks else query[:0]
    cu_seqlens_q = torch.tensor(cu_q, device=device, dtype=torch.int32)
    w = torch.tensor(w_list, device=device, dtype=torch.int32)

    # Gather full prompt keys for the selected requests into a packed [T, Hk, D].
    prompt_lens_sel = prompt_lens.index_select(0, sel).to(torch.int32)
    topk_keep_sel = topk_keep.index_select(0, sel).to(torch.int32)
    cu_seqlens_k = torch.zeros((int(prompt_lens_sel.numel()) + 1, ),
                               device=device,
                               dtype=torch.int32)
    if int(prompt_lens_sel.numel()) > 0:
        cu_seqlens_k[1:] = torch.cumsum(prompt_lens_sel, dim=0)

    block_table_sel = block_table.index_select(0, sel).to(torch.int32)
    key_cache_view = paged_k_cache_view_for_triton_gather(
        key_cache=key_cache,
        block_size=int(block_size),
    )

    from vllm.v1.kv_compression.kv_cache_triton import (
        gather_k_to_packed_triton)

    k_packed = gather_k_to_packed_triton(
        key_cache_view,
        block_table_sel,
        prompt_lens_sel,
        cu_seqlens_k,
    )

    token_scores = snapkv_query_aware_token_scores(
        query=q_packed,
        key=k_packed,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k,
        window=w,
        sm_scale=float(sm_scale),
    )

    idx_sorted, keep_len = _prompt_end_topk_keep_indices(
        token_scores=token_scores,
        prompt_lens=prompt_lens_sel,
        topk_keep=topk_keep_sel,
        protected_prefix=protected_prefix,
        protected_suffix=protected_suffix,
        keep_last_token=keep_last,
        topk_keep_max=topk_keep_max,
    )

    return {
        "req_indices": sel.to(torch.int32),
        "idx_sorted": idx_sorted,  # [B_sel, K_max] int32
        "keep_len": keep_len,  # [B_sel] int32
        "prompt_lens": prompt_lens_sel,  # [B_sel] int32
    }
