from __future__ import annotations

import os
import random
from collections import deque
from contextlib import nullcontext
from enum import Enum
from typing import (
    TYPE_CHECKING,
    Iterable,
    List,
    Literal,
    Optional,
    Tuple,
    Type,
    overload,
)

import numpy as np
import torch
import torch.distributed as dist

from sglang.srt.configs.model_config import get_dsa_index_topk
from sglang.srt.disaggregation.base import KVPoll
from sglang.srt.environ import envs
from sglang.srt.utils import is_hip, is_npu

if TYPE_CHECKING:
    from sglang.srt.disaggregation.base.conn import KVArgs, StateType
    from sglang.srt.disaggregation.common.conn import (
        CommonKVBootstrapServer,
        CommonKVManager,
        CommonKVReceiver,
        CommonKVSender,
    )
    from sglang.srt.managers.schedule_batch import Req
    from sglang.srt.server_args import ServerArgs

if is_npu():
    from sglang.srt.hardware_backend.npu.dsv4.dsv4_memory_pool import (
        DSV4NPUTokenToKVPool,
    )

#########################
# Constants & Enums
#########################
FAKE_BOOTSTRAP_HOST = "2.2.2.2"
_IS_HIP = is_hip()


def poll_and_all_reduce_pp(
    rids: Iterable[str],
    ready_poll: int,
    pp_good_rids: Optional[List[str]] = None,
    pp_bad_rids: Optional[List[str]] = None,
) -> List[Optional[int]]:
    """Map authoritative PP consensus to poll states without polling again."""
    if pp_good_rids is None or pp_bad_rids is None:
        raise ValueError("PP consensus is required")

    good_rids = set(pp_good_rids)
    bad_rids = set(pp_bad_rids)
    return [
        KVPoll.Failed if rid in bad_rids else ready_poll if rid in good_rids else None
        for rid in rids
    ]


def get_dsa_seed_metadata_dim(hf_config) -> int:
    """Return the model-defined PD seed width, independent of local spec mode."""
    if not getattr(hf_config, "index_share_for_mtp_iteration", False):
        return 0
    return get_dsa_index_topk(hf_config)


def is_dsv4_c128_online_enabled() -> bool:
    """Return whether DSV4 C128 uses request-scoped online state."""
    return not _IS_HIP and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()


def get_dsv4_c128_state_indices(
    req_pool_idx: int,
    seq_len: int,
    *,
    online: bool,
    ring_size: int,
) -> np.ndarray:
    """Return the PD transfer row/page indices for DSV4 C128 state."""
    if seq_len == 0 or seq_len % 128 == 0:
        return np.empty((0,), dtype=np.int32)
    if online:
        return np.array([int(req_pool_idx)], dtype=np.int32)

    assert ring_size % 128 == 0, f"C128 ring_size must be 128-aligned, got {ring_size}"
    pages_per_req = ring_size // 128
    page = int(req_pool_idx) * pages_per_req + ((seq_len - 1) % ring_size) // 128
    return np.array([page], dtype=np.int32)


class DisaggregationMode(Enum):
    NULL = "null"
    PREFILL = "prefill"
    DECODE = "decode"

    @staticmethod
    def to_engine_type(mode: str) -> str:
        if mode == DisaggregationMode.PREFILL.value:
            return "prefill"
        elif mode == DisaggregationMode.DECODE.value:
            return "decode"
        return "unified"


#########################
# Synchronization
#########################


def _get_failure_prob() -> float:
    try:
        return float(envs.SGLANG_TEST_DISAGG_FAILURE_PROB.get())
    except Exception:
        # fallback to legacy env var
        return float(os.getenv("DISAGGREGATION_TEST_FAILURE_PROB", "0"))


def _poll_with_failure_injection(pollers) -> List[int]:
    if (failure_prob := _get_failure_prob()) > 0:
        return [
            int(KVPoll.Failed) if random.random() < failure_prob else int(poller.poll())
            for poller in pollers
        ]
    return [int(poller.poll()) for poller in pollers]


def _is_fake_transfer(req: Req, server_args: ServerArgs) -> bool:
    return req.bootstrap_host == FAKE_BOOTSTRAP_HOST or (
        req.bootstrap_host is None
        and server_args.disaggregation_transfer_backend == "fake"
    )


def _apply_metadata_gate(polls, decode_reqs, metadata_buffers, server_args) -> None:
    """Downgrade Success → Transferring for requests whose metadata hasn't landed.

    Mutates `polls` in-place. Called before all-reduce so that MIN across TP
    ranks naturally prevents any rank from committing before all ranks are ready.
    """
    for i, poll_val in enumerate(polls):
        if poll_val == int(KVPoll.Success):
            decode_req = decode_reqs[i]
            if _is_fake_transfer(decode_req.req, server_args):
                continue
            actual_room = metadata_buffers.bootstrap_room[
                decode_req.metadata_buffer_index, 0
            ].item()
            if actual_room == 0:
                polls[i] = int(KVPoll.Transferring)


def poll_and_all_reduce(
    pollers,
    gloo_group: dist.ProcessGroup,
    decode_reqs=None,
    metadata_buffers: Optional[MetadataBuffers] = None,
    server_args: Optional[ServerArgs] = None,
):
    # at a certain prob, the poll is failed to simulate failure
    polls = _poll_with_failure_injection(pollers)

    # Apply metadata gate on the decode requests to downgrade Success → Transferring for requests whose metadata hasn't landed.
    if (
        decode_reqs is not None
        and metadata_buffers is not None
        and server_args is not None
    ):
        _apply_metadata_gate(polls, decode_reqs, metadata_buffers, server_args)
    tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
    dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=gloo_group)
    return tensor_to_reduce.tolist()


