# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Paged attention."""

from collections.abc import Sequence

import torch
from torch import nn

from tokenspeed.runtime.execution.context import ForwardContext


class PagedAttention(nn.Module):
    """
    The attention layer implementation.
    """

    def __init__(
        self,
        num_heads: int,
        head_dim: int,
        scaling: float,
        num_kv_heads: int,
        layer_id: int,
        logit_cap: float = 0.0,
        v_head_dim: int = -1,
        sliding_window_size: int = -1,
        group_id: str = "",
    ):
        super().__init__()
        self.tp_q_head_num = num_heads
        self.tp_k_head_num = num_kv_heads
        self.tp_v_head_num = num_kv_heads
        self.head_dim = head_dim
        self.qk_head_dim = head_dim
        self.v_head_dim = v_head_dim if v_head_dim != -1 else head_dim
        self.scaling = scaling
        self.layer_id = layer_id
        self.logit_cap = logit_cap
        self.sliding_window_size = sliding_window_size or -1
        # Flat KV-cache group ("" -> single-table fallback in the backend).
        # TODO(radix-removal): make group_id mandatory once flat is the only path.
        self.group_id = group_id
        self.k_scale = None
        self.v_scale = None

    def forward(
        self,
        q,
        k,
        v,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        save_kv_cache: bool = True,
        **kwargs,
    ):
        if k is not None:
            # For cross-layer sharing, kv can be None
            if v is None:
                raise ValueError("v must be provided when k is provided.")
            if "k_pe" not in kwargs:
                k = k.view(-1, self.tp_k_head_num, self.qk_head_dim)
                v = v.view(-1, self.tp_v_head_num, self.v_head_dim)
            else:
                k = k.view(-1, self.tp_k_head_num, self.v_head_dim)
                v = v.view(-1, self.tp_v_head_num, self.v_head_dim)

        return ctx.attn_backend.forward(
            q,
            k,
            v,
            self,
            out_cache_loc,
            ctx.token_to_kv_pool,
            ctx.forward_mode,
            ctx.bs,
            save_kv_cache,
            **kwargs,
        )


def validate_paged_cache_group_ids(
    model: nn.Module,
    paged_cache_group_specs: Sequence,
) -> None:
    """Fail fast (ValueError) when a pool publishing more than one paged-cache
    group meets a PagedAttention layer whose group_id is empty or unknown --
    instead of a KeyError deep in the backend, possibly during graph capture.
    """
    group_ids = {str(spec.group_id) for spec in paged_cache_group_specs}
    if len(group_ids) <= 1:
        return
    model_name = type(model).__name__
    for name, module in model.named_modules():
        if not isinstance(module, PagedAttention):
            continue
        if not module.group_id:
            raise ValueError(
                f"{model_name}: attention layer {name!r} (layer_id="
                f"{module.layer_id}) has empty group_id but the KV pool "
                f"publishes {len(group_ids)} paged-cache groups "
                f"{sorted(group_ids)}; pass group_id=<layer_type> to "
                "PagedAttention (see gpt_oss.py)."
            )
        if module.group_id not in group_ids:
            raise ValueError(
                f"{model_name}: attention layer {name!r} (layer_id="
                f"{module.layer_id}) has group_id={module.group_id!r} which "
                "is not among the KV pool's paged-cache groups "
                f"{sorted(group_ids)}."
            )