from dataclasses import dataclass
from enum import Enum
from typing import Iterable, List, Optional, Set, Tuple, Union

import torch

from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.utils.hf_transformers_utils import AutoConfig


@dataclass
class MoELoRABatchInfo:
    # Per-request segment indptrs used by MoE LoRA routing, shape (bs + 1,).
    seg_indptr: torch.Tensor

    # Per-request adapter index used by MoE LoRA routing, shape (bs,).
    req_to_lora: torch.Tensor

    # A mask indicating if lora adapter is enabled. Shape (num_loras,)
    adapter_enabled: torch.Tensor

    # A mapping of which lora adapter is used for each token. Shape (num_tokens,)
    # If a token has no lora adapter, the value is -1.
    token_lora_mapping: torch.Tensor


@dataclass
class LoRABatchInfo:
    # The forward mode is using CUDA Graph.
    use_cuda_graph: bool

    # Batch size
    bs: int

    # Number of segments. For triton backend, it is equal to batch size.
    num_segments: int

    # Indice pointers of each segment in shape (num_segments + 1, )
    seg_indptr: torch.Tensor

    # The index of lora adapter used by each segment, in shape (num_segments,)
    weight_indices: torch.Tensor

    # ranks of each lora adapter, in shape (lora_num,)
    lora_ranks: torch.Tensor

    # scaling of each lora adapter, in shape (lora_num,)
    scalings: torch.Tensor

    # Maximum segment length of current batch
    max_len: Optional[int]

    # Lengths of each segments in shape (num_segments,)
    seg_lens: Optional[torch.Tensor]

    # The logical (re)ordering of input rows (tokens), in shape (num_tokens,)
    permutation: Optional[torch.Tensor]

    # Total number of tokens this batch info expects (host-side int).
    # Used by lm_head LoRA to validate input shape without GPU sync.
    expected_tokens: Optional[int] = None

    # CPU-side flag: True when at least one request uses a LoRA adapter.
    # Computed from Python lists in prepare_lora_batch to avoid GPU sync.
    has_active_lora: bool = False

    # Per-request segment indptrs, shape (bs + 1,). Required by MoE virtual
    # experts which map tokens to requests regardless of the dense-LoRA
    # backend's internal segmentation.  For the triton backend these are
    # identical to seg_indptr/weight_indices; for csgmv they differ because
    # its segments are chunked across adapters.
    req_seg_indptr: Optional[torch.Tensor] = None

    # Per-request adapter index, shape (bs,).
    req_weight_indices: Optional[torch.Tensor] = None

    # MoE LoRA batch info
    moe_lora_info: Optional[MoELoRABatchInfo] = None


class LoRAType(Enum):
    LORA_A = 0
    LORA_B = 1


def copy_weight_into_buffer(
    buffer_view: torch.Tensor,
    weight: torch.Tensor,
) -> None:
    """
    Copy a LoRA weight tensor into a destination buffer.

    When a pinned CPU source has a dtype mismatch with a device destination,
    cast on the destination device instead of doing the conversion on CPU.
    """
    if weight.dtype == buffer_view.dtype:
        buffer_view.copy_(weight, non_blocking=True)
        return

    if weight.device.type == "cpu" and buffer_view.device.type != "cpu":
        weight = weight.to(device=buffer_view.device, non_blocking=True)

    buffer_view.copy_(weight.to(dtype=buffer_view.dtype), non_blocking=True)


