import os

from .base import MusicGenerationProvider


def get_provider() -> MusicGenerationProvider:
    backend = os.environ.get("GENERATION_BACKEND", "mock")

    if backend == "mock":
        from .mock_provider import MockMusicGenerationProvider

        return MockMusicGenerationProvider()

    if backend == "diffrhythm":
        from .diffrhythm_provider import DiffRhythmProvider

        return DiffRhythmProvider()

    raise ValueError(f"Unknown GENERATION_BACKEND: {backend}")