# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/model_loader/utils.py

"""Utilities for selecting and loading models."""

import concurrent.futures
import contextlib
import logging
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type

import torch
import transformers
from torch import nn
from transformers.dynamic_module_utils import get_class_from_dynamic_module

from sglang.srt.configs.model_config import ModelConfig, ModelImpl
from sglang.srt.layers import deep_gemm_wrapper

logger = logging.getLogger(__name__)


@contextlib.contextmanager
def set_default_torch_dtype(dtype: torch.dtype):
    """Sets the default torch dtype to the given dtype."""
    old_dtype = torch.get_default_dtype()
    torch.set_default_dtype(dtype)
    yield
    torch.set_default_dtype(old_dtype)


def _is_moe_model(model_config: ModelConfig, architectures: list[str]) -> bool:
    lowered_arches = [arch.lower() for arch in architectures]
    if any("moe" in arch or "mixtral" in arch for arch in lowered_arches):
        return True

    text_config = model_config.hf_text_config
    expert_attrs = (
        "num_local_experts",
        "num_experts",
        "num_experts_per_tok",
        "moe_intermediate_size",
        "n_routed_experts",
    )
    for attr in expert_attrs:
        value = getattr(text_config, attr, None)
        if value is None:
            continue
        if isinstance(value, bool):
            if value:
                return True
            continue
        if isinstance(value, (int, float)):
            threshold = 0 if attr == "moe_intermediate_size" else 1
            if value > threshold:
                return True
            continue
        if isinstance(value, (list, tuple, set, dict)):
            if len(value) > 0:
                return True
            continue
        if isinstance(value, str) and value == "":
            continue
        if value is not None:
            return True
    return False


def _is_sequence_classification_model(architectures: list[str]) -> bool:
    return any(
        "sequenceclassification" in lowered or "rewardmodel" in lowered
        for lowered in (arch.lower() for arch in architectures)
    )


def _get_transformers_backend_arch(
    model_config: ModelConfig, architectures: list[str]
) -> str:
    is_pooling = not model_config.is_generation
    is_multimodal = model_config.is_multimodal or (
        model_config.hf_config is not model_config.hf_text_config
    )
    is_moe = _is_moe_model(model_config, architectures)
    base_arch = "ForCausalLM"
    if is_pooling:
        base_arch = (
            "ForSequenceClassification"
            if _is_sequence_classification_model(architectures)
            else "EmbeddingModel"
        )

    arch = "Transformers"
    if is_multimodal:
        arch += "MultiModal"
    if is_moe:
        arch += "MoE"
    return arch + base_arch


def _model_impl_from_architecture(architecture: str) -> ModelImpl:
    if architecture.startswith("Transformers"):
        return ModelImpl.TRANSFORMERS
    if architecture.startswith("MindSpore"):
        return ModelImpl.MINDSPORE
    return ModelImpl.SGLANG


def resolve_transformers_arch(model_config: ModelConfig, architectures: list[str]):
    backend_arch = _get_transformers_backend_arch(model_config, architectures)

    for arch in architectures:
        if arch.startswith("Transformers"):
            continue
        auto_map: dict[str, str] = (
            getattr(model_config.hf_config, "auto_map", None) or dict()
        )
        # Make sure that config class is always initialized before model class,
        # otherwise the model class won't be able to access the config class,
        # the expected auto_map should have correct order like:
        # "auto_map": {
        #     "AutoConfig": "<your-repo-name>--<config-name>",
        #     "AutoModel": "<your-repo-name>--<config-name>",
        #     "AutoModelFor<Task>": "<your-repo-name>--<config-name>",
        # },
        auto_modules = {}
        try:
            auto_modules = {
                name: get_class_from_dynamic_module(
                    module, model_config.model_path, revision=model_config.revision
                )
                for name, module in sorted(auto_map.items(), key=lambda x: x[0])
            }
        except Exception as e:
            logger.warning(
                "Failed to load dynamic modules from auto_map for '%s': %s. "
                "Skipping remote model compatibility checks.",
                arch,
                e,
            )
        model_module = getattr(transformers, arch, None)
        if model_module is None:
            has_auto_model = "AutoModel" in auto_modules
            if not has_auto_model and model_config.model_impl == ModelImpl.TRANSFORMERS:
                logger.warning(
                    "Cannot resolve model class for '%s' and no auto_map.AutoModel "
                    "is present. Skipping compatibility gate because "
                    "--model-impl=transformers is explicitly requested.",
                    arch,
                )
                continue
            if not has_auto_model and "AutoModel" not in auto_map:
                raise ValueError(
                    f"Cannot find model module. '{arch}' is not a registered "
                    "model in the Transformers library (only relevant if the "
                    "model is meant to be in Transformers) and 'AutoModel' is "
                    "not present in the model config's 'auto_map' (relevant "
                    "if the model is custom)."
                )
            if not has_auto_model:
                raise ValueError(
                    f"Cannot find model module. '{arch}' is not a registered "
                    "model in the Transformers library and loading the custom "
                    f"model from auto_map failed. The remote model code may be "
                    f"incompatible with the installed transformers version."
                )
            model_module = auto_modules["AutoModel"]
        if model_config.model_impl == ModelImpl.TRANSFORMERS:
            if hasattr(model_module, "is_backend_compatible") and (
                not model_module.is_backend_compatible()
            ):
                logger.warning(
                    "The Transformers implementation of %s reports it is not "
                    "backend-compatible (_supports_attention_backend=False). "
                    "Proceeding anyway because --model-impl=transformers was "
                    "explicitly requested. The model may not work correctly.",
                    arch,
                )
        if model_config.model_impl == ModelImpl.AUTO:
            if hasattr(model_module, "is_backend_compatible") and (
                not model_module.is_backend_compatible()
            ):
                raise ValueError(
                    f"{arch} has no SGlang implementation and the Transformers "
                    "implementation is not compatible with SGLang."
                )
            logger.warning(
                "%s has no SGLang implementation, falling back to Transformers "
                "implementation. Some features may not be supported and "
                "performance may not be optimal.",
                arch,
            )
    return [backend_arch]