def get_hidden_dim(
    module_name: str,
    config: AutoConfig,
    base_model: torch.nn.Module,
    layer_idx: int,
    lora_added_vocab_size: int = 0,
) -> Tuple[int]:
    """
    Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
    """

    if hasattr(base_model, "get_hidden_dim"):
        return base_model.get_hidden_dim(module_name, layer_idx)
    else:
        """
        WARNING: get_hidden_dim() is not defined,
        which is used to get the hidden dim for different lora modules
        Use the default one, but please check if it is correct for your model.
        Please implement the function in the model class if it is not.
        You can reference this function in llama.py.
        """
        head_dim = getattr(
            config, "head_dim", config.hidden_size // config.num_attention_heads
        )
        if module_name == "qkv_proj":
            return config.hidden_size, head_dim * (
                config.num_attention_heads + config.num_key_value_heads * 2
            )
        elif module_name == "o_proj":
            o_head_dim = getattr(config, "v_head_dim", None) or head_dim
            return (
                o_head_dim * config.num_attention_heads,
                config.hidden_size,
            )
        elif module_name == "gate_up_proj":
            inter = config.intermediate_size
            first_k = getattr(config, "first_k_dense_replace", None)
            moe_freq = getattr(config, "moe_layer_freq", 1)
            if (
                first_k is not None
                and layer_idx >= first_k
                and layer_idx % moe_freq == 0
            ):
                moe_inter = getattr(config, "moe_intermediate_size", None)
                n_shared = getattr(config, "n_shared_experts", None)
                if moe_inter is not None and n_shared is not None:
                    inter = moe_inter * n_shared
            return config.hidden_size, inter * 2
        elif module_name == "down_proj":
            inter = config.intermediate_size
            first_k = getattr(config, "first_k_dense_replace", None)
            moe_freq = getattr(config, "moe_layer_freq", 1)
            if (
                first_k is not None
                and layer_idx >= first_k
                and layer_idx % moe_freq == 0
            ):
                moe_inter = getattr(config, "moe_intermediate_size", None)
                n_shared = getattr(config, "n_shared_experts", None)
                if moe_inter is not None and n_shared is not None:
                    inter = moe_inter * n_shared
            return inter, config.hidden_size
        elif module_name == "fused_qkv_a_proj_with_mqa":
            q_lora_rank = getattr(config, "q_lora_rank", None) or 0
            kv_lora_rank = config.kv_lora_rank
            qk_rope_head_dim = config.qk_rope_head_dim
            return (
                config.hidden_size,
                q_lora_rank + kv_lora_rank + qk_rope_head_dim,
            )
        elif module_name == "q_b_proj":
            return (
                config.q_lora_rank,
                config.num_attention_heads
                * (config.qk_nope_head_dim + config.qk_rope_head_dim),
            )
        elif module_name == "kv_b_proj":
            return (
                config.kv_lora_rank,
                config.num_attention_heads
                * (config.qk_nope_head_dim + config.v_head_dim),
            )
        elif module_name in DSA_INDEXER_LORA_NAMES:
            from sglang.srt.configs.model_config import (
                get_dsa_index_head_dim,
                get_dsa_index_n_heads,
            )

            if module_name == "indexer.wq_b":
                return (
                    config.q_lora_rank,
                    get_dsa_index_n_heads(config) * get_dsa_index_head_dim(config),
                )
            elif module_name == "indexer.wk":
                return config.hidden_size, get_dsa_index_head_dim(config)
            else:  # indexer.weights_proj
                return config.hidden_size, get_dsa_index_n_heads(config)
        elif module_name == "gate_up_proj_moe":
            moe_inter = (
                getattr(config, "moe_intermediate_size", None)
                or config.intermediate_size
            )
            return config.hidden_size, moe_inter * 2
        elif module_name == "down_proj_moe":
            moe_inter = (
                getattr(config, "moe_intermediate_size", None)
                or config.intermediate_size
            )
            return moe_inter, config.hidden_size
        elif module_name == "embed_tokens":
            # For embedding: input is vocab_size (as embedding lookup), output is hidden_size
            # if contain extra tokens will be added; otherwise is 0.
            return config.vocab_size + lora_added_vocab_size, config.hidden_size
        elif module_name == "lm_head":
            # For lm_head: input is hidden_size, output is vocab_size
            # if contain extra tokens will be added; otherwise is 0.
            return config.hidden_size, config.vocab_size + lora_added_vocab_size
        else:
            raise NotImplementedError(
                "get_hidden_dim not implemented for " + module_name
            )


