import logging
import os
import base64
import uuid
import wave
from typing import Optional, List, Dict, Any, TYPE_CHECKING
from google import genai
from google.genai import types as genai_types

from autobyteus.multimedia.audio.base_audio_client import BaseAudioClient
from autobyteus.multimedia.utils.response_types import SpeechGenerationResponse
from autobyteus.utils.gemini_helper import initialize_gemini_client_with_runtime
from autobyteus.utils.gemini_model_mapping import resolve_model_for_runtime

if TYPE_CHECKING:
    from autobyteus.multimedia.audio.audio_model import AudioModel
    from autobyteus.multimedia.utils.multimedia_config import MultimediaConfig

logger = logging.getLogger(__name__)


_AUDIO_TEMP_DIR = "/tmp/autobyteus_audio"

_AUDIO_MIME_EXTENSION_MAP = {
    "audio/wav": "wav",
    "audio/x-wav": "wav",
    "audio/mpeg": "mp3",
    "audio/mp3": "mp3",
    "audio/ogg": "ogg",
    "audio/webm": "webm",
}


def _save_audio_bytes_to_wav(pcm_bytes: bytes, channels=1, rate=24000, sample_width=2) -> str:
    """Saves PCM audio bytes to a temporary WAV file and returns the path."""
    os.makedirs(_AUDIO_TEMP_DIR, exist_ok=True)
    file_path = os.path.join(_AUDIO_TEMP_DIR, f"{uuid.uuid4()}.wav")
    
    try:
        with wave.open(file_path, "wb") as wf:
            wf.setnchannels(channels)
            wf.setsampwidth(sample_width)
            wf.setframerate(rate)
            wf.writeframes(pcm_bytes)
        logger.info(f"Successfully saved generated audio to {file_path}")
        return file_path
    except Exception as e:
        logger.error(f"Failed to save audio to WAV file at {file_path}: {e}")
        raise


def _save_audio_bytes(audio_bytes: bytes, extension: Optional[str]) -> str:
    """Saves audio bytes to a temporary file and returns the path."""
    os.makedirs(_AUDIO_TEMP_DIR, exist_ok=True)
    suffix = (extension or "bin").lstrip(".")
    file_path = os.path.join(_AUDIO_TEMP_DIR, f"{uuid.uuid4()}.{suffix}")
    try:
        with open(file_path, "wb") as audio_file:
            audio_file.write(audio_bytes)
        logger.info(f"Successfully saved generated audio to {file_path}")
        return file_path
    except Exception as e:
        logger.error(f"Failed to save audio to file at {file_path}: {e}")
        raise


def _parse_mime_type(mime_type: Optional[str]) -> tuple[str, Dict[str, str]]:
    if not mime_type:
        return "", {}
    parts = [part.strip() for part in mime_type.split(";") if part.strip()]
    base = parts[0].lower() if parts else ""
    params: Dict[str, str] = {}
    for part in parts[1:]:
        if "=" in part:
            key, value = part.split("=", 1)
            params[key.strip().lower()] = value.strip()
    return base, params


def _coerce_audio_bytes(audio_data: Any) -> bytes:
    if audio_data is None:
        return b""
    if isinstance(audio_data, bytes):
        return audio_data
    if isinstance(audio_data, bytearray):
        return bytes(audio_data)
    if isinstance(audio_data, memoryview):
        return audio_data.tobytes()
    if isinstance(audio_data, str):
        return base64.b64decode(audio_data)
    return bytes(audio_data)