def poll_and_all_reduce_attn_cp_tp_group(
    pollers,
    attn_cp_cpu_group: dist.ProcessGroup,
    attn_tp_cpu_group: dist.ProcessGroup,
):
    # First sync across attn-tp ranks so all TP participants for a given (dp, cp)
    # shard observe the same status transitions.
    polls = poll_and_all_reduce(pollers, attn_tp_cpu_group)

    # Then sync across attn-cp ranks, so all TPxCP participants in one DP shard
    # converge to the same global status.
    tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
    dist.all_reduce(
        tensor_to_reduce,
        op=dist.ReduceOp.MIN,
        group=attn_cp_cpu_group,
    )
    return tensor_to_reduce.tolist()


def poll_and_all_reduce_with_staging(
    decode_reqs,
    staging_handler,
    gloo_group: dist.ProcessGroup,
    metadata_buffers: Optional[MetadataBuffers] = None,
    server_args: Optional[ServerArgs] = None,
):
    """Staging-aware polling: advance scatter, demote incomplete transfers, all_reduce."""
    for decode_req in decode_reqs:
        if decode_req.kv_receiver.require_staging and not staging_handler.is_done(
            decode_req
        ):
            staging_handler.advance_scatter(decode_req)

    # allow test injection of failure probability at runtime
    receivers = [dr.kv_receiver for dr in decode_reqs]
    raw_polls = _poll_with_failure_injection(receivers)
    for i, decode_req in enumerate(decode_reqs):
        if decode_req.kv_receiver.require_staging and staging_handler.is_failed(
            decode_req
        ):
            # Staging completion timed out; KVPoll.Failed == 0 propagates
            # through the MIN all_reduce.
            raw_polls[i] = int(KVPoll.Failed)
            continue
        if raw_polls[i] == int(KVPoll.Success):
            if decode_req.kv_receiver.require_staging and not staging_handler.is_done(
                decode_req
            ):
                raw_polls[i] = int(KVPoll.Transferring)
    # Apply metadata gate on the decode requests to downgrade Success → Transferring for requests whose metadata hasn't landed.
    if metadata_buffers is not None and server_args is not None:
        _apply_metadata_gate(raw_polls, decode_reqs, metadata_buffers, server_args)
    poll_tensor = torch.tensor(raw_polls, dtype=torch.uint8, device="cpu")
    dist.all_reduce(poll_tensor, op=dist.ReduceOp.MIN, group=gloo_group)
    return poll_tensor.tolist()


#########################
# Metadata Buffers
#########################


class ReqToMetadataIdxAllocator:
    """A memory pool that maps a request to its first output token location."""

    def __init__(
        self,
        size: int,
    ):
        self.size = size
        self.free_slots = deque(list(range(size)))

    def available_size(self):
        return len(self.free_slots)

    def alloc(self) -> Optional[int]:
        if len(self.free_slots) == 0:
            return None

        return self.free_slots.popleft()

    def free(self, free_index: int):
        self.free_slots.append(free_index)


