# 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.

"""Forward mode enums and position computation helpers."""

from __future__ import annotations

from enum import IntEnum, auto

import torch
import triton
import triton.language as tl


class ForwardMode(IntEnum):
    # Extend a sequence. The KV cache of the beginning part of the sequence
    # is already computed (e.g., system prompt).
    EXTEND = auto()
    # Decode one or more tokens per request.
    DECODE = auto()
    # Contains both EXTEND and DECODE tokens in one batch.
    MIXED = auto()
    # No sequence to forward; used for data parallel attention idle ranks.
    IDLE = auto()

    def is_extend(self):
        return self == ForwardMode.EXTEND

    def is_decode(self):
        return self == ForwardMode.DECODE

    def is_mixed(self):
        return self == ForwardMode.MIXED

    def is_idle(self):
        return self == ForwardMode.IDLE

    def is_extend_or_mixed(self):
        return self == ForwardMode.EXTEND or self == ForwardMode.MIXED

    def is_decode_or_idle(self):
        return self == ForwardMode.DECODE or self == ForwardMode.IDLE

    @staticmethod
    def from_num_extends(num_extends: int, batch_size: int) -> "ForwardMode":
        if batch_size <= 0:
            return ForwardMode.IDLE
        elif num_extends > 0:
            return ForwardMode.MIXED if num_extends < batch_size else ForwardMode.EXTEND
        else:
            return ForwardMode.DECODE


class CaptureHiddenMode(IntEnum):
    NULL = auto()
    # Capture hidden states of all tokens.
    FULL = auto()
    # Capture a hidden state of the last token.
    LAST = auto()

    def need_capture(self):
        return self != CaptureHiddenMode.NULL

    def is_full(self):
        return self == CaptureHiddenMode.FULL

    def is_last(self):
        return self == CaptureHiddenMode.LAST


def compute_position_triton(
    extend_prefix_lens: torch.Tensor,
    extend_seq_lens: torch.Tensor,
    extend_seq_lens_sum,
    out: torch.Tensor | None = None,
):
    batch_size = extend_seq_lens.shape[0]
    if out is None:
        positions = torch.empty(
            extend_seq_lens_sum, dtype=torch.int64, device=extend_seq_lens.device
        )
    else:
        if out.numel() < extend_seq_lens_sum:
            raise ValueError(
                "compute_position_triton out buffer too small: "
                f"{out.numel()} < {extend_seq_lens_sum}"
            )
        positions = out
    extend_start_loc = torch.empty(
        batch_size, dtype=torch.int32, device=extend_seq_lens.device
    )
    has_prefix = extend_prefix_lens.shape[0] == batch_size
    # Launch kernel
    compute_position_kernel[(batch_size,)](
        positions, extend_start_loc, extend_prefix_lens, extend_seq_lens, has_prefix
    )

    return positions, extend_start_loc


@triton.jit
def compute_position_kernel(
    positions,
    extend_start_loc,
    extend_prefix_lens,
    extend_seq_lens,
    has_prefix: tl.constexpr,
):
    BLOCK_SIZE: tl.constexpr = 512
    pid = tl.program_id(0).to(tl.int64)

    prefix_len = tl.load(extend_prefix_lens + pid) if has_prefix else 0
    seq_len = tl.load(extend_seq_lens + pid)

    #  This can be slow for large bs
    cumsum_start = tl.cast(0, tl.int64)
    for i in range(pid):
        cumsum_start += tl.load(extend_seq_lens + i)

    num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
    for i in range(num_loop):
        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
        tl.store(
            positions + cumsum_start + offset,
            prefix_len + offset,
            mask=offset < seq_len,
        )
    tl.store(extend_start_loc + pid, cumsum_start)