class GeminiAudioClient(BaseAudioClient):
    """
    An audio client that uses Google's Gemini models for audio tasks.

    **Setup Requirements:**
    1.  **Vertex AI Express Mode:** Set `VERTEX_AI_API_KEY`.
    2.  **Vertex AI Mode:** Set `VERTEX_AI_PROJECT` and `VERTEX_AI_LOCATION`.
    3.  **AI Studio Mode:** Set `GEMINI_API_KEY`.
    """

    def __init__(self, model: "AudioModel", config: "MultimediaConfig"):
        super().__init__(model, config)
        
        try:
            self.client, self.runtime_info = initialize_gemini_client_with_runtime()
            self.async_client = self.client.aio
            logger.info(f"GeminiAudioClient initialized for model '{self.model.name}'.")
        except Exception as e:
            logger.error(f"Failed to configure Gemini client: {e}")
            raise RuntimeError(f"Failed to configure Gemini client: {e}")


    async def generate_speech(
        self,
        prompt: str,
        generation_config: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> SpeechGenerationResponse:
        """
        Generates spoken audio from text using a Gemini TTS model, supporting single-speaker,
        multi-speaker, and style-controlled generation.
        """
        try:
            final_config = self.config.to_dict().copy()
            if generation_config:
                final_config.update(generation_config)
            
            # Handle style instructions by prepending them to the prompt
            style_instructions = final_config.get("style_instructions")
            final_prompt = f"{style_instructions}: {prompt}" if style_instructions else prompt
            logger.debug(f"Final prompt for TTS: '{final_prompt[:100]}...'")

            speech_config = None
            mode = final_config.get("mode", "single-speaker")

            # Handle multi-speaker generation
            if mode == "multi-speaker":
                speaker_mapping_list = final_config.get("speaker_mapping")
                if not speaker_mapping_list or not isinstance(speaker_mapping_list, list):
                    raise ValueError("Multi-speaker mode requires a 'speaker_mapping' list in generation_config.")
                
                logger.info(f"Configuring multi-speaker TTS with mapping: {speaker_mapping_list}")
                speaker_voice_configs = []
                for mapping_item in speaker_mapping_list:
                    speaker = mapping_item.get("speaker")
                    voice_name = mapping_item.get("voice")
                    if not speaker or not voice_name:
                        logger.warning(f"Skipping invalid item in speaker_mapping list: {mapping_item}")
                        continue
                    
                    speaker_voice_configs.append(
                        genai_types.SpeakerVoiceConfig(
                            speaker=speaker,
                            voice_config=genai_types.VoiceConfig(
                                prebuilt_voice_config=genai_types.PrebuiltVoiceConfig(voice_name=voice_name)
                            )
                        )
                    )
                
                if not speaker_voice_configs:
                    raise ValueError("The 'speaker_mapping' list was empty or contained no valid mappings.")

                speech_config = genai_types.SpeechConfig(
                    multi_speaker_voice_config=genai_types.MultiSpeakerVoiceConfig(speaker_voice_configs=speaker_voice_configs)
                )

            # Handle single-speaker generation (default)
            else:
                voice_name = final_config.get("voice_name", "Kore") # A default voice
                logger.info(f"Configuring single-speaker TTS with voice: '{voice_name}'")
                speech_config = genai_types.SpeechConfig(
                    voice_config=genai_types.VoiceConfig(
                        prebuilt_voice_config=genai_types.PrebuiltVoiceConfig(voice_name=voice_name)
                    )
                )

            # The google-genai library's TTS endpoint uses a synchronous call.
            # FIX: Ensure no 'models/' prefix is used here.
            runtime_adjusted_model = resolve_model_for_runtime(
                self.model.value,
                modality="tts",
                runtime=getattr(self, "runtime_info", None) and self.runtime_info.runtime,
            )
            logger.info(
                "Generating speech with Gemini TTS model '%s' (requested '%s').",
                runtime_adjusted_model,
                self.model.value,
                )
            resp = self.client.models.generate_content(
                model=runtime_adjusted_model,
                contents=final_prompt,
                config=genai_types.GenerateContentConfig(
                    response_modalities=["AUDIO"],
                    speech_config=speech_config
                ),
            )
            
            part = resp.candidates[0].content.parts[0]
            inline_data = part.inline_data
            if not inline_data or not inline_data.data:
                raise ValueError("Gemini TTS response did not include audio data.")

            mime_type, mime_params = _parse_mime_type(inline_data.mime_type)
            audio_bytes = _coerce_audio_bytes(inline_data.data)
            if not audio_bytes:
                raise ValueError("Gemini TTS returned empty audio data.")

            logger.info(
                "Received Gemini TTS audio payload (mime_type='%s', bytes=%d).",
                mime_type or "unknown",
                len(audio_bytes),
            )

            if not mime_type or mime_type.startswith("audio/pcm") or mime_type == "audio/l16":
                rate = 24000
                channels = 1
                if "rate" in mime_params:
                    try:
                        rate = int(mime_params["rate"])
                    except ValueError:
                        logger.warning("Invalid sample rate in mime_type '%s'; using default 24000.", inline_data.mime_type)
                if "channels" in mime_params:
                    try:
                        channels = int(mime_params["channels"])
                    except ValueError:
                        logger.warning("Invalid channel count in mime_type '%s'; using default 1.", inline_data.mime_type)

                audio_path = _save_audio_bytes_to_wav(audio_bytes, channels=channels, rate=rate, sample_width=2)
            else:
                extension = _AUDIO_MIME_EXTENSION_MAP.get(mime_type, "bin")
                audio_path = _save_audio_bytes(audio_bytes, extension)

            return SpeechGenerationResponse(audio_urls=[audio_path])

        except Exception as e:
            logger.error(f"Error during Google Gemini speech generation: {str(e)}")
            raise ValueError(f"Google Gemini speech generation failed: {str(e)}")

    async def cleanup(self):
        logger.debug("GeminiAudioClient cleanup called.")