class MetadataBuffers:
    def __init__(
        self,
        size: int,
        hidden_size: int,
        hidden_states_dtype: torch.dtype,
        max_top_logprobs_num: int = 128,
        max_sampling_mask_tokens: Optional[int] = None,
        custom_mem_pool: torch.cuda.MemPool = None,
        output_dsa_topk_indices_dim: int = 0,
    ):
        self.custom_mem_pool = custom_mem_pool
        self.output_dsa_topk_indices_dim = output_dsa_topk_indices_dim
        if max_sampling_mask_tokens is None:
            max_sampling_mask_tokens = (
                envs.SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.get()
            )
        self.enable_sampling_mask = max_sampling_mask_tokens > 0
        bootstrap_room_dtype = torch.uint64
        device = "cpu"
        if is_npu():
            # For ascend backend, output tokens are placed in the NPU and will be transferred by D2D channel.
            device = "npu"
            # TODO: Fix me when npu backend supports torch.uint64
            bootstrap_room_dtype = torch.int64
        elif self.custom_mem_pool:
            # TODO(shangming): Fix me (use 'cuda') when nvlink_transport of Mooncake is bug-free
            device = "cpu"
        elif envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() == "INTRA_NODE_NVLINK":
            device = "cuda"
        with (
            torch.cuda.use_mem_pool(self.custom_mem_pool)
            if self.custom_mem_pool
            else nullcontext()
        ):
            # TODO: abort top_logprobs_num > 128 in PD

            # We transfer the metadata of first output token to decode
            # The minimal size for RDMA is 64Bytes, so we pad it to > 64Bytes
            self.output_ids = torch.zeros((size, 16), dtype=torch.int32, device=device)
            self.cached_tokens = torch.zeros(
                (size, 16), dtype=torch.int32, device=device
            )
            self.output_token_logprobs_val = torch.zeros(
                (size, 16), dtype=torch.float32, device=device
            )
            self.output_token_logprobs_idx = torch.zeros(
                (size, 16), dtype=torch.int32, device=device
            )
            self.output_top_logprobs_val = torch.zeros(
                (size, max_top_logprobs_num), dtype=torch.float32, device=device
            )
            self.output_top_logprobs_idx = torch.zeros(
                (size, max_top_logprobs_num), dtype=torch.int32, device=device
            )
            self.output_token_sampling_mask_len = None
            self.output_token_sampling_mask_idx = None
            self.output_token_sampling_logprobs = None
            if self.enable_sampling_mask:
                self.output_token_sampling_mask_len = torch.zeros(
                    (size, 16), dtype=torch.int32, device=device
                )
                self.output_token_sampling_mask_idx = torch.zeros(
                    (size, max_sampling_mask_tokens), dtype=torch.int32, device=device
                )
                self.output_token_sampling_logprobs = torch.zeros(
                    (size, 16), dtype=torch.float32, device=device
                )
            # For PD + spec decode
            self.output_topk_p = torch.zeros(
                (size, 16), dtype=torch.float32, device=device
            )
            self.output_topk_index = torch.zeros(
                (size, 16), dtype=torch.int64, device=device
            )
            self.output_hidden_states = torch.zeros(
                (size, hidden_size), dtype=hidden_states_dtype, device=device
            )
            if self.output_dsa_topk_indices_dim > 0:
                self.output_dsa_topk_indices = torch.full(
                    (size, self.output_dsa_topk_indices_dim),
                    -1,
                    dtype=torch.int32,
                    device=device,
                )
            else:
                self.output_dsa_topk_indices = None
            # Request validation: store bootstrap_room to detect metadata corruption
            self.bootstrap_room = torch.zeros(
                (size, 8), dtype=bootstrap_room_dtype, device=device
            )

    def get_buf_infos(self):
        bufs = [
            self.output_ids,
            self.cached_tokens,
            self.output_token_logprobs_val,
            self.output_token_logprobs_idx,
            self.output_top_logprobs_val,
            self.output_top_logprobs_idx,
        ]
        if self.enable_sampling_mask:
            bufs.extend(
                [
                    self.output_token_sampling_mask_len,
                    self.output_token_sampling_mask_idx,
                    self.output_token_sampling_logprobs,
                ]
            )
        bufs.extend(
            [
                self.output_topk_p,
                self.output_topk_index,
                self.output_hidden_states,
            ]
        )
        if self.output_dsa_topk_indices is not None:
            bufs.append(self.output_dsa_topk_indices)
        bufs.append(self.bootstrap_room)
        ptrs = [buf.data_ptr() for buf in bufs]
        data_lens = [buf.nbytes for buf in bufs]
        item_lens = [buf[0].nbytes for buf in bufs]
        return ptrs, data_lens, item_lens

    def get_buf(self, idx: int):
        sampling_mask_len = None
        sampling_mask_idx = None
        sampling_logprobs = None
        if self.enable_sampling_mask:
            sampling_mask_len = self.output_token_sampling_mask_len[idx].clone()
            sampling_mask_idx = self.output_token_sampling_mask_idx[idx].clone()
            sampling_logprobs = self.output_token_sampling_logprobs[idx].clone()
        return (
            self.output_ids[idx].clone(),
            self.cached_tokens[idx].clone(),
            self.output_token_logprobs_val[idx].clone(),
            self.output_token_logprobs_idx[idx].clone(),
            self.output_top_logprobs_val[idx].clone(),
            self.output_top_logprobs_idx[idx].clone(),
            sampling_mask_len,
            sampling_mask_idx,
            sampling_logprobs,
            self.output_topk_p[idx].clone(),
            self.output_topk_index[idx].clone(),
            self.output_hidden_states[idx].clone(),
            (
                self.output_dsa_topk_indices[idx].clone()
                if self.output_dsa_topk_indices is not None
                else None
            ),
            self.bootstrap_room[idx].clone(),
        )

    def set_buf(self, req: Req):

        self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0]
        # The cached_tokens buffer is (size, 16); slots 0-3 hold cached token
        # counts and slots 4-6 are reused for multimodal prompt token counts
        # (slots 7-15 remain spare). This avoids adding new RDMA buffers.
        # Slot map: 0=cached 1=device 2=host 3=storage 4=image 5=audio 6=video.
        self.cached_tokens[req.metadata_buffer_index][0] = req.cached_tokens
        self.cached_tokens[req.metadata_buffer_index][1] = req.cached_tokens_device
        self.cached_tokens[req.metadata_buffer_index][2] = req.cached_tokens_host
        self.cached_tokens[req.metadata_buffer_index][3] = req.cached_tokens_storage

        # Compute multimodal prompt token counts on the prefill node so decode
        # can report them in usage.
        if req.multimodal_inputs:
            image_t, audio_t, video_t = req.multimodal_inputs.compute_mm_token_counts()
        else:
            image_t = audio_t = video_t = 0
        self.cached_tokens[req.metadata_buffer_index][4] = image_t
        self.cached_tokens[req.metadata_buffer_index][5] = audio_t
        self.cached_tokens[req.metadata_buffer_index][6] = video_t
        if req.return_logprob:
            if req.logprob.output_token_logprobs_val:  # not none or empty list
                self.output_token_logprobs_val[req.metadata_buffer_index][0] = (
                    req.logprob.output_token_logprobs_val[0]
                )
            if req.logprob.output_token_logprobs_idx:  # not none or empty list
                self.output_token_logprobs_idx[req.metadata_buffer_index][0] = (
                    req.logprob.output_token_logprobs_idx[0]
                )

            if req.logprob.output_top_logprobs_val:  # not none or empty list
                top_logprobs_len = len(req.logprob.output_top_logprobs_val[0])
                max_top_logprobs_len = self.output_top_logprobs_val.shape[1]
                if top_logprobs_len > max_top_logprobs_len:
                    raise RuntimeError(
                        f"top_logprobs_num {top_logprobs_len} exceeds "
                        f"disaggregation metadata capacity {max_top_logprobs_len}. "
                        "Lower top_logprobs_num or increase the metadata buffer."
                    )
                self.output_top_logprobs_val[req.metadata_buffer_index][
                    : len(req.logprob.output_top_logprobs_val[0])
                ] = torch.tensor(
                    req.logprob.output_top_logprobs_val[0],
                    dtype=torch.float32,
                    device="cpu",
                )
            if req.logprob.output_top_logprobs_idx:  # not none or empty list
                self.output_top_logprobs_idx[req.metadata_buffer_index][
                    : len(req.logprob.output_top_logprobs_idx[0])
                ] = torch.tensor(
                    req.logprob.output_top_logprobs_idx[0],
                    dtype=torch.int32,
                    device="cpu",
                )
        if req.return_sampling_mask:
            if not self.enable_sampling_mask:
                raise RuntimeError(
                    "return_sampling_mask with disaggregation requires "
                    "SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS > 0."
                )
            # Sentinel -1: the decode side records None for this handoff token.
            self.output_token_sampling_mask_len[req.metadata_buffer_index][0] = -1
            sampling_masks = req.output_token_sampling_mask
            sampling_logprobs = req.output_token_sampling_logprobs
            if sampling_masks:
                sampling_mask = sampling_masks[0]
                sampling_logprob = sampling_logprobs[0] if sampling_logprobs else None
                if sampling_mask is not None and sampling_logprob is not None:
                    mask_len = len(sampling_mask)
                    max_mask_len = self.output_token_sampling_mask_idx.shape[1]
                    if mask_len > max_mask_len:
                        raise RuntimeError(
                            f"Sampling mask length {mask_len} exceeds disaggregation "
                            f"metadata capacity {max_mask_len}. Increase "
                            "SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS."
                        )
                    self.output_token_sampling_mask_len[req.metadata_buffer_index][
                        0
                    ] = mask_len
                    if mask_len:
                        self.output_token_sampling_mask_idx[
                            req.metadata_buffer_index, :mask_len
                        ].copy_(
                            torch.tensor(
                                sampling_mask,
                                dtype=torch.int32,
                                device=self.output_token_sampling_mask_idx.device,
                            )
                        )
                    self.output_token_sampling_logprobs[req.metadata_buffer_index][
                        0
                    ] = float(sampling_logprob)
        # For PD + spec decode
        if req.hidden_states_tensor is not None:
            # speculative_eagle_topk should not be greater than 16 currently
            topk = req.output_topk_p.size(0)

            self.output_topk_p[req.metadata_buffer_index, :topk].copy_(
                req.output_topk_p
            )
            self.output_topk_index[req.metadata_buffer_index, :topk].copy_(
                req.output_topk_index
            )
            self.output_hidden_states[req.metadata_buffer_index].copy_(
                req.hidden_states_tensor
            )
            if self.output_dsa_topk_indices is not None:
                dsa_topk_indices = req.output_dsa_topk_indices
                if dsa_topk_indices is not None:
                    self.output_dsa_topk_indices[req.metadata_buffer_index].copy_(
                        dsa_topk_indices
                    )
                else:
                    self.output_dsa_topk_indices[req.metadata_buffer_index].fill_(-1)
        # Store bootstrap_room for validation on decode side
        self.bootstrap_room[req.metadata_buffer_index, 0] = (
            req.bootstrap_room if req.bootstrap_room is not None else 0
        )