def get_normalized_target_modules(
    target_modules: Union[str, Iterable[str]],
) -> set[str]:
    """
    Mapping a list of target module name to names of the normalized LoRA weights.
    Handles both base module names (e.g., "gate_proj") and prefixed module names (e.g., "feed_forward.gate_proj").

    Also handles PEFT shorthand strings like "all-linear" or "all" by returning
    {"all"} as a sentinel value.  Callers that need a concrete module set
    should use :func:`auto_detect_lora_target_modules` to resolve the shorthand
    against the loaded base model.
    """
    # Handle PEFT shorthand strings — return {"all"} as sentinel.
    # Callers can resolve to concrete names via auto_detect_lora_target_modules().
    if isinstance(target_modules, str):
        if target_modules not in ["all", "all-linear"]:
            raise ValueError(
                "Only 'all' or 'all-linear' can be used as the string for target module"
            )
        return {"all"}

    params_mapping = {
        "q_proj": "qkv_proj",
        "k_proj": "qkv_proj",
        "v_proj": "qkv_proj",
        "gate_proj": "gate_up_proj",
        "up_proj": "gate_up_proj",
        "out_proj": "out_proj",
        "embed_tokens": "embed_tokens",
        "vocab_emb": "embed_tokens",
        "embeddings": "embed_tokens",
        "word_embeddings": "embed_tokens",
        "lm_head": "lm_head",
        "output": "lm_head",
        "unembed_tokens": "lm_head",
        "q_a_proj": "fused_qkv_a_proj_with_mqa",
        "kv_a_proj_with_mqa": "fused_qkv_a_proj_with_mqa",
        # DSA indexer projections are qualified with their parent module name
        # because the bare leaf names collide with unrelated modules in other
        # models (e.g. DeepSeek-V4 attention `wq_b`, Pixtral vision `wk`).
        "wq_b": "indexer.wq_b",
        "wk": "indexer.wk",
        "weights_proj": "indexer.weights_proj",
    }

    result = set()
    for name in target_modules:
        base_name = name.split(".")[-1]
        normalized_name = params_mapping.get(base_name, base_name)
        result.add(normalized_name)
    return result


def get_stacked_multiply(
    module_name: str, base_model: Optional[torch.nn.Module] = None
) -> int:
    """
    Mapping a lora module name to its magnification at output dimension.
    Models can override via a get_stacked_multiply(module_name) method.
    """
    if base_model is not None and hasattr(base_model, "get_stacked_multiply"):
        return base_model.get_stacked_multiply(module_name)
    stacked_rank = {
        "qkv_proj": 3,
        "in_proj_qkvz": 4,  # GDN packed input projection
        "gate_up_proj": 2,
        "gate_up_proj_moe": 2,
        "in_proj": 2,
        "fused_qkv_a_proj_with_mqa": 2,
    }
    return stacked_rank[module_name] if module_name in stacked_rank else 1


def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> str:
    """
    Get the target module name in target_modules that can match full_module_name.

    If there is a target module name in target_modules that can match full_module_name, return this name
    Else raise ValueError.

    When multiple target modules match (e.g. both "up_proj" and "gate_up_proj"
    are substrings), the longest match wins to avoid ambiguity.
    """
    best = None
    for target_module in target_modules:
        if target_module in full_module_name:
            if best is None or len(target_module) > len(best):
                best = target_module
    if best is not None:
        return best
    raise ValueError(
        f"Cannot find target module name for {full_module_name} in {target_modules}"
    )


EMBEDDING_NAMES = ["embed_tokens", "lm_head"]
ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "out_proj", "down_proj", "down_proj_moe"]
DSA_INDEXER_LORA_NAMES = frozenset(
    {"indexer.wq_b", "indexer.wk", "indexer.weights_proj"}
)
REPLICATED_LINEAR_LORA_NAMES = [
    "fused_qkv_a_proj_with_mqa",
    "fc1_latent_proj",
    "fc2_latent_proj",
    *DSA_INDEXER_LORA_NAMES,
]

# Normalized module names that the LoRA system fully supports
# (i.e. get_hidden_dim, init_buffers, and init_lora_modules can handle them).
_KNOWN_LORA_TARGET_MODULES = frozenset(
    {
        "qkv_proj",
        "o_proj",
        "out_proj",
        "in_proj",
        "in_proj_qkvz",
        "up_proj",
        "gate_up_proj",
        "down_proj",
        "fc1_latent_proj",
        "fc2_latent_proj",
        "embed_tokens",
        "lm_head",
        "fused_qkv_a_proj_with_mqa",
        "q_b_proj",
        "kv_b_proj",
    }
    | DSA_INDEXER_LORA_NAMES
)


