# SPDX-License-Identifier: Apache-2.0
"""Registry of TTS serving adapters, and the model detection built on it.

Adapters register themselves by their ``name`` (the model-type discriminator)
via ``@register_tts_adapter``, and declare the deployment metadata that
identifies them: ``stage_keys``, plus ``model_archs`` for the few models a stage
key alone cannot distinguish.

That metadata is the single source of truth. :func:`detect_tts_model_type` and
:func:`all_tts_stage_keys` derive the whole stage -> model-type mapping from it,
so a new TTS model is one adapter file and needs no edit to ``serving_speech.py``.
"""

from vllm.logger import init_logger

from vllm_omni.entrypoints.openai.tts_adapters.base import (
    ARTTSAdapter,
    DiffusionTTSAdapter,
    OutputPolicy,
    PreparedRequest,
    SpeechServingContext,
    TTSModelAdapter,
)

logger = init_logger(__name__)

TTS_ADAPTER_REGISTRY: dict[str, type[TTSModelAdapter]] = {}


def register_tts_adapter(cls: type[TTSModelAdapter]) -> type[TTSModelAdapter]:
    """Class decorator: index ``cls`` under its ``name`` (model-type)."""
    if cls.name in TTS_ADAPTER_REGISTRY:
        raise ValueError(
            f"TTS adapter name {cls.name!r} already registered to "
            f"{TTS_ADAPTER_REGISTRY[cls.name].__qualname__}; {cls.__qualname__} conflicts."
        )
    TTS_ADAPTER_REGISTRY[cls.name] = cls
    return cls


def all_tts_model_types() -> frozenset[str]:
    """All registered model-type names."""
    return frozenset(TTS_ADAPTER_REGISTRY)


def resolve_adapter(model_type: str | None) -> type[TTSModelAdapter] | None:
    """Return the adapter for a detected model-type, or ``None``."""
    if model_type is None:
        return None
    return TTS_ADAPTER_REGISTRY.get(model_type)


def iter_tts_detectors() -> list[type[TTSModelAdapter]]:
    """Every detector, in resolution order.

    Sorted by ``detect_priority`` then ``name``, so the order is total and does
    not depend on adapter import order.
    """
    detectors = list(TTS_ADAPTER_REGISTRY.values())
    detectors.sort(key=lambda d: (d.detect_priority, d.name))
    return detectors


def detect_tts_model_type(model_stage: str | None, model_arch: str | None) -> str | None:
    """Resolve a deployed stage to its TTS model-type, or ``None``.

    The first detector that matches wins; ties are impossible for distinct model
    types because overlapping detectors must declare an explicit
    ``detect_priority`` (enforced by ``test_tts_detection.py``).
    """
    for detector in iter_tts_detectors():
        if detector.matches(model_stage, model_arch):
            return detector.name
    return None


def all_tts_stage_keys() -> frozenset[str]:
    """Every ``model_stage`` value served by ``/v1/audio/speech``."""
    keys: set[str] = set()
    for detector in iter_tts_detectors():
        keys |= detector.stage_keys
    return frozenset(keys)


def tts_entry_stage_archs() -> frozenset[str]:
    """``model_arch`` values that identify an AR entry stage on their own.

    Only for models that own no ``model_stage`` key; stage discovery falls back
    to these. Deliberately *not* every declared ``model_arch``: a model that also
    owns a stage key is found by that key, and widening this would change which
    stage is selected in mixed deployments.
    """
    archs: set[str] = set()
    for detector in iter_tts_detectors():
        if detector.arch_identifies_entry_stage:
            archs |= detector.model_archs
    return frozenset(archs)


# Import adapter modules for their registration side effects. Keep at the bottom
# so the registry helpers above are defined first.
from vllm_omni.entrypoints.openai.tts_adapters import (  # noqa: E402,F401
    audex,
    audex_tta,
    cosyvoice3,
    covo_audio,
    fish_speech,
    glm_tts,
    higgs_audio_v2,
    higgs_audio_v3,
    indextts2,
    ming_flash_omni_tts,
    ming_tts,
    minimax_music3,
    moss_tts,
    omnivoice,
    qwen3_tts,
    step_audio2,
    voxcpm2,
    voxtral,
)

__all__ = [
    "ARTTSAdapter",
    "DiffusionTTSAdapter",
    "OutputPolicy",
    "PreparedRequest",
    "SpeechServingContext",
    "TTSModelAdapter",
    "TTS_ADAPTER_REGISTRY",
    "all_tts_model_types",
    "all_tts_stage_keys",
    "detect_tts_model_type",
    "iter_tts_detectors",
    "register_tts_adapter",
    "resolve_adapter",
    "tts_entry_stage_archs",
]