#########################
# Transfer Backend
#########################


class TransferBackend(Enum):
    MOONCAKE = "mooncake"
    MORI = "mori"
    NIXL = "nixl"
    ASCEND = "ascend"
    FAKE = "fake"


class KVClassType(Enum):
    KVARGS = "kvargs"
    MANAGER = "manager"
    SENDER = "sender"
    RECEIVER = "receiver"
    BOOTSTRAP_SERVER = "bootstrap_server"


@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.KVARGS]
) -> Type[KVArgs]: ...
@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.MANAGER]
) -> Type[CommonKVManager]: ...
@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.SENDER]
) -> Type[CommonKVSender]: ...
@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.RECEIVER]
) -> Type[CommonKVReceiver]: ...
@overload
def get_kv_class(
    transfer_backend: TransferBackend, class_type: Literal[KVClassType.BOOTSTRAP_SERVER]
) -> Type[CommonKVBootstrapServer]: ...


def get_kv_class(
    transfer_backend: TransferBackend, class_type: KVClassType
) -> Optional[Type]:
    from sglang.srt.disaggregation.fake import FakeKVReceiver, FakeKVSender

    if transfer_backend == TransferBackend.MOONCAKE:
        from sglang.srt.disaggregation.base import KVArgs
        from sglang.srt.disaggregation.mooncake import (
            MooncakeKVBootstrapServer,
            MooncakeKVManager,
            MooncakeKVReceiver,
            MooncakeKVSender,
        )

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: MooncakeKVManager,
            KVClassType.SENDER: MooncakeKVSender,
            KVClassType.RECEIVER: (MooncakeKVReceiver),
            KVClassType.BOOTSTRAP_SERVER: MooncakeKVBootstrapServer,
        }
        return class_mapping.get(class_type)
    elif transfer_backend == TransferBackend.MORI:
        from sglang.srt.disaggregation.base import KVArgs
        from sglang.srt.disaggregation.mori import (
            MoriKVBootstrapServer,
            MoriKVManager,
            MoriKVReceiver,
            MoriKVSender,
        )

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: MoriKVManager,
            KVClassType.SENDER: MoriKVSender,
            KVClassType.RECEIVER: (MoriKVReceiver),
            KVClassType.BOOTSTRAP_SERVER: MoriKVBootstrapServer,
        }
        return class_mapping.get(class_type)
    elif transfer_backend == TransferBackend.ASCEND:
        from sglang.srt.disaggregation.ascend import (
            AscendKVBootstrapServer,
            AscendKVManager,
            AscendKVReceiver,
            AscendKVSender,
        )
        from sglang.srt.disaggregation.base import KVArgs

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: AscendKVManager,
            KVClassType.SENDER: AscendKVSender,
            KVClassType.RECEIVER: (AscendKVReceiver),
            KVClassType.BOOTSTRAP_SERVER: AscendKVBootstrapServer,
        }
        return class_mapping.get(class_type)
    elif transfer_backend == TransferBackend.NIXL:
        from sglang.srt.disaggregation.base import KVArgs
        from sglang.srt.disaggregation.nixl import (
            NixlKVBootstrapServer,
            NixlKVManager,
            NixlKVReceiver,
            NixlKVSender,
        )

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: NixlKVManager,
            KVClassType.SENDER: NixlKVSender,
            KVClassType.RECEIVER: (NixlKVReceiver),
            KVClassType.BOOTSTRAP_SERVER: NixlKVBootstrapServer,
        }
        return class_mapping.get(class_type)
    elif transfer_backend == TransferBackend.FAKE:
        from sglang.srt.disaggregation.base import KVArgs
        from sglang.srt.disaggregation.fake import (
            FakeKVManager,
            FakeKVReceiver,
            FakeKVSender,
        )

        class_mapping = {
            KVClassType.KVARGS: KVArgs,
            KVClassType.MANAGER: FakeKVManager,
            KVClassType.SENDER: FakeKVSender,
            KVClassType.RECEIVER: (FakeKVReceiver),
        }
        return class_mapping.get(class_type)

    raise ValueError(f"Unsupported transfer backend: {transfer_backend}")