def auto_detect_lora_target_modules(model: "torch.nn.Module") -> set:
    """Discover LoRA-compatible modules by inspecting the base model.

    Walks the model graph and returns the set of *normalized* target-module
    names that (a) actually exist in the model and (b) the LoRA memory pool
    can handle.  This is used to resolve PEFT shorthands like ``"all-linear"``
    without requiring the user to enumerate modules on the CLI.
    """
    from sglang.srt.layers.linear import LinearBase
    from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
    from sglang.srt.layers.vocab_parallel_embedding import (
        ParallelLMHead,
        VocabParallelEmbedding,
    )

    raw_names: set = set()
    dsa_indexer_leaf_names = {
        target_name.split(".")[-1] for target_name in DSA_INDEXER_LORA_NAMES
    }
    for name, module in model.named_modules():
        if isinstance(module, FusedMoE):
            raw_names.add("gate_up_proj")
            raw_names.add("down_proj")
        elif isinstance(module, ParallelLMHead):
            raw_names.add("lm_head")
        elif isinstance(module, VocabParallelEmbedding):
            raw_names.add("embed_tokens")
        elif isinstance(module, LinearBase):
            parts = name.split(".")
            leaf_name = parts[-1]
            parent_qualified_name = ".".join(parts[-2:])
            if parent_qualified_name in DSA_INDEXER_LORA_NAMES:
                raw_names.add(parent_qualified_name)
            elif leaf_name in dsa_indexer_leaf_names:
                # Bare DSA indexer leaf names are ambiguous across model
                # families. Only auto-detect them when the actual module path
                # proves they are under an `indexer` parent.
                continue
            else:
                raw_names.add(leaf_name)

    normalized = get_normalized_target_modules(raw_names)
    result = normalized & _KNOWN_LORA_TARGET_MODULES

    # Allow models to declare additional LoRA-compatible modules that
    # cannot be auto-discovered or need to bypass normalization
    # (e.g. Mamba in_proj, non-gated up_proj).
    if hasattr(model, "supported_lora_modules"):
        result.update(set(model.supported_lora_modules) & _KNOWN_LORA_TARGET_MODULES)

    return result


def get_lm_head_lora_b_shard_size(output_dim: int, shard_indices=None) -> int:
    """Get the LoRA B output dimension for lm_head, accounting for TP.

    lm_head is column-parallel, so its LoRA B must be sharded along the
    vocab dimension to match the base output.  When shard_indices is
    provided, the returned size reflects the base model's actual per-rank
    vocab partition.

    Args:
        output_dim: Full (unsharded) output dimension (vocab_size).
        shard_indices: VocabParallelEmbeddingShardIndices from the base
            ParallelLMHead layer.  When provided, returns the per-rank
            org vocab size from the base model's actual sharding.
    """
    if shard_indices is not None:
        return shard_indices.num_org_elements
    return output_dim


def generate_sequence_lengths(
    forward_batch: ForwardBatch, device: Optional[torch.device] = None
) -> torch.Tensor:

    device = torch.get_default_device() if device is None else device
    with torch.device(device):
        if forward_batch.forward_mode.is_decode():
            seg_lens = torch.ones(forward_batch.batch_size, dtype=torch.int32)
        elif forward_batch.forward_mode.is_target_verify():
            seg_lens = torch.full(
                size=(forward_batch.batch_size,),
                fill_value=forward_batch.spec_info.draft_token_num,
                dtype=torch.int32,
            )
        elif forward_batch.forward_mode.is_extend():
            seg_lens = (
                forward_batch.extend_seq_lens
                if forward_batch.extend_seq_lens.device == device
                else torch.tensor(
                    forward_batch.extend_seq_lens_cpu,
                    dtype=torch.int32,
                )
            )
        else:
            raise ValueError(f"Unsupported forward mode: {forward_batch.forward_mode}")
    return seg_lens


