# SPDX-License-Identifier: Apache-2.0
"""Base contract for per-model TTS serving adapters.

This package factors the per-model ``if self._tts_model_type == ...`` dispatch
in ``serving_speech.py`` into one adapter class per model. Each adapter owns its
model's request normalization, validation, prompt/param building, sampling
overrides, and output policy, so adding a model means writing one adapter file
instead of editing the shared serving module in ~10 scattered places.

See the RFC for the full design (issue #4327). This is the foundation landed in
the first migration PR; Qwen3-TTS is the first model routed through it while the
remaining models stay on the legacy path until individually migrated.
"""

from abc import ABC, abstractmethod
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar

if TYPE_CHECKING:
    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest


_conditioning_cache_salt_fn: "Callable[..., str] | None" = None


def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: dict) -> str:
    """Return the conditioning cache salt for ``request`` + ``tts_params``.

    Lazily imports and caches ``serving_speech._conditioning_cache_salt`` on first
    use: the import is deferred to break the adapters<->serving_speech import
    cycle, and cached so it resolves once instead of on every ``build()`` call.
    """
    global _conditioning_cache_salt_fn
    if _conditioning_cache_salt_fn is None:
        from vllm_omni.entrypoints.openai.serving_speech import _conditioning_cache_salt

        _conditioning_cache_salt_fn = _conditioning_cache_salt
    return _conditioning_cache_salt_fn(request, tts_params)


@dataclass
class OutputPolicy:
    """How the orchestrator aggregates engine output for a model.

    ``accumulate_nonstreaming`` enables MOSS-style cross-step accumulation in the
    non-streaming path. Streaming cumulative/delta semantics stay engine-side,
    keyed by request id, and are unaffected by this flag.
    """

    accumulate_nonstreaming: bool = False


@dataclass
class PreparedRequest:
    """Everything the generic orchestrator needs to call ``<engine>.generate()``.

    The fields mirror what ``_prepare_speech_generation`` assembled inline:
    ``prompt`` is the engine prompt dict, ``tts_params`` the per-model parameter
    dict, ``model_type`` the discriminator used for logging, and
    ``output_policy`` controls non-streaming aggregation.
    """

    prompt: dict[str, Any]
    tts_params: dict[str, Any] = field(default_factory=dict)
    model_type: str = "generic"
    output_policy: OutputPolicy = field(default_factory=OutputPolicy)
    #: Cross-cutting per-request state the orchestrator still owns (e.g. the
    #: Qwen3-TTS ref-audio warmup artifact key tracked after ``generate()``).
    warmup_artifact_key: str | None = None


@dataclass
class SpeechServingContext:
    """Shared state and helpers an adapter may use.

    During the incremental migration this holds a back-reference to the owning
    ``OmniOpenAIServingSpeech`` instance (``server``) so adapters can reuse the
    existing, battle-tested helper implementations (ref-audio resolution,
    uploaded-speaker handling, prompt-length estimation, speaker storage)
    without relocating them. Follow-up PRs may narrow this to explicit fields as
    more models migrate. Exactly one of ``engine_client`` / ``diffusion_engine``
    is set, matching the deployment's serving backend.
    """

    server: Any
    engine_client: Any | None = None
    diffusion_engine: Any | None = None