def _get_cp_rank_page_bounds(
    total_pages: int, cp_rank: int, cp_size: int
) -> Tuple[int, int]:
    base = total_pages // cp_size
    rem = total_pages % cp_size
    local_start = cp_rank * base + min(cp_rank, rem)
    n_pages = base + (1 if cp_rank < rem else 0)
    return local_start, local_start + n_pages


def page_indices_to_cp_rank_page_indices(
    page_indices: np.ndarray,
    total_pages: int,
    cp_rank: int,
    cp_size: int,
) -> np.ndarray:
    """
    Filter page_indices (which are *global* page ids in the KV pool) to those
    belonging to the given CP rank for this request.

    For a single request, its pages occupy a contiguous global range
    [first_page, first_page + total_pages). We first compute the local
    split [0, total_pages) across cp_size ranks, then shift that local
    range by first_page back into the global page id space and take
    the intersection with page_indices.

    Returns:
        Subset of page_indices that fall in this rank's global
        [start_page, end_page) slice for the given CP rank.
    """
    if cp_size <= 1:
        return page_indices

    if page_indices.size == 0:
        return np.asarray(page_indices)

    first_page = int(page_indices.min())
    base = total_pages // cp_size
    rem = total_pages % cp_size

    if rem == 0:
        local_start = cp_rank * base
        local_end = local_start + base
    else:
        local_start = cp_rank * base + min(cp_rank, rem)
        n_pages = base + (1 if cp_rank < rem else 0)
        local_end = local_start + n_pages

    # Map back to global page ids.
    start_page = first_page + local_start
    end_page = first_page + local_end

    mask = (page_indices >= start_page) & (page_indices < end_page)
    return np.asarray(page_indices)[mask]


def filter_kv_indices_for_cp_rank(
    kv_mgr: CommonKVManager,
    kv_indices: np.ndarray,
    index_slice: slice,
    total_pages: Optional[int] = None,
) -> Tuple[np.ndarray, slice]:
    """Filters kv_indices and index_slice for the current CP rank."""
    if total_pages is None:
        total_pages = len(kv_indices)
    cp_rank = kv_mgr.attn_cp_rank
    cp_size = kv_mgr.attn_cp_size

    if cp_size <= 1:
        return kv_indices, index_slice

    rank_start, rank_end = _get_cp_rank_page_bounds(total_pages, cp_rank, cp_size)
    chunk_start = index_slice.start if index_slice.start is not None else 0
    chunk_end = index_slice.stop if index_slice.stop is not None else total_pages
    first_pos = max(rank_start, chunk_start) - chunk_start
    last_pos = min(rank_end, chunk_end) - chunk_start

    if last_pos <= first_pos:
        new_kv_indices = kv_indices[:0]
        new_index_slice = slice(chunk_start, chunk_start)
    else:
        new_kv_indices = kv_indices[first_pos:last_pos]
        new_index_slice = slice(
            chunk_start + first_pos,
            chunk_start + last_pos,
        )
    return new_kv_indices, new_index_slice


#########################
# Misc
#########################


def is_mla_backend(target_kv_pool) -> bool:
    from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
    from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool

    return isinstance(target_kv_pool, (MLATokenToKVPool, DeepSeekV4TokenToKVPool))


def compute_mamba_state_slice_blocks(
    src_dim: int,
    dst_dim: int,
    src_attn_tp_size: int,
    dst_attn_tp_size: int,
    dst_tp_rank_in_group: int,
    local_tp_rank_in_group: int,
    conv_shard_groups: Optional[List[int]] = None,
) -> List[Tuple[int, int, int]]:
    """Blocks to copy one mamba state item across differing attn-TP sizes.

    Returns ``(src_dim_start, dst_dim_start, num_dims)`` triples in units of the
    sliceable (3rd) dimension. Single-axis states (temporal_state, or when
    ``conv_shard_groups`` is None) return one contiguous block -- byte-identical to
    the legacy behavior.

    GDN conv_state is ``cat([query | key | value])`` where each sub-block (full
    dims == ``conv_shard_groups``, e.g. ``[key_dim, key_dim, value_dim]``) is
    head-sharded INDEPENDENTLY across attn-TP. In the SCATTER direction
    (1 prefill rank -> several decode ranks) a single contiguous slice straddles
    the q/k/v boundaries and delivers wrong channels. The AGGREGATION direction
    (several prefill ranks -> 1 decode rank) has the symmetric problem: a single
    contiguous write interleaves the sub-blocks by writer. Both directions emit one
    block per sub-block for conv_state; temporal_state and non-GDN states (when
    ``conv_shard_groups`` is None) keep the single contiguous slice.
    """
    use_subdims = (
        conv_shard_groups is not None
        and sum(conv_shard_groups) == src_dim * src_attn_tp_size
    )

    if src_attn_tp_size > dst_attn_tp_size:
        # Aggregation: several prefill ranks each write their shard into one decode slot.
        writers_per_decode = src_attn_tp_size // dst_attn_tp_size
        local_writer_idx = local_tp_rank_in_group % writers_per_decode
        if not use_subdims:
            return [(0, local_writer_idx * src_dim, src_dim)]
        # conv_state: a plain contiguous write would interleave the sub-blocks by
        # writer ([q0,k0,v0,q1,k1,v1,...]); place this writer's shard of each
        # independently head-sharded sub-block at its grouped offset so the decode
        # buffer is [q0,q1,...,k0,k1,...,v0,v1,...].
        blocks: List[Tuple[int, int, int]] = []
        src_off = 0
        dst_off = 0
        for full_sd in conv_shard_groups:
            src_sub = full_sd // src_attn_tp_size
            dst_sub = full_sd // dst_attn_tp_size
            blocks.append((src_off, dst_off + local_writer_idx * src_sub, src_sub))
            src_off += src_sub
            dst_off += dst_sub
        return blocks

    # Scatter: 1 prefill rank feeds several decode ranks.
    if not use_subdims:
        src_dim_start = (dst_tp_rank_in_group * dst_dim) % src_dim
        return [(src_dim_start, 0, dst_dim)]

    # conv_state: gather the decode rank's [q | k | v] shard from the three
    # independently head-sharded sub-blocks of the src tensor. dst is contiguous.
    blocks: List[Tuple[int, int, int]] = []
    src_off = 0
    dst_off = 0
    for full_sd in conv_shard_groups:
        src_sub = full_sd // src_attn_tp_size  # this prefill rank's shard of sub-block
        dst_sub = full_sd // dst_attn_tp_size  # this decode rank's shard of sub-block
        src_start = src_off + (dst_tp_rank_in_group * dst_sub) % src_sub
        blocks.append((src_start, dst_off, dst_sub))
        src_off += src_sub
        dst_off += dst_sub
    return blocks


