"""Adapt four allowlisted ZenMux JSON ASR models to LiteLLM's audio endpoint."""
import base64

from litellm.llms.base_llm.audio_transcription.transformation import AudioTranscriptionRequestData
from litellm.llms.openai.transcriptions.whisper_transformation import OpenAIWhisperAudioTranscriptionConfig
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.types.utils import TranscriptionResponse, TranscriptionUsageDurationObject

MODELS = frozenset({'openai/gpt-transcribe', 'qwen/qwen-audio-3.1-asr-flash',
    'bytedance/doubao-seed-asr-2.0', 'xiaomi/mimo-v2.5-asr'})
API_BASE = 'https://zenmux.ai/api/v1'


class ZenMuxTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig):
    @property
    def has_native_transcription_endpoint(self):
        return True

    def get_supported_openai_params(self, model):
        return ['language', 'response_format', 'prompt']

    def map_openai_params(self, non_default_params, optional_params, model, drop_params):
        params = super().map_openai_params(non_default_params, optional_params, model, drop_params)
        if params.get('prompt'):
            raise ValueError('ZenMux comparison routes do not accept prompts; leave prompt empty')
        if params.get('response_format') not in (None, 'json', 'text'):
            raise ValueError('ZenMux comparison routes support json/text only')
        return params

    def validate_environment(self, headers, model, messages, optional_params,
                             litellm_params, api_key=None, api_base=None):
        if not api_key:
            raise ValueError('ZenMux subscription API key is required')
        return {**headers, 'Authorization': 'Bearer ' + api_key, 'Content-Type': 'application/json'}

    def get_complete_url(self, api_base, api_key, model, optional_params, litellm_params, stream=None):
        if (api_base or '').rstrip('/') != API_BASE or model not in MODELS:
            raise ValueError('ZenMux ASR adapter is restricted to the configured endpoint and model allowlist')
        if stream:
            raise ValueError('Streaming is not enabled for comparison dictation routes')
        return API_BASE + '/audio/transcriptions'

    def transform_audio_transcription_request(self, model, audio_file, optional_params, litellm_params):
        if model not in MODELS:
            raise ValueError('Unknown ZenMux ASR model')
        if optional_params.get('timestamp_granularities') or optional_params.get('prompt'):
            raise ValueError('Timestamps and prompt injection are disabled on comparison routes')
        audio = process_audio_file(audio_file)
        formats = {'audio/wav': 'wav', 'audio/x-wav': 'wav', 'audio/wave': 'wav',
            'audio/mpeg': 'mp3', 'audio/mp3': 'mp3', 'audio/flac': 'flac',
            'audio/x-flac': 'flac', 'audio/m4a': 'm4a', 'audio/mp4': 'm4a',
            'audio/x-m4a': 'm4a', 'audio/ogg': 'ogg', 'audio/webm': 'webm',
            'video/webm': 'webm', 'audio/aac': 'aac'}
        mime = audio.content_type.split(';', 1)[0].lower()
        if mime not in formats:
            raise ValueError('Unsupported audio MIME type for ZenMux ASR')
        body = {'model': model, 'input_audio': {
            'data': base64.b64encode(audio.file_content).decode(), 'format': formats[mime]},
            'stream': False}
        language = optional_params.get('language')
        if language:
            body['language'] = language
        return AudioTranscriptionRequestData(data=body)

    def transform_audio_transcription_response(self, raw_response):
        body = raw_response.json()
        if not isinstance(body.get('text'), str):
            raise ValueError('ZenMux ASR returned no transcript text')
        response = TranscriptionResponse(text=body['text'])
        usage = body.get('usage') or {}
        if usage.get('seconds') is not None:
            # ZenMux has no audio/text token breakdown; report duration without inventing one.
            response.usage = TranscriptionUsageDurationObject(type='duration', seconds=usage['seconds'])
            response['duration'] = usage['seconds']
        return response


def install():
    import litellm
    from litellm.utils import ProviderConfigManager
    original = ProviderConfigManager.get_provider_audio_transcription_config
    if getattr(original, '_zenmux_asr_adapter', False):
        return

    def select(model, provider):
        if provider == litellm.LlmProviders.OPENAI and model in MODELS:
            return ZenMuxTranscriptionConfig()
        return original(model, provider)

    select._zenmux_asr_adapter = True
    ProviderConfigManager.get_provider_audio_transcription_config = staticmethod(select)