class TTSModelAdapter(ABC):
    """Mandatory base class for a TTS model served via ``/v1/audio/speech``.

    One concrete subclass per model, registered by stage key. The serving layer
    resolves exactly one adapter at startup and dispatches all per-model
    decisions to it. Adapters reuse shared helpers through ``ctx`` rather than
    re-implementing them.
    """

    #: Stable discriminator string (the model-type from detection); registry key.
    name: ClassVar[str]
    #: Engine ``model_stage`` key(s) this model uses. Load-bearing: the serving
    #: layer discovers the TTS stage and resolves the model type from these.
    stage_keys: ClassVar[frozenset[str]] = frozenset()
    #: Engine ``model_arch`` value(s) that identify this model type. Checked
    #: before ``stage_keys`` in :meth:`matches`. Only needed by models whose
    #: ``model_stage`` alone is ambiguous or absent.
    model_archs: ClassVar[frozenset[str]] = frozenset()
    #: Set when the model has no dedicated ``model_stage`` value and its AR
    #: entry stage must be discovered by ``model_archs`` instead (Ming dense).
    arch_identifies_entry_stage: ClassVar[bool] = False
    #: Detection order; lower runs first. Only set this to break a genuine
    #: overlap with another adapter — see ``tests/.../test_tts_detection.py``,
    #: which fails if two same-priority detectors can match the same input.
    detect_priority: ClassVar[int] = 100
    #: Serving backend: ``"ar"`` (engine_client) or ``"diffusion"``.
    backend: ClassVar[str] = "ar"
    #: Whether the model consumes ``request.speed`` in its native parameters.
    native_speed_control: ClassVar[bool] = False

    max_new_tokens_min = 1

    max_new_tokens_max = 4096

    def __init__(self, ctx: SpeechServingContext) -> None:
        self.ctx = ctx

    @classmethod
    def matches(cls, model_stage: str | None, model_arch: str | None) -> bool:
        """Whether a deployed stage with this ``model_stage``/``model_arch``
        is served by this adapter.

        Architecture wins over stage key, so a model that declares both is
        recognized even when deployed under a stage key it shares with another
        model. Override only for match rules that are not a set membership
        test (see ``covo_audio``).
        """
        if model_arch is not None and model_arch in cls.model_archs:
            return True
        return model_stage is not None and model_stage in cls.stage_keys

    @classmethod
    def stage_serves_speech(cls, model_stage: str | None, all_stage_keys: frozenset[str]) -> bool:
        """Whether a matched stage really accepts ``/v1/audio/speech`` in *this*
        deployment.

        Lets a model that is speech-capable only in some topologies say so
        (Audex: its omni thinker is text-final unless the speech decoder is
        deployed alongside). Default: always.
        """
        return True

    def normalize(self, request: "OpenAICreateSpeechRequest") -> None:
        """In-place request normalization/mutation (e.g. infer task type,
        lowercase voice). Default: no-op."""

    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
        """Return an error string, or ``None`` if the request is valid.

        Should be free of new side effects beyond what ``normalize`` did.
        """
        return None

    @abstractmethod
    async def build(
        self,
        request: "OpenAICreateSpeechRequest",
        sampling_params_list: list,
        has_inline_ref_audio: bool,
    ) -> PreparedRequest:
        """Build the engine prompt + tts_params for this request.

        ``sampling_params_list`` is passed read-only for models (e.g. MOSS) that
        fold the resolved seed into ``additional_information`` at build time.

        ``has_inline_ref_audio`` is captured by the orchestrator *before*
        ``validate()`` runs, because ``_apply_uploaded_speaker`` (invoked inside
        several adapters' ``validate``) sets ``request.ref_audio`` in place.
        Recomputing it here would misclassify uploaded voices as inline and drop
        the ``voice_name`` / ``voice_created_at`` metadata.
        """

    def apply_sampling_overrides(
        self,
        sampling_params_list: list,
        request: "OpenAICreateSpeechRequest",
    ) -> list:
        """Apply model-specific sampling mutations.

        The orchestrator guarantees the call order:
        stream-coercion -> extra_params -> THIS -> seed. Default: identity.
        """
        return sampling_params_list


class ARTTSAdapter(TTSModelAdapter):
    """Adapter for models served through the AR ``engine_client`` path.

    Covers pure-AR codec models as well as AR-base-LM + diffusion-side-computation
    hybrids (e.g. VoxCPM2, Ming) whose internal diffusion is invisible to the
    serving layer.
    """

    backend: ClassVar[str] = "ar"


class DiffusionTTSAdapter(TTSModelAdapter):
    """Adapter for pure-diffusion pipelines served through the diffusion engine.

    In scope today: OmniVoice. Bridges to the diffusion pipeline parameter
    contract (#3572) when present; see ``extra_body_params``.
    """

    backend: ClassVar[str] = "diffusion"

    #: Backing diffusion pipeline class (for EXTRA_BODY_PARAMS lookup).
    pipeline_cls: ClassVar[type | None] = None

    @classmethod
    def extra_body_params(cls) -> frozenset[str]:
        """Fallback-safe access to the pipeline's declared body params.

        Returns the pipeline's ``EXTRA_BODY_PARAMS`` if the #3572 contract is
        present, else an empty frozenset (the adapter then uses its own inline
        parameter logic).
        """
        params = getattr(cls.pipeline_cls, "EXTRA_BODY_PARAMS", None)
        return frozenset(params) if params is not None else frozenset()


@dataclass(frozen=True)
class LegacyDetector:
    """A model type the serving layer detects but has no adapter for yet.

    Detection has to keep naming these models so ``serving_speech.py`` can route
    them down its legacy path, but they own none of the adapter contract. Each
    entry is a migration debt: delete it when the model gets an adapter. The
    pre-commit gate (``tools/pre_commit/check_tts_adapter.py``) fails if the list
    grows.

    Exposes the subset of the :class:`TTSModelAdapter` class surface that
    detection reads, so both kinds live in one sorted sequence.
    """

    name: str
    stage_keys: frozenset[str] = frozenset()
    model_archs: frozenset[str] = frozenset()
    arch_identifies_entry_stage: bool = False
    detect_priority: int = 100

    def matches(self, model_stage: str | None, model_arch: str | None) -> bool:
        if model_arch is not None and model_arch in self.model_archs:
            return True
        return model_stage is not None and model_stage in self.stage_keys


# Re-exported here to avoid import cycles at call sites.
__all__ = [
    "ARTTSAdapter",
    "DiffusionTTSAdapter",
    "LegacyDetector",
    "OutputPolicy",
    "PreparedRequest",
    "SpeechServingContext",
    "TTSModelAdapter",
]