def compute_mamba_state_slice_byte_blocks(
    *,
    src_item_len: int,
    dst_item_len: int,
    src_dim: int,
    dst_dim: int,
    outer_count: int,
    src_attn_tp_size: int,
    dst_attn_tp_size: int,
    dst_tp_rank_in_group: int,
    local_tp_rank_in_group: int,
    conv_shard_groups: Optional[List[int]] = None,
) -> List[Tuple[int, int, int]]:
    """Convert logical TP slices into physical byte blocks for one state slot.

    ``outer_count`` is one for the usual ``[slice_dim, ...]`` layout. Kimi
    conv state is ``[K - 1, slice_dim]``, so each logical channel slice expands
    into one byte block per convolution row.
    """
    src_bytes_per_dim = src_item_len // (src_dim * outer_count)
    dst_bytes_per_dim = dst_item_len // (dst_dim * outer_count)
    logical_blocks = compute_mamba_state_slice_blocks(
        src_dim=src_dim,
        dst_dim=dst_dim,
        src_attn_tp_size=src_attn_tp_size,
        dst_attn_tp_size=dst_attn_tp_size,
        dst_tp_rank_in_group=dst_tp_rank_in_group,
        local_tp_rank_in_group=local_tp_rank_in_group,
        conv_shard_groups=conv_shard_groups,
    )

    blocks = []
    for outer_idx in range(outer_count):
        src_row_offset = outer_idx * src_dim * src_bytes_per_dim
        dst_row_offset = outer_idx * dst_dim * dst_bytes_per_dim
        for src_dim_start, dst_dim_start, num_dims in logical_blocks:
            blocks.append(
                (
                    src_row_offset + src_dim_start * src_bytes_per_dim,
                    dst_row_offset + dst_dim_start * dst_bytes_per_dim,
                    num_dims * src_bytes_per_dim,
                )
            )
    return blocks


def build_transfer_entry_pairs(
    src_layer_ids: List[int],
    dst_layer_ids: List[int],
    n_src: int,
    n_dst: int,
    allow_positional_fallback: bool = False,
) -> List[Tuple[int, int]]:
    """Pair prefill-local transfer entries with decode entries by layer id."""
    if n_src == 0:
        return []
    if bool(src_layer_ids) != bool(dst_layer_ids):
        if not allow_positional_fallback:
            raise RuntimeError(
                "Layer metadata must be provided by both PD peers or neither"
            )
        src_layer_ids = []
        dst_layer_ids = []
    if src_layer_ids:
        if len(src_layer_ids) != n_src or len(dst_layer_ids) != n_dst:
            raise RuntimeError(
                "Layer metadata length must match transfer entries: "
                f"src metadata={len(src_layer_ids)} entries={n_src}, "
                f"dst metadata={len(dst_layer_ids)} entries={n_dst}"
            )
        # Layer ids can repeat across tensor groups (for example K/V or multiple
        # state tensors), so pair occurrences in order rather than by plain lookup.
        dst_pos = {}
        for j, lid in enumerate(dst_layer_ids):
            dst_pos.setdefault(lid, deque()).append(j)
        pairs = []
        for i, lid in enumerate(src_layer_ids):
            if not dst_pos.get(lid):
                raise RuntimeError(
                    f"Decode peer is missing a transfer entry for model layer {lid}"
                )
            pairs.append((i, dst_pos[lid].popleft()))
        return pairs
    if n_dst < n_src or (n_src != n_dst and not allow_positional_fallback):
        # Without layer ids a positional pairing would silently transfer the
        # wrong layers (e.g. PP prefill peered with a stale decode server).
        raise RuntimeError(
            "PP-heterogeneous transfer requires layer ids on "
            f"both peers; got src={n_src} dst={n_dst} entries"
        )
    return [(i, i) for i in range(n_src)]


def resolve_dcp_dst_entry_indices(
    src_layer_ids: List[int],
    dst_layer_ids: List[int],
    n_src: int,
    n_dst: int,
) -> List[int]:
    """Destination entry index for each local KV entry, for a DCP relayout.

    DCP re-splits the KV by context while PP re-splits it by layer, so the two
    index spaces only line up when neither peer is pipelined. Both backends
    need the same resolution, hence the shared helper.
    """
    if not src_layer_ids and not dst_layer_ids:
        # Legacy/non-PP layout. n_dst may exceed n_src when the decode side
        # runs speculative decoding and the prefill side does not.
        return list(range(n_src))
    # A one-sided mapping is rejected by build_transfer_entry_pairs itself.
    return [
        j
        for _, j in build_transfer_entry_pairs(
            src_layer_ids, dst_layer_ids, n_src, n_dst
        )
    ]


