"""Bridge OpenAI file uploads to Beijing Bailian's native ASR API."""
import base64
import io
import json
from pathlib import Path
from urllib.parse import urlsplit
import wave

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

MODEL = 'qwen-audio-3.1-asr-flash'
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/mp4':'m4a','audio/m4a':'m4a','audio/x-m4a':'m4a',
    'audio/ogg':'ogg','audio/webm':'webm','video/webm':'webm','audio/aac':'aac'}
VOCABULARY_PATHS = frozenset({'/etc/litellm/personal-vocabulary.json',
                              '/credentials/personal-vocabulary.json'})


def vocabulary(path):
    if path not in VOCABULARY_PATHS:
        raise ValueError('Unexpected personal vocabulary path')
    words = json.loads(Path(path).read_text())
    if not isinstance(words,list) or not 1 <= len(words) <= 100:
        raise ValueError('Personal vocabulary must contain 1-100 terms')
    if any(not isinstance(w,str) or not w.strip() or len(w)>100 for w in words):
        raise ValueError('Invalid personal vocabulary term')
    return {w.strip():3 for w in words}


class BailianTranscriptionConfig(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):
        result=super().map_openai_params(non_default_params,optional_params,model,drop_params)
        if result.get('response_format') not in (None,'json','text'):
            raise ValueError('Bailian dictation supports json/text without timestamps')
        return result

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

    def get_complete_url(self,api_base,api_key,model,optional_params,litellm_params,stream=None):
        parsed=urlsplit(api_base or '')
        if (model != MODEL or parsed.scheme != 'https' or
            not (parsed.hostname or '').endswith('.cn-beijing.maas.aliyuncs.com') or
            parsed.path.rstrip('/') != '/api/v1' or parsed.port not in (None,443) or
            parsed.username or parsed.password or parsed.query or parsed.fragment):
            raise ValueError('Bailian ASR restricted to Beijing workspace native endpoints')
        if stream:raise ValueError('Streaming not enabled on this dictation route')
        return api_base.rstrip('/')+'/services/aigc/multimodal-generation/generation'

    def transform_audio_transcription_request(self,model,audio_file,optional_params,litellm_params):
        if model != MODEL or optional_params.get('timestamp_granularities'):
            raise ValueError('Unsupported Bailian dictation model or timestamps')
        audio=process_audio_file(audio_file)
        mime=audio.content_type.split(';',1)[0].lower()
        if mime not in FORMATS:raise ValueError('Unsupported Bailian audio MIME type')
        if len(audio.file_content)>10_000_000:raise ValueError('Bailian inline audio exceeds 10 MB')
        parameters={'format':FORMATS[mime]}
        if parameters['format']=='wav':
            with wave.open(io.BytesIO(audio.file_content),'rb') as wav:
                parameters['sample_rate']=str(wav.getframerate())
        language=optional_params.get('language')
        if language:parameters['language_hints']=[language]
        # LiteLLM classifies DashScope as OpenAI-compatible and nests custom kwargs.
        extra=optional_params.get('extra_body') or {}
        path=optional_params.get('bailian_vocabulary_path') or extra.get('bailian_vocabulary_path')
        if path:parameters['vocabulary']=vocabulary(path)
        prompt=optional_params.get('prompt')
        messages=[]
        if prompt:
            if not isinstance(prompt,str) or len(prompt)>400:
                raise ValueError('Bailian context prompt must be at most 400 characters')
            messages.append({'role':'user','content':[{'type':'input_text','text':prompt}]})
        # Always preserve the original bytes; only WAV metadata supplies sample rate.
        uri_mime='audio/mp3' if parameters['format']=='mp3' else mime
        data_uri='data:'+uri_mime+';base64,'+base64.b64encode(audio.file_content).decode()
        if len(data_uri)>10_000_000:
            raise ValueError('Bailian inline audio exceeds 10 MB after Base64 encoding')
        messages.append({'role':'user','content':[{'type':'input_audio','input_audio':{
            'data':data_uri}}]})
        return AudioTranscriptionRequestData(data={'model':MODEL,'input':{'messages':messages},
                                                  'parameters':parameters})

    def transform_audio_transcription_response(self,raw_response):
        body=raw_response.json()
        text=body.get('output',{}).get('text',body.get('text'))
        if not isinstance(text,str):raise ValueError('Bailian returned no transcript')
        result=TranscriptionResponse(text=text)
        duration=(body.get('usage') or {}).get('duration')
        if duration is not None:
            result.usage=TranscriptionUsageDurationObject(type='duration',seconds=duration)
            result['duration']=duration
        return result


def install():
    import litellm
    from litellm.utils import ProviderConfigManager
    original=ProviderConfigManager.get_provider_audio_transcription_config
    if getattr(original,'_bailian_asr_adapter',False):return
    def select(model,provider):
        if provider==litellm.LlmProviders.DASHSCOPE and model==MODEL:
            return BailianTranscriptionConfig()
        return original(model,provider)
    select._bailian_asr_adapter=True
    ProviderConfigManager.get_provider_audio_transcription_config=staticmethod(select)