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

from __future__ import annotations

from typing import Any

import vllm.envs as envs
from vllm.v1.kv_compression.budget import (compute_prompt_keep_len,
                                           compute_topk_budget_step,
                                           count_prompt_must_keep_in_range)


def maybe_init_num_kv_tokens_on_running_transition(
    *,
    request: Any,
    num_computed_tokens: int,
    kv_compression_enabled: bool,
) -> None:
    """在 request 切换为 RUNNING 时，必要时初始化 `request.num_kv_tokens`。

    - 未开启 KV compression：KV 实际长度始终等于逻辑长度，直接令
      `num_kv_tokens = num_computed_tokens` 即可。
    - 开启 KV compression：大多数请求从 0 token 开始（`num_computed_tokens == 0`），
      不需要额外初始化；但某些路径（例如 KV connector / cache hit）可能让一个请求在
      进入 RUNNING 时已经“预先拥有”一段已计算的 token。如果不把 `num_kv_tokens`
      初始化到同样的值，后续 KV 写入偏移（基于 `num_kv_tokens`）会从 0 开始，导致
      slot_mapping/KV cache 写入错位。
    """
    if not kv_compression_enabled:
        request.num_kv_tokens = num_computed_tokens
        return
    if getattr(request, "num_kv_tokens", 0) == 0 and num_computed_tokens > 0:
        request.num_kv_tokens = num_computed_tokens


