import contextlib
import typing

import httpx
from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ..core.http_response import AsyncHttpResponse, HttpResponse
from ..core.request_options import RequestOptions
from ..core.serialization import convert_and_respect_annotation_metadata
from ..types.stream_tts_inference_options import StreamTtsInferenceOptions
from ..types.stream_tts_output_configuration import StreamTtsOutputConfiguration
from ..types.stream_tts_voice_settings import StreamTtsVoiceSettings
from .types.create_stream_tts_request_payload_language import CreateStreamTtsRequestPayloadLanguage
from .types.create_stream_tts_request_payload_speech_model import CreateStreamTtsRequestPayloadSpeechModel

OMIT = typing.cast(typing.Any, ...)


@contextlib.contextmanager
def baseten_tts(
    client_wrapper: SyncClientWrapper,
    *,
    text: str,
    language: CreateStreamTtsRequestPayloadLanguage,
    voice_id: typing.Optional[int] = OMIT,
    speech_model: typing.Optional[CreateStreamTtsRequestPayloadSpeechModel] = OMIT,
    user_instructions: typing.Optional[str] = OMIT,
    enhance_named_entities_pronunciation: typing.Optional[bool] = OMIT,
    output_configuration: typing.Optional[StreamTtsOutputConfiguration] = OMIT,
    voice_settings: typing.Optional[StreamTtsVoiceSettings] = OMIT,
    inference_options: typing.Optional[StreamTtsInferenceOptions] = OMIT,
    request_options: typing.Optional[RequestOptions] = None,
) -> typing.Iterator[httpx.Response]:
    # Retrieve API key from provider_params
    provider_params = client_wrapper.provider_params or {}
    api_key = provider_params.get("api_key", "")
    mars_pro_url = provider_params.get("mars_pro_url") or provider_params.get("mars_url")
    if not mars_pro_url:
        raise ValueError("mars_url is required for using Baseten as provider")

    headers = {
        "Authorization": f"Api-Key {api_key}" if api_key else "",
        "Content-Type": "application/json",
    }
    # Construct Payload
    # 1. Basic Fields
    payload = {
        "text": text,
        "language": str(language).lower().replace("_", "-"),
        "output_format": "wav",
        "stream": True,
        "apply_ner_nlp": False,
    }

    # 2. Output Configuration
    if output_configuration and output_configuration is not OMIT:
        if output_configuration.format:
            payload["output_format"] = str(output_configuration.format)

    # 3. Voice Settings
    if voice_settings and voice_settings is not OMIT:
        if voice_settings.enhance_reference_audio_quality is not None:
             payload["apply_ref_mpsenet"] = voice_settings.enhance_reference_audio_quality
        if voice_settings.maintain_source_accent:
             payload["accent_nudge"] = 0.8

    # 4. Inference Options
    if inference_options and inference_options is not OMIT:
        if inference_options.temperature is not None:
             payload["temperature"] = inference_options.temperature
        if inference_options.inference_steps is not None:
             payload["inference_steps"] = inference_options.inference_steps
            
        if inference_options.speaker_similarity is not None:
             # Formula from user snippet:
             s = max(0.0, min(0.7, inference_options.speaker_similarity))
             payload["campp_speaker_nudge"] = 1.5 * (1 - s / 0.7)

    # 5. Extract additional params (reference_audio, reference_language) from request_options if present
    #    This allows passing 'reference_audio' without breaking the explicit signature for now.
    extra_body = {}
    if request_options and request_options.get("additional_body_parameters"):
        extra_body = request_options.get("additional_body_parameters")
    
    if "reference_audio" not in extra_body:
        raise ValueError("reference_audio is required in additional_body_parameters for Baseten provider")
    if "reference_language" not in extra_body:
        raise ValueError("reference_language is required in additional_body_parameters for Baseten provider")

    payload["reference_language"] = extra_body["reference_language"]
    payload["audio_ref"] = extra_body["reference_audio"]
    payload["reference_audio"] = extra_body["reference_audio"]

    timeout = None
    if request_options and request_options.get("timeout_in_seconds") is not None:
        timeout = request_options.get("timeout_in_seconds")

    # Use the raw httpx client to avoid SDK wrapper injecting unwanted headers/params
    # that might interfere with Baseten's strict endpoint.
    with client_wrapper.httpx_client.httpx_client.stream(
        "POST",
        mars_pro_url,
        json=payload,
        headers=headers,
        timeout=timeout
    ) as _response:
        # Check status manually since we bypassed the wrapper's check
        if not (200 <= _response.status_code < 300):
            # Try to read error body
            _response.read()
            raise Exception(f"Baseten API Error: {_response.status_code} - {_response.text}")

        yield HttpResponse(
            response=_response,
            data=(_chunk for _chunk in _response.iter_bytes(chunk_size=None)),
        )