def get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module], str]:
    from sglang.srt.models.registry import ModelRegistry

    architectures = getattr(model_config.hf_config, "architectures", [])
    # Special handling for quantized Mixtral.
    # FIXME(woosuk): This is a temporary hack.
    mixtral_supported = [
        "fp8",
        "compressed-tensors",
        "gptq_marlin",
        "awq_marlin",
        "quark_int4fp8_moe",
    ]

    if (
        model_config.quantization is not None
        and model_config.quantization not in mixtral_supported
        and "MixtralForCausalLM" in architectures
    ):
        architectures = ["QuantMixtralForCausalLM"]

    supported_archs = ModelRegistry.get_supported_archs()
    is_native_supported = any(arch in supported_archs for arch in architectures)

    if model_config.model_impl == ModelImpl.MINDSPORE:
        architectures = ["MindSporeForCausalLM"]
    elif not is_native_supported or model_config.model_impl == ModelImpl.TRANSFORMERS:
        architectures = resolve_transformers_arch(model_config, architectures)
    model_cls, resolved_arch = ModelRegistry.resolve_model_cls(architectures)
    setattr(model_config, "_resolved_model_arch", resolved_arch)
    setattr(
        model_config,
        "_resolved_model_impl",
        _model_impl_from_architecture(resolved_arch),
    )
    return model_cls, resolved_arch


def get_resolved_model_impl(model_config: ModelConfig) -> ModelImpl:
    resolved_model_impl = getattr(model_config, "_resolved_model_impl", None)
    if resolved_model_impl is not None:
        return resolved_model_impl

    resolved_arch = getattr(model_config, "_resolved_model_arch", None)
    if resolved_arch is None:
        _, resolved_arch = get_model_architecture(model_config)

    resolved_model_impl = _model_impl_from_architecture(resolved_arch)
    setattr(model_config, "_resolved_model_arch", resolved_arch)
    setattr(model_config, "_resolved_model_impl", resolved_model_impl)
    return resolved_model_impl


def get_architecture_class_name(model_config: ModelConfig) -> str:
    return get_model_architecture(model_config)[1]


def should_deepgemm_weight_requant_ue8m0(
    weight_block_size, output_dtype=None, weight_shape=None
):
    """Should we requant fp8 weights into UE8M0 format when loading the model.

    When output_dtype or weight_shape are provided, also checks that DeepGEMM
    can actually run this layer at runtime (bf16 output, N%64==0, K%128==0).
    Without these checks, scales would be converted to UE8M0 but the GEMM would
    fall back to triton which expects float32 scales, causing wrong results.
    """
    if not (
        deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
        and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
        and weight_block_size is not None
    ):
        return False
    if output_dtype is not None and output_dtype != torch.bfloat16:
        return False
    if weight_shape is not None and (
        weight_shape[0] % 64 != 0 or weight_shape[1] % 128 != 0
    ):
        return False
    return True


def should_async_load(weight: torch.Tensor) -> bool:
    """Return True if we should load the given weight asynchronously.

    For host (CPU) tensors, using a threadpool can overlap H2D copies
    and improve throughput. For device tensors, threading often adds overhead
    (e.g., GIL contention) without benefit, so we do it synchronously.
    """
    device = getattr(weight, "device", None)
    if device is None:
        return False
    return device.type == "cpu"


def maybe_executor_submit(
    *,
    executor: concurrent.futures.ThreadPoolExecutor,
    futures: List[concurrent.futures.Future],
    use_async: bool,
    func: Callable[..., Any],
    func_args: Iterable[Any] = (),
    func_kwargs: Optional[Dict[str, Any]] = None,
) -> None:
    """Submit a task to the executor if async loading is enabled.

    Parameters (keyword-only):
    - executor: ThreadPoolExecutor used to submit background tasks
    - futures: a list collecting the submitted Future objects
    - use_async: whether to submit to executor or run inline
    - func: the callable to run
    - func_args: positional args for the callable (defaults to empty tuple)
    - func_kwargs: keyword args for the callable (defaults to empty dict)
    """
    if func_kwargs is None:
        func_kwargs = {}
    if use_async:
        futures.append(executor.submit(func, *func_args, **func_kwargs))
    else:
        func(*func_args, **func_kwargs)