def update_num_kv_tokens_after_schedule(
    *,
    request: Any,
    start_pos: int,
    num_scheduled_token: int,
    chunked_prefill_enabled: bool,
    kv_compression_enabled: bool,
) -> None:
    """在一次调度（一个 step）之后推进 `request.num_kv_tokens`。

    这是 KV compression 的“调度侧记账”函数（不做打分/TopK/重写 KV），目的仅是
    让 Scheduler 维护出“KV cache 实际长度”，以便后续：
    - KV block 分配（allocate_slots）
    - Worker 侧 slot_mapping / KV 写入偏移
    - attention metadata 里的 `seq_lens`
    都基于正确的 KV 长度工作。

    关键概念（针对单个 request）：
      - `num_computed_tokens`：逻辑进度（token 位置 / RoPE index）。
      - `num_kv_tokens`：KV cache 中“实际保留/存储”的 token 数。
        开启 KV compression 后，prompt KV 可能被压缩（只保留一部分），因此
        `num_kv_tokens` 可能小于 `num_computed_tokens`；但 decode token 始终全保留。

    参数含义：
      - `start_pos`：本 step 开始前的逻辑位置；本次新调度的逻辑区间为
        `[start_pos, end_pos)`。
      - `num_scheduled_token`：该 request 在本 step 被调度的 token 数（可能同时包含
        prompt token 和 decode token）。
      - `chunked_prefill_enabled`：是否启用 chunked prefill。
      - `kv_compression_enabled`：调度器 gate 后的“是否启用 KV compression”。

    注意：
      - 如果这里记账错了，worker 可能用错误的 KV 基址/长度生成 slot_mapping，
        从而导致 KV cache 读写错位甚至越界。
      - chunked prefill 模式下，为避免“下一段 prefill 看不到完整历史”导致质量崩溃，
        prompt 的 KV compaction 会延后到 prompt 结束后一次性执行（prompt-end one-shot）。
        因此 `num_kv_tokens` 的更新策略与非 chunked 模式不同。
    """
    if num_scheduled_token <= 0:
        return

    if not kv_compression_enabled:
        # 未开启压缩：KV cache 长度与逻辑长度始终一致，直接累加即可。
        request.num_kv_tokens += num_scheduled_token
        return

    # 开启 KV compression 后：prompt token 只保留子集；decode token 始终全保留。
    prompt_ratio = envs.VLLM_KV_COMPRESSION_PROMPT_RATIO
    prompt_budget = envs.VLLM_KV_COMPRESSION_PROMPT_BUDGET
    protected_prefix = envs.VLLM_KV_COMPRESSION_PROTECTED_PREFIX
    protected_suffix = envs.VLLM_KV_COMPRESSION_PROTECTED_SUFFIX
    keep_last = envs.VLLM_KV_COMPRESSION_KEEP_LAST_TOKEN

    end_pos = start_pos + int(num_scheduled_token)
    prompt_end = int(getattr(request, "num_prompt_tokens", 0))

    # Chunked prefill：prefill 过程中不要“边 ingest 边压缩 prompt KV”。
    # 否则下一段 chunk prefill 会注意力看不到完整历史（语义变化/质量崩溃）。
    # 正确策略是：prefill 阶段暂时全量保留 prompt KV；等 prompt ingest 完成后，
    # 在第一次 decode 前做一次性 prompt compaction（prompt-end one-shot）。
    if chunked_prefill_enabled:
        if start_pos >= prompt_end:
            # 纯 decode 段：decode token 始终全保留，KV 长度直接累加。
            request.num_kv_tokens += num_scheduled_token
            return

        if end_pos < prompt_end:
            # prompt 还没 ingest 完：暂时先全保留（不做 mid-prefill 压缩），KV 长度累加。
            request.num_kv_tokens += num_scheduled_token
            return

        # 重要：这里是“重置”而不是“累加”。
        # 因为 prompt 结束后 KV 长度会发生不连续跳变：
        #   - prompt ingest 过程中：KV cache 中存的是“完整 prompt 前缀”
        #   - prompt ingest 完成后：KV cache 中应变为“压缩后的 prompt”
        # 实际的 in-place compaction 在 worker 侧 decode 前执行；这里先把记账值更新到位。
        kept_prompt_total = compute_prompt_keep_len(
            prompt_len=prompt_end,
            protected_prefix=protected_prefix,
            protected_suffix=protected_suffix,
            keep_last_token=keep_last,
            prompt_ratio=prompt_ratio,
            prompt_budget=prompt_budget,
        )
        # 如果本 step 跨过 prompt_end，prompt_end 之后的 token 属于 decode，仍需全保留。
        kept_decode = max(0, end_pos - max(start_pos, prompt_end))
        request.num_kv_tokens = kept_prompt_total + kept_decode
        return

    # 非 chunked prefill（scheme 1/2）：每个 step 内做 token-shared 的选择。
    # - decode token：始终全保留；
    # - prompt token：只保留 must-keep（protected prefix/suffix/可选最后token）
    #   + 本 step Top-K 选中的部分。
    decode_start = max(start_pos, prompt_end)
    kept_decode = max(0, end_pos - decode_start)

    # 本 step 的逻辑区间内，prompt token 里“必须保留”的部分：
    # protected_prefix / protected_suffix /（可选）最后一个 prompt token。
    kept_prompt_must_keep = count_prompt_must_keep_in_range(
        prompt_len=prompt_end,
        start_pos=start_pos,
        end_pos=end_pos,
        protected_prefix=protected_prefix,
        protected_suffix=protected_suffix,
        keep_last_token=keep_last,
    )
    # 本 step 通过 Top-K 策略“额外保留”的 prompt token 数。
    # 预算定义在 prompt 的“非保护区”上，并由 `compute_topk_budget_step` 按 step 分摊。
    kept_prompt_topk = compute_topk_budget_step(
        prompt_len=prompt_end,
        start_pos=start_pos,
        end_pos=end_pos,
        protected_prefix=protected_prefix,
        protected_suffix=protected_suffix,
        keep_last_token=keep_last,
        prompt_ratio=prompt_ratio,
        prompt_budget=prompt_budget,
    )

    # 本 step 结束后：KV cache 实际长度按“保留的 KV 条目数”推进。
    request.num_kv_tokens += (kept_decode + kept_prompt_must_keep +
                              kept_prompt_topk)