def append_state_component(
    kv_args: KVArgs,
    state_type: StateType,
    data_ptrs: List[int],
    data_lens: List[int],
    item_lens: List[int],
    dim_per_tensor: Optional[List[int]] = None,
    conv_shard_groups: Optional[List[Optional[List[int]]]] = None,
    slice_outer_counts: Optional[List[int]] = None,
    layer_ids: Optional[List[int]] = None,
) -> None:
    """Append one state component. Caller orders state_types consistently
    on prefill and decode sides."""
    kv_args.state_types.append(state_type)
    kv_args.state_data_ptrs.append(data_ptrs)
    kv_args.state_data_lens.append(data_lens)
    kv_args.state_item_lens.append(item_lens)
    kv_args.state_dim_per_tensor.append(dim_per_tensor or [])
    kv_args.state_conv_shard_groups.append(conv_shard_groups or [])
    kv_args.state_slice_outer_counts.append(slice_outer_counts or [])
    kv_args.state_layer_ids.append(layer_ids or [])


def setup_state_kv_args(
    kv_args: KVArgs,
    token_to_kv_pool,
    draft_token_to_kv_pool=None,
    total_kv_layers: int = None,
    req_to_token_pool=None,
) -> None:
    """Populate ``kv_args`` state-buffer fields from the given pool.
    Shared by prefill and decode bootstrap paths so the state_type dispatch
    lives in one place.
    """
    from sglang.srt.disaggregation.base.conn import StateType
    from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool
    from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
    from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
    from sglang.srt.mem_cache.memory_pool import (
        DSATokenToKVPool,
        HybridLinearKVPool,
        MiniMaxSparseKVPool,
    )

    kv_args.state_types = []
    kv_args.state_data_ptrs = []
    kv_args.state_data_lens = []
    kv_args.state_item_lens = []
    kv_args.state_dim_per_tensor = []
    kv_args.state_slice_outer_counts = []
    kv_args.state_layer_ids = []
    kv_args.is_hybrid_mla_backend = False
    kv_args.state_conv_shard_groups = []

    if is_npu() and isinstance(token_to_kv_pool, DSV4NPUTokenToKVPool):
        # Pool ships each sub-pool as its own page-indexed component (fixed order
        # so prefill and decode register identically); skips get_state_buf_infos.
        for (
            st,
            comp_ptrs,
            comp_lens,
            comp_item_lens,
        ) in token_to_kv_pool.get_pd_state_components():
            append_state_component(kv_args, st, comp_ptrs, comp_lens, comp_item_lens)
    elif isinstance(token_to_kv_pool, MiniMaxSparseKVPool):
        if token_to_kv_pool.index_kv_pool is not None:
            raise NotImplementedError(
                "PD disaggregation for MiniMax sparse layers with index value "
                "(index_kv_pool) is not yet supported; only K-only sparse layers are."
            )
        if token_to_kv_pool.index_k_pool is not None:
            dp, dl, il = token_to_kv_pool.get_index_k_state_buf_infos()
            append_state_component(kv_args, StateType.MINIMAX_INDEX_K, dp, dl, il)
    elif hasattr(token_to_kv_pool, "get_state_buf_infos"):
        data_ptrs, data_lens, item_lens = token_to_kv_pool.get_state_buf_infos()

        # DeepSeekV4TokenToKVPool inherits BaseSWAKVPool; its heterogeneous
        # state list is described per-entry via get_state_buf_infos.
        if isinstance(token_to_kv_pool, BaseSWAKVPool):
            append_state_component(
                kv_args, StateType.SWA, data_ptrs, data_lens, item_lens
            )
            # unified_kv: the SWA ring lives in the unified buffers (no separate
            # swa_kv_pool) and is addressed per-row, so ship it as SWA_RING.
            if getattr(token_to_kv_pool, "_unified_kv", False) and hasattr(
                token_to_kv_pool, "get_unified_swa_ring_buf_infos"
            ):
                ring_ptrs, ring_lens, ring_item_lens = (
                    token_to_kv_pool.get_unified_swa_ring_buf_infos()
                )
                if ring_ptrs:
                    append_state_component(
                        kv_args,
                        StateType.SWA_RING,
                        ring_ptrs,
                        ring_lens,
                        ring_item_lens,
                    )
            if hasattr(token_to_kv_pool, "get_c128_state_buf_infos"):
                c128_ptrs, c128_lens, c128_item_lens = (
                    token_to_kv_pool.get_c128_state_buf_infos()
                )
                if c128_ptrs:
                    append_state_component(
                        kv_args,
                        StateType.C128_STATE,
                        c128_ptrs,
                        c128_lens,
                        c128_item_lens,
                    )
        elif isinstance(token_to_kv_pool, HybridLinearKVPool):
            dim = (
                token_to_kv_pool.get_state_dim_per_tensor()
                if hasattr(token_to_kv_pool, "get_state_dim_per_tensor")
                else None
            )
            kv_args.is_hybrid_mla_backend = is_mla_backend(
                token_to_kv_pool.full_kv_pool
            )
            conv_shard_groups = (
                token_to_kv_pool.get_state_conv_shard_groups()
                if hasattr(token_to_kv_pool, "get_state_conv_shard_groups")
                else None
            )
            slice_outer_counts = (
                token_to_kv_pool.get_state_slice_outer_counts()
                if hasattr(token_to_kv_pool, "get_state_slice_outer_counts")
                else None
            )
            # Global layer ids let the sender pair src/dst entries when the
            # prefill PP stage registers only its own subset of mamba layers.
            layer_ids = token_to_kv_pool.get_state_layer_ids()
            append_state_component(
                kv_args,
                StateType.MAMBA,
                data_ptrs,
                data_lens,
                item_lens,
                dim,
                conv_shard_groups,
                slice_outer_counts,
                layer_ids,
            )
        elif isinstance(token_to_kv_pool, (DSATokenToKVPool, NPUMLATokenToKVPool)):
            if draft_token_to_kv_pool is not None and isinstance(
                draft_token_to_kv_pool, DSATokenToKVPool
            ):
                (
                    draft_data_ptrs,
                    draft_data_lens,
                    draft_item_lens,
                ) = draft_token_to_kv_pool.get_state_buf_infos()
                data_ptrs = data_ptrs + draft_data_ptrs
                data_lens = data_lens + draft_data_lens
                item_lens = item_lens + draft_item_lens
            if isinstance(token_to_kv_pool, NPUMLATokenToKVPool):
                kv_args.kv_buf_groups = (
                    len(kv_args.kv_data_ptrs) // token_to_kv_pool.layer_num
                )
                kv_args.total_kv_layers = total_kv_layers
            else:
                append_state_component(
                    kv_args, StateType.DSA, data_ptrs, data_lens, item_lens
                )

    # DSV4 NextN shares the target allocator, so target and draft use the same
    # local SWA indices. Keep draft buffers in a separate positional component
    # to avoid mixing them into the target's heterogeneous state layout, while
    # reusing the existing SWA transport dispatch. NPU has a different paged
    # state layout and is intentionally left unchanged.
    if (
        not is_npu()
        and isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
        and isinstance(draft_token_to_kv_pool, DeepSeekV4TokenToKVPool)
    ):
        if not draft_token_to_kv_pool.compression_ratios or not all(
            ratio == 0 for ratio in draft_token_to_kv_pool.compression_ratios
        ):
            raise RuntimeError(
                "DSV4 draft state transfer expects SWA-only NextN layers"
            )
        if token_to_kv_pool._unified_kv != draft_token_to_kv_pool._unified_kv:
            raise RuntimeError(
                "DSV4 target and draft pools must use the same unified-KV mode"
            )

        if token_to_kv_pool._unified_kv:
            target_geometry = (
                token_to_kv_pool.unified_swa_window,
                token_to_kv_pool.unified_swa_ring_size,
                token_to_kv_pool.unified_swa_pages,
            )
            draft_geometry = (
                draft_token_to_kv_pool.unified_swa_window,
                draft_token_to_kv_pool.unified_swa_ring_size,
                draft_token_to_kv_pool.unified_swa_pages,
            )
            if target_geometry != draft_geometry:
                raise RuntimeError(
                    "DSV4 target and draft pools must share SWA ring geometry: "
                    f"target={target_geometry}, draft={draft_geometry}"
                )
            draft_ptrs, draft_lens, draft_item_lens = (
                draft_token_to_kv_pool.get_unified_swa_ring_buf_infos()
            )
            draft_state_type = StateType.SWA_RING
        else:
            if (
                token_to_kv_pool.full_to_swa_index_mapping
                is not draft_token_to_kv_pool.full_to_swa_index_mapping
            ):
                raise RuntimeError(
                    "DSV4 target and draft pools must share the SWA index mapping"
                )
            target_geometry = (
                token_to_kv_pool.page_size,
                token_to_kv_pool.sliding_window,
            )
            draft_geometry = (
                draft_token_to_kv_pool.page_size,
                draft_token_to_kv_pool.sliding_window,
            )
            if target_geometry != draft_geometry:
                raise RuntimeError(
                    "DSV4 target and draft pools must share paged SWA geometry: "
                    f"target={target_geometry}, draft={draft_geometry}"
                )
            draft_ptrs, draft_lens, draft_item_lens = (
                draft_token_to_kv_pool.get_state_buf_infos()
            )
            draft_state_type = StateType.SWA

        if draft_ptrs:
            append_state_component(
                kv_args,
                draft_state_type,
                draft_ptrs,
                draft_lens,
                draft_item_lens,
            )

    if (
        StateType.MAMBA not in kv_args.state_types
        and req_to_token_pool is not None
        and hasattr(req_to_token_pool, "get_state_buf_infos")
    ):
        data_ptrs, data_lens, item_lens = req_to_token_pool.get_state_buf_infos()
        if data_ptrs:
            dim = (
                req_to_token_pool.get_state_dim_per_tensor()
                if hasattr(req_to_token_pool, "get_state_dim_per_tensor")
                else None
            )
            conv_shard_groups = (
                req_to_token_pool.get_state_conv_shard_groups()
                if hasattr(req_to_token_pool, "get_state_conv_shard_groups")
                else None
            )
            slice_outer_counts = (
                req_to_token_pool.get_state_slice_outer_counts()
                if hasattr(req_to_token_pool, "get_state_slice_outer_counts")
                else None
            )
            append_state_component(
                kv_args,
                StateType.MAMBA,
                data_ptrs,
                data_lens,
                item_lens,
                dim,
                conv_shard_groups,
                slice_outer_counts,
            )


def prepare_abort(req: Req, error_message: str, status_code=None):
    from sglang.srt.managers.schedule_batch import FINISH_ABORT

    # populate finish metadata and stream output
    req.finished_reason = FINISH_ABORT(error_message, status_code)

    if req.return_logprob:
        req.logprob.input_token_logprobs_val = []
        req.logprob.input_token_logprobs_idx = []
        req.logprob.input_top_logprobs_val = []
        req.logprob.input_top_logprobs_idx = []
        req.logprob.input_token_ids_logprobs_val = []
        req.logprob.input_token_ids_logprobs_idx = []


def is_aborted(req: Req) -> bool:
    from sglang.srt.managers.schedule_batch import FINISH_ABORT

    return isinstance(req.to_finish, FINISH_ABORT) or isinstance(
        req.finished_reason, FINISH_ABORT
    )