def get_lm_head_pruned_lens(
    forward_batch: ForwardBatch,
) -> Optional[List[int]]:
    """
    Compute per-sequence pruned lengths for lm_head LoRA.

    Returns a list of pruned lengths (one per sequence) if pruning applies,
    or None if lm_head pruning is not applicable for this batch.

    Pruning rules:
    - Extend without logprobs: 1 token per sequence
    - Extend with logprobs: max(extend_len - logprob_start_len, 1) per sequence
    - Decode / target_verify / draft_extend_v2: no pruning

    IMPORTANT: This must stay in sync with LogitsProcessor._get_pruned_states()
    in sglang/srt/layers/logits_processor.py, which determines how many tokens
    per sequence are passed to lm_head. If the pruning conditions or lengths
    there change, this function must be updated to match, otherwise the
    lm_head LoRA will operate on incorrectly shaped inputs.
    """
    lm_head_pruning = (
        forward_batch.forward_mode.is_extend()
        and not forward_batch.forward_mode.is_target_verify()
        and not forward_batch.forward_mode.is_draft_extend_v2()
    )

    if not lm_head_pruning:
        return None

    if forward_batch.return_logprob:
        pruned_lens = []
        for ext_len, start_len in zip(
            forward_batch.extend_seq_lens_cpu,
            forward_batch.extend_logprob_start_lens_cpu,
        ):
            pruned_lens.append(1 if ext_len == start_len else ext_len - start_len)
    else:
        pruned_lens = [1] * forward_batch.batch_size

    return pruned_lens


def merge_and_chunk_segments(
    weight_indices: list[int],
    pruned_lens: List[int],
    chunk_size: int,
) -> Tuple[List[int], List[int]]:
    """
    Merge consecutive same-adapter sequences and chunk at chunk_size boundaries.

    Merges consecutive sequences that use the same adapter into single
    segments, splitting any segment that exceeds chunk_size.

    Args:
        weight_indices: Per-sequence adapter indices.
        pruned_lens: Per-sequence pruned token counts.
        chunk_size: Maximum segment length before splitting.

    Returns:
        (seg_weight_indices, seg_lens): Merged and chunked segments.
    """
    seg_weight_indices: List[int] = []
    seg_lens: List[int] = []
    for wi, pl in zip(weight_indices, pruned_lens):
        if seg_weight_indices and seg_weight_indices[-1] == wi:
            seg_lens[-1] += pl
        else:
            seg_weight_indices.append(wi)
            seg_lens.append(pl)
        # Split the last segment if it exceeds chunk_size
        while seg_lens[-1] > chunk_size:
            remainder = seg_lens[-1] - chunk_size
            seg_lens[-1] = chunk_size
            seg_weight_indices.append(wi)
            seg_lens.append(remainder)

    return seg_weight_indices, seg_lens


def build_lm_head_pass_segments(
    weight_indices: List[int],
    pruned_lens: List[int],
    logprobs_chunk_size: int,
) -> List[Tuple[List[int], List[int]]]:
    """
    Precompute per-pass segment info for lm_head LoRA logprobs processing.

    When LogitsProcessor uses chunked logprobs processing
    (process_input_logprobs_by_chunk), pruned hidden states are split into
    fixed-size passes.  Each pass needs its own segmentation
    (weight_indices, seg_lens) so that lm_head LoRA operates on the
    correct adapter assignments per pass.

    Args:
        weight_indices: Per-sequence adapter indices.
        pruned_lens: Per-sequence pruned token counts.
        logprobs_chunk_size: Fixed pass size used by LogitsProcessor.

    Returns:
        List of (seg_weight_indices, seg_lens) tuples, one per pass.
    """
    # Expand to per-token weight index
    token_wi: List[int] = []
    for wi, pl in zip(weight_indices, pruned_lens):
        token_wi.extend([wi] * pl)
    total = len(token_wi)
    num_passes = (total + logprobs_chunk_size - 1) // logprobs_chunk_size

    result: List[Tuple[List[int], List[int]]] = []
    for i in range(num_passes):
        start = i * logprobs_chunk_size
        end = min((i + 1) * logprobs_chunk_size, total)

        # Run-length encode the pass's adapter indices
        seg_wi: List[int] = []
        seg_lens: List[int] = []
        for t in range(start, end):
            if seg_wi and seg_wi[-1] == token_wi[t]:
                seg_lens[-1] += 1
            else:
                seg_wi.append(token_wi[t])
                seg_lens.append(1)
        result.append((seg_wi, seg_lens))

    return result