# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0

import importlib
import os
import pkgutil
import traceback
from abc import ABC
from typing import Any, Type

import torch
from diffusers import AutoModel
from torch import nn
from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer

from sglang.multimodal_gen.configs.models import ModelConfig
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.layers.attention.selector import (
    component_attn_backend_context_manager,
    get_component_attn_backend_context,
)
from sglang.multimodal_gen.runtime.loader.utils import (
    _normalize_component_type,
    component_name_to_loader_cls,
    get_memory_usage_of_component,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    configure_layerwise_offload_modules,
    is_layerwise_offloaded_module,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
    LAYERWISE_OFFLOAD_ALL_COMPONENTS,
    LAYERWISE_OFFLOAD_DIT_GROUP,
    layerwise_component_matches_any_selection,
    normalize_layerwise_offload_components,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
    get_hf_config,
    prepare_diffusers_component_path_for_loading,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision

logger = init_logger(__name__)


def _load_auto_tokenizer_with_roberta_processing_compat(*args, **kwargs):
    from tokenizers import processors

    roberta_processing = processors.RobertaProcessing

    def roberta_processing_compat(*processor_args, **processor_kwargs):
        if "sep" in processor_kwargs and "cls" in processor_kwargs:
            sep = processor_kwargs.pop("sep")
            cls_token = processor_kwargs.pop("cls")
            return roberta_processing(
                sep, cls_token, *processor_args, **processor_kwargs
            )
        return roberta_processing(*processor_args, **processor_kwargs)

    processors.RobertaProcessing = roberta_processing_compat
    try:
        return AutoTokenizer.from_pretrained(*args, **kwargs)
    finally:
        processors.RobertaProcessing = roberta_processing


class ComponentLoader(ABC):
    """Base class for loading a specific type of model component."""

    # the list of possible name of the component in model_index.json, e.g., scheduler
    component_names: list[str] = []

    # diffusers or transformers
    expected_library: str = ""

    _loaders_registered = False

    def __init_subclass__(cls, **kwargs):
        """
        register loaders, called when subclass is imported
        """
        super().__init_subclass__(**kwargs)
        for component_name in cls.component_names:
            component_name_to_loader_cls[component_name] = cls

    def __init__(self, device=None) -> None:
        self.device = device
        self.component_architecture: str | None = None

    def should_offload(
        self, server_args: ServerArgs, model_config: ModelConfig | None = None
    ):
        # not offload by default
        return False

    def target_device(self, should_offload):
        if should_offload:
            return (
                torch.device("mps")
                if current_platform.is_mps()
                else torch.device("cpu")
            )
        else:
            return get_local_torch_device()

    def customized_load_kwargs_for_component(
        self, _server_args: ServerArgs, _component_name: str
    ) -> dict[str, Any]:
        return {}

    def should_raise_customized_load_error(
        self, _server_args: ServerArgs, _component_name: str
    ) -> bool:
        return False

    @staticmethod
    def _is_component_set_as_layerwise_load(
        server_args: ServerArgs, component_name: str
    ) -> bool:
        """if a component should be loaded in a layerwise-fashion"""
        selected_component_names = normalize_layerwise_offload_components(
            server_args.layerwise_offload_components
        )
        if selected_component_names is None:
            return False
        selected_component_names = set(selected_component_names)
        if LAYERWISE_OFFLOAD_ALL_COMPONENTS in selected_component_names:
            return True
        explicit_component_names = selected_component_names - {
            LAYERWISE_OFFLOAD_DIT_GROUP
        }
        return layerwise_component_matches_any_selection(
            component_name, explicit_component_names
        )

    def _maybe_configure_layerwise_after_startup_cpu_staging(
        self,
        component: AutoModel,
        server_args: ServerArgs,
        component_name: str,
        load_kwargs: dict[str, Any],
    ) -> AutoModel:
        if not load_kwargs.get("cpu_offload_flag"):
            return component
        if not isinstance(component, nn.Module):
            return component

        # try to configure layerwise-offload with the component
        configured_components = configure_layerwise_offload_modules(
            {component_name: component},
            server_args,
            component_names=server_args.layerwise_offload_components,
            warn_missing=False,
        )
        if is_layerwise_offloaded_module(component):
            logger.info(
                "Configured layerwise offload for %s immediately after startup CPU staging",
                component_name,
            )
            return component

        logger.warning(
            "Layerwise startup CPU staging was requested for %s, but the loaded "
            "module did not enable layerwise offload. Moving it to GPU.",
            component_name,
        )
        # ensures the module is on GPU
        if component_name in configured_components:
            return component
        return component.to(get_local_torch_device())

    def _load_customized_with_context(
        self,
        component_model_path: str,
        server_args: ServerArgs,
        component_name: str,
        attn_backend: Any,
        component_attn_name: str | None,
    ) -> AutoModel:
        with component_attn_backend_context_manager(
            attn_backend, component_name=component_attn_name
        ):
            load_kwargs = self.customized_load_kwargs_for_component(
                server_args, component_name
            )
            component = self.load_customized(
                component_model_path, server_args, component_name, **load_kwargs
            )
            return self._maybe_configure_layerwise_after_startup_cpu_staging(
                component, server_args, component_name, load_kwargs
            )

    def _load_native_with_context(
        self,
        component_model_path: str,
        server_args: ServerArgs,
        component_name: str,
        transformers_or_diffusers: str,
        attn_backend: Any,
        component_attn_name: str | None,
    ) -> AutoModel:
        with component_attn_backend_context_manager(
            attn_backend, component_name=component_attn_name
        ):
            component = self.load_native(
                component_model_path,
                server_args,
                transformers_or_diffusers,
                component_name,
            )
        should_offload = self.should_offload(server_args)
        target_device = self.target_device(should_offload)
        return component.to(device=target_device)

    def load(
        self,
        component_model_path: str,
        server_args: ServerArgs,
        component_name: str,
        transformers_or_diffusers: str,
    ) -> tuple[AutoModel, float]:
        """
        Template method that standardizes logging around the core load implementation.
        The priority of loading method is:
            1. load customized component
            2. load native diffusers/transformers component
        If all of the above methods failed, an error will be thrown

        """
        gpu_mem_before_loading = current_platform.get_available_gpu_memory()
        logger.info(
            "Loading %s from %s. avail mem: %.2f GB",
            component_name,
            component_model_path,
            gpu_mem_before_loading,
        )
        attn_backend = None
        component_attn_name = None
        if get_component_attn_backend_context() is None:
            attn_backend, matched_backend_key = (
                server_args.resolve_component_attention_backend(component_name)
            )
            component_attn_name = matched_backend_key or component_name
            if attn_backend is not None:
                logger.info(
                    "Using %s backend for component: %s",
                    attn_backend.name.lower(),
                    matched_backend_key,
                )
        try:
            component = self._load_customized_with_context(
                component_model_path,
                server_args,
                component_name,
                attn_backend,
                component_attn_name,
            )
            source = "sgl-diffusion"
        except Exception as e:
            if self.should_raise_customized_load_error(server_args, component_name):
                traceback.print_exc()
                raise RuntimeError(
                    f"Failed to load customized {component_name}; native fallback "
                    "is disabled for this component configuration."
                ) from e
            if "Unsupported model architecture" in str(e):
                logger.info(
                    f"Component: {component_name} doesn't have a customized version yet, using native version"
                )
            else:
                traceback.print_exc()
                logger.error(
                    f"Error while loading customized {component_name}, falling back to native version"
                )
            # fallback to native version
            component = self._load_native_with_context(
                component_model_path,
                server_args,
                component_name,
                transformers_or_diffusers,
                attn_backend,
                component_attn_name,
            )
            source = "native"
            logger.warning(
                "Native component %s: %s is loaded, performance may be sub-optimal",
                component_name,
                component.__class__.__name__,
            )

        if component is None:
            logger.error("Load %s failed", component_name)
            consumed = 0.0
        else:
            if isinstance(component, nn.Module):
                component = component.eval()
            current_gpu_mem = current_platform.get_available_gpu_memory()
            model_size = get_memory_usage_of_component(component) or "NA"
            consumed = gpu_mem_before_loading - current_gpu_mem
            logger.info(
                f"Loaded %s: %s ({source} version). model size: %s GB, consumed GPU mem: %.2f GB, avail GPU mem: %.2f GB",
                component_name,
                component.__class__.__name__,
                model_size,
                consumed,
                current_gpu_mem,
            )
        return component, consumed

    def load_native(
        self,
        component_model_path: str,
        server_args: ServerArgs,
        transformers_or_diffusers: str,
        component_name: str | None = None,
    ) -> AutoModel:
        """
        Load the component using the native library (transformers/diffusers).
        """
        precision = (
            resolve_component_precision(server_args, component_name)
            if component_name is not None
            else None
        )
        load_kwargs = {}
        if precision is not None:
            load_kwargs["torch_dtype"] = precision

        if transformers_or_diffusers == "transformers":
            from transformers import AutoModel

            config = get_hf_config(
                component_model_path,
                trust_remote_code=server_args.trust_remote_code,
                revision=server_args.revision,
            )
            return AutoModel.from_pretrained(
                component_model_path,
                config=config,
                trust_remote_code=server_args.trust_remote_code,
                revision=server_args.revision,
                **load_kwargs,
            )
        elif transformers_or_diffusers == "diffusers":
            from diffusers import AutoModel

            component_model_path = prepare_diffusers_component_path_for_loading(
                component_model_path
            )
            return AutoModel.from_pretrained(
                component_model_path,
                revision=server_args.revision,
                trust_remote_code=server_args.trust_remote_code,
                **load_kwargs,
            )
        else:
            raise ValueError(f"Unsupported library: {transformers_or_diffusers}")

    def load_customized(
        self, component_model_path: str, server_args: ServerArgs, component_name: str
    ):
        """
        Load the customized version component, implemented and optimized in SGL-diffusion
        """
        raise NotImplementedError(
            f"load_customized not implemented for {self.__class__.__name__}"
        )

    @classmethod
    def _ensure_loaders_registered(cls):
        """
        avoid multiple registration
        """
        if cls._loaders_registered:
            return

        package_dir = os.path.dirname(__file__)
        package_name = (
            __package__
            or "sglang.multimodal_gen.runtime.loader.component_loaders.component_loaders"
        )

        for _, name, _ in pkgutil.iter_modules([package_dir]):
            # skip importing self to avoid circular dependency issues
            if name == "component_loader":
                continue
            try:
                importlib.import_module(f".{name}", package=package_name)
            except ImportError as e:
                logger.warning(f"Failed to import loader component {name}: {e}")

        cls._loaders_registered = True

    @classmethod
    def resolve_transformers_or_diffusers(
        self, transformers_or_diffusers: str, component_name: str
    ) -> str:
        # NOTE(FlamingoPg): special for LTX-2 models
        if component_name == "vocoder" or component_name == "connectors":
            transformers_or_diffusers = "diffusers"

        # NOTE(CloudRipple): special for MOVA models
        # TODO(CloudRipple): remove most of these special cases after unifying the loading logic
        if component_name in [
            "audio_vae",
            "audio_dit",
            "dual_tower_bridge",
            "video_dit",
        ]:
            transformers_or_diffusers = "diffusers"

        if (
            component_name == "scheduler"
            and transformers_or_diffusers == "mova.diffusion.schedulers.flow_match_pair"
        ):
            transformers_or_diffusers = "diffusers"

        return transformers_or_diffusers

    @classmethod
    def for_component_type(
        cls,
        component_name: str,
        transformers_or_diffusers: str,
        component_architecture: str | None = None,
    ) -> "ComponentLoader":
        """
        Factory method to create a component loader for a specific component type.

        Args:
            component_name: Type of component (e.g., "vae", "text_encoder", "transformer", "scheduler")
            transformers_or_diffusers: Whether the component is from transformers or diffusers
        """
        cls._ensure_loaders_registered()

        # Map of component types to their loader classes and expected library
        component_name = _normalize_component_type(component_name)

        transformers_or_diffusers = cls.resolve_transformers_or_diffusers(
            transformers_or_diffusers, component_name
        )

        if component_name in component_name_to_loader_cls:
            loader_cls: Type[ComponentLoader] = component_name_to_loader_cls[
                component_name
            ]
            expected_library = loader_cls.expected_library
            # Assert that the library matches what's expected for this component type
            assert (
                transformers_or_diffusers == expected_library
            ), f"{component_name} must be loaded from {expected_library}, got {transformers_or_diffusers}"
            loader = loader_cls()
            loader.component_architecture = component_architecture
            return loader

        # For unknown component types, use a generic loader
        logger.warning(
            "No specific loader found for component type: %s. Using generic loader.",
            component_name,
        )
        return GenericComponentLoader(transformers_or_diffusers, component_architecture)


class ImageProcessorLoader(ComponentLoader):
    """Loader for image processor."""

    component_names = ["image_processor"]
    expected_library = "transformers"

    def load_customized(
        self, component_model_path: str, server_args: ServerArgs, component_name: str
    ) -> Any:
        return AutoImageProcessor.from_pretrained(component_model_path, use_fast=True)


class AutoProcessorLoader(ComponentLoader):
    """Loader for auto processor."""

    component_names = ["processor"]
    expected_library = "transformers"

    def load_customized(
        self, component_model_path: str, server_args: ServerArgs, component_name: str
    ) -> Any:
        return AutoProcessor.from_pretrained(component_model_path)


class TokenizerLoader(ComponentLoader):
    """Loader for tokenizers."""

    component_names = ["tokenizer", "text_tokenizer"]
    expected_library = "transformers"

    def load_customized(
        self, component_model_path: str, server_args: ServerArgs, component_name: str
    ) -> Any:
        # Some pipelines keep the slot name `tokenizer` in model_index.json even
        # when the declared class is a processor. e.g. FLUX.2:
        # `tokenizer: ["transformers", "PixtralProcessor"]`.
        # Honor the declared component class instead of guessing from the slot name.
        if (
            self.component_architecture is not None
            and self.component_architecture.endswith("Processor")
        ):
            return AutoProcessor.from_pretrained(component_model_path)

        # Qwen-Image's model_index declares Qwen2Tokenizer; using the fast class
        # changes text preprocessing and shifts official GT comparisons.
        use_fast = self.component_architecture != "Qwen2Tokenizer"
        try:
            return AutoTokenizer.from_pretrained(
                component_model_path,
                padding_side="right",
                use_fast=use_fast,
            )
        except TypeError as e:
            # tokenizers>=0.21 removed the `cls` kwarg from RobertaProcessing,
            # but some transformers CLIPTokenizer builds still pass it. Fall back
            # to the pure-Python (slow) tokenizer which avoids the rust path.
            if "RobertaProcessing" in str(e) and use_fast:
                logger.warning(
                    "Fast tokenizer failed (%s), retrying with use_fast=False", e
                )
                return _load_auto_tokenizer_with_roberta_processing_compat(
                    component_model_path,
                    padding_side="right",
                    use_fast=False,
                )
            raise


class GenericComponentLoader(ComponentLoader):
    """Generic loader for components that don't have a specific loader."""

    def __init__(
        self, library="transformers", component_architecture: str | None = None
    ) -> None:
        super().__init__()
        self.library = library
        self.component_architecture = component_architecture


class PipelineComponentLoader:
    """
    Utility class for loading the components in a pipeline.
    """

    @staticmethod
    def load_component(
        component_name: str,
        component_model_path: str,
        transformers_or_diffusers: str,
        server_args: ServerArgs,
        component_architecture: str | None = None,
    ):
        """
        Load a pipeline component.

        Args:
            component_name: Name of the component (e.g., "vae", "text_encoder", "transformer", "scheduler")
            component_model_path: Path to the component model
            transformers_or_diffusers: Whether the component is from transformers or diffusers
            component_architecture: the class name of the module
        """

        # Get the appropriate loader for this component type
        loader = ComponentLoader.for_component_type(
            component_name, transformers_or_diffusers, component_architecture
        )

        try:
            # Load the component
            return loader.load(
                component_model_path,
                server_args,
                component_name,
                transformers_or_diffusers,
            )
        except Exception as e:
            logger.error(
                f"Error while loading component: {component_name}, {component_model_path=}"
            )
            raise e