@contextlib.asynccontextmanager
async def async_baseten_tts(
    client_wrapper: AsyncClientWrapper,
    *,
    text: str,
    language: CreateStreamTtsRequestPayloadLanguage,
    voice_id: typing.Optional[int] = OMIT,
    speech_model: typing.Optional[CreateStreamTtsRequestPayloadSpeechModel] = OMIT,
    user_instructions: typing.Optional[str] = OMIT,
    enhance_named_entities_pronunciation: typing.Optional[bool] = OMIT,
    output_configuration: typing.Optional[StreamTtsOutputConfiguration] = OMIT,
    voice_settings: typing.Optional[StreamTtsVoiceSettings] = OMIT,
    inference_options: typing.Optional[StreamTtsInferenceOptions] = OMIT,
    request_options: typing.Optional[RequestOptions] = None,
) -> typing.AsyncIterator[AsyncHttpResponse[typing.AsyncIterator[bytes]]]:
    # Retrieve API key from provider_params
    provider_params = client_wrapper.provider_params or {}
    api_key = provider_params.get("api_key", "")
    mars_pro_url = provider_params.get("mars_pro_url") or provider_params.get("mars_url")
    if not mars_pro_url:
        raise ValueError("mars_url is required for using Baseten as provider")
    api_key_header_val = f"Api-Key {api_key}"

    # Construct Payload
    # 1. Basic Fields
    payload = {
        "text": text,
        "language": str(language).lower().replace("_", "-"),
        "stream": True,
        "output_format": "wav",  # Default
        "apply_ner_nlp": False, # Default based on doc
    }

    # 2. Output Configuration
    if output_configuration and output_configuration is not OMIT:
        if output_configuration.format:
            payload["output_format"] = str(output_configuration.format)

    # 3. Voice Settings
    if voice_settings and voice_settings is not OMIT:
        if voice_settings.enhance_reference_audio_quality is not None:
             payload["apply_ref_mpsenet"] = voice_settings.enhance_reference_audio_quality
        if voice_settings.maintain_source_accent:
             payload["accent_nudge"] = 0.8

    # 4. Inference Options
    if inference_options and inference_options is not OMIT:
        if inference_options.temperature is not None:
             payload["temperature"] = inference_options.temperature
        if inference_options.inference_steps is not None:
             payload["inference_steps"] = inference_options.inference_steps
            
        if inference_options.speaker_similarity is not None:
             # Formula from user snippet:
             s = max(0.0, min(0.7, inference_options.speaker_similarity))
             payload["campp_speaker_nudge"] = 1.5 * (1 - s / 0.7)

    # 5. Extract additional params (reference_audio, reference_language) from request_options
    extra_body = {}
    if request_options and request_options.get("additional_body_parameters"):
        extra_body = request_options.get("additional_body_parameters")
    
    if "reference_audio" not in extra_body:
        raise ValueError("reference_audio is required in additional_body_parameters for Baseten provider")
    if "reference_language" not in extra_body:
        raise ValueError("reference_language is required in additional_body_parameters for Baseten provider")

    payload["reference_language"] = extra_body["reference_language"]
    payload["audio_ref"] = extra_body["reference_audio"]
    payload["reference_audio"] = extra_body["reference_audio"]

    timeout = None
    if request_options and request_options.get("timeout_in_seconds") is not None:
        timeout = request_options.get("timeout_in_seconds")
    
    # Use the raw httpx client to avoid SDK wrapper injecting unwanted headers/params
    # that might interfere with Baseten's strict endpoint.
    async with client_wrapper.httpx_client.httpx_client.stream(
        "POST",
        mars_pro_url,
        json=payload,
        headers={
            "Authorization": api_key_header_val,
            "content-type": "application/json",
        },
        timeout=timeout
    ) as _response:
        # Check status manually since we bypassed the wrapper's check
        if not (200 <= _response.status_code < 300):
             # Try to read error body
             await _response.aread()
             raise Exception(f"Baseten API Error: {_response.status_code} - {_response.text}")

        yield AsyncHttpResponse(
            response=_response,
            data=(_chunk async for _chunk in _response.aiter_bytes(chunk_size=None)),
        )
