"""Minimal configuration shipped inside MOSS-Music HuggingFace weight folders.

This file is copied into each released checkpoint by ``convert_hf_checkpoint.py``
so that ``AutoConfig.from_pretrained(..., trust_remote_code=True)`` works out of
the box. The full developer-facing configuration lives in
``src/configuration_moss_music.py``.
"""
from transformers import PretrainedConfig, Qwen3Config


class MossMusicConfig(PretrainedConfig):
    model_type = "moss_music"
    is_composition = True

    def __init__(
        self,
        audio_config=None,
        language_config=None,
        adapter_hidden_size=8192,
        ignore_index=-100,
        deepstack_num_inject_layers=None,
        **kwargs,
    ):
        if isinstance(language_config, dict):
            language_config = Qwen3Config(**language_config)
        elif language_config is None:
            language_config = Qwen3Config()

        self.audio_config = audio_config
        self.language_config = language_config
        self.adapter_hidden_size = adapter_hidden_size
        self.ignore_index = ignore_index
        self.deepstack_num_inject_layers = deepstack_num_inject_layers

        for key in ("num_hidden_layers", "eos_token_id", "bos_token_id", "vocab_size"):
            kwargs.setdefault(key, getattr(language_config, key, None))

        super().__init__(**kwargs)


def _register_moss_music_with_sglang() -> None:
    """Register MossMusic* as aliases of the built-in MossAudio* classes inside
    SGLang. This module is imported by ``AutoConfig.from_pretrained`` via
    ``auto_map`` + ``trust_remote_code``, which fires before SGLang looks up
    ``architectures[0]`` in its ``ModelRegistry`` / ``PROCESSOR_MAPPING``.
    """
    try:
        from sglang.srt.configs import model_config as _sg_model_config
        from sglang.srt.managers.multimodal_processor import PROCESSOR_MAPPING
        from sglang.srt.models.moss_audio import MossAudioModel
        from sglang.srt.models.registry import ModelRegistry
        from sglang.srt.multimodal.processors.moss_audio import (
            MossAudioMultimodalProcessor,
        )
    except Exception:
        return

    if "MossMusicModel" not in ModelRegistry.models:
        alias_cls = type("MossMusicModel", (MossAudioModel,), {})
        ModelRegistry.models["MossMusicModel"] = alias_cls
        PROCESSOR_MAPPING[alias_cls] = MossAudioMultimodalProcessor

    # SGLang also dispatches on the *string* architecture name in several
    # helper functions (multimodal detection, always-process-mm-data, ...).
    # The Python-class alias above does not cover those, so we also extend
    # the name-based whitelists that mirror ``MossAudioModel``.
    try:
        if "MossMusicModel" not in _sg_model_config.multimodal_model_archs:
            _sg_model_config.multimodal_model_archs.append("MossMusicModel")
    except Exception:
        pass

    try:
        if not getattr(
            _sg_model_config.is_always_process_mm_data_model,
            "_moss_music_patched",
            False,
        ):
            _orig_always_mm = _sg_model_config.is_always_process_mm_data_model

            def _patched_always_mm(model_architectures):
                if "MossMusicModel" in model_architectures:
                    return True
                return _orig_always_mm(model_architectures)

            _patched_always_mm._moss_music_patched = True  # type: ignore[attr-defined]
            _sg_model_config.is_always_process_mm_data_model = _patched_always_mm
    except Exception:
        pass


_register_moss_music_with_sglang()