import asyncio
import base64
import io
import wave
from typing import Any, Callable, Coroutine

import sounddevice as sd
import soundfile as sf
from google import genai
from google.genai import types
from typing_extensions import TypedDict
from zrb import AnyContext

from zrb_extras.llm.tool.google.client import get_client
from zrb_extras.llm.tool.google.default_config import TTS_MODEL, VOICE_NAME


class MultiSpeakerVoice(TypedDict):
    speaker: str
    voice: str


def create_speak_tool(
    client: genai.Client | None = None,
    api_key: str | None = None,
    tts_model: str = TTS_MODEL,
    voice_name: str | list[MultiSpeakerVoice] | None = VOICE_NAME,
    sample_rate_out: int = 24000,
    safety_settings: list[types.SafetySetting] | None = None,
    tool_name: str | None = None,
    tool_description: str | None = None,
) -> Callable[
    [AnyContext, str, str | list[MultiSpeakerVoice] | None], Coroutine[Any, Any, bool]
]:
    async def speak(
        ctx: AnyContext,
        text: str,
        voice_name: str | list[MultiSpeakerVoice] | None = voice_name,
    ) -> bool:
        """Converts text to speech and plays it aloud.

        Use this tool to verbally communicate with the user.
        The system will play the generated audio and return `True` upon completion.
        Keep the text concise (a sentence or two) for a faster response.

        You can control the speech's style by describing how you want the text to be spoken.
        For example: 'Say in a cheerful, enthusiastic voice: "Good morning, everyone!"'

        Args:
          text: The text to be spoken. Can be plain text or a controllable prompt.
          voice_name: The voice or voices to use.
            - For a single speaker, provide a voice name string (e.g., "Sulafat").
            - For multiple speakers (up to two), provide a list of speaker-voice mappings.
              Example: `[{"speaker": "User", "voice": "Aoede"}, {"speaker": "Agent", "voice": "Puck"}]`

          Available voices:
          - Zephyr (Female, Bright)
          - Puck (Male, Upbeat)
          - Charon (Male, Informative)
          - Kore (Female, Firm)
          - Fenrir (Male, Excitable)
          - Leda (Female, Youthful)
          - Orus (Male, Firm)
          - Aoede (Female, Breezy)
          - Callirrhoe (Female, Easy-going)
          - Autonoe (Female, Bright)
          - Enceladus (Male, Breathy)
          - Iapetus (Male, Clear)
          - Umbriel (Male, Easy-going)
          - Algieba (Male, Smooth)
          - Despina (Female, Smooth)
          - Erinome (Female, Clear)
          - Algenib (Male, Gravelly)
          - Rasalgethi (Male, Informative)
          - Laomedeia (Female, Upbeat)
          - Achernar (Female, Soft)
          - Alnilam (Male, Firm)
          - Schedar (Male, Even)
          - Gacrux (Female, Mature)
          - Pulcherrima (Female, Forward)
          - Achird (Male, Friendly)
          - Zubenelgenubi (Male, Casual)
          - Vindemiatrix (Female, Gentle)
          - Sadachbia (Male, Lively)
          - Sadaltager (Male, Knowledgeable)
          - Sulafat (Female, Warm)
        Returns:
          True if speech was successfully generated and played, False otherwise.
        """
        return await _synthesize_and_play(
            ctx,
            client=get_client(client, api_key),
            text=text,
            tts_model=tts_model,
            voice_name=voice_name,
            sample_rate_out=sample_rate_out,
            safety_settings=safety_settings,
        )

    if tool_name is not None:
        speak.__name__ = tool_name
    if tool_description is not None:
        speak.__doc__ = tool_description
    return speak


async def _synthesize_and_play(
    ctx: AnyContext,
    client: genai.Client,
    text: str,
    tts_model: str,
    voice_name: str | list[MultiSpeakerVoice] | None = None,
    sample_rate_out: int = 24000,
    safety_settings: list[types.SafetySetting] | None = None,
):
    if not text:
        text = "I have nothing to say."
    ctx.print("Requesting TTS...", plain=True)
    resp = client.models.generate_content(
        model=tts_model,
        contents=text,
        config=types.GenerateContentConfig(
            response_modalities=["AUDIO"],
            safety_settings=safety_settings,
            speech_config=(
                types.SpeechConfig(
                    multi_speaker_voice_config=types.MultiSpeakerVoiceConfig(
                        speaker_voice_configs=[
                            types.SpeakerVoiceConfig(
                                speaker=config.get("speaker", f"Speaker {idx + 1}"),
                                voice_config=types.VoiceConfig(
                                    prebuilt_voice_config=types.PrebuiltVoiceConfig(
                                        voice_name=config.get("voice", VOICE_NAME)
                                    )
                                ),
                            )
                            for idx, config in enumerate(voice_name)
                        ]
                    )
                )
                if isinstance(voice_name, list)
                else types.SpeechConfig(
                    voice_config=types.VoiceConfig(
                        prebuilt_voice_config=types.PrebuiltVoiceConfig(
                            voice_name=voice_name or VOICE_NAME
                        )
                    )
                )
            ),
        ),
    )
    # Extract audio blob
    ctx.print("Extracting audio...", plain=True)
    try:
        part = resp.candidates[0].content.parts[0].inline_data
        b64 = part.data
        mime = getattr(part, "mime_type", None)
    except Exception as e:
        raise RuntimeError(f"No audio found in response: {e}")
    if isinstance(b64, str) and b64.startswith("data:"):
        comma = b64.find(",")
        if comma != -1:
            b64 = b64[comma + 1 :]
    raw = base64.b64decode(b64) if isinstance(b64, str) else bytes(b64)
    # Prepare an in-memory WAV buffer
    ctx.print("Preparing audio buffer...", plain=True)
    buf = io.BytesIO()
    if mime and "wav" in mime.lower():
        # Already WAV or RIFF container
        buf.write(raw)
        buf.seek(0)
    else:
        # Assume raw PCM16 mono at 24kHz and wrap it
        with wave.open(buf, "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)  # 16-bit
            wf.setframerate(sample_rate_out)
            wf.writeframes(raw)
        buf.seek(0)
    # Now decode and play directly from buffer
    data, sr = sf.read(buf)
    ctx.print("Playing audio...", plain=True)
    sd.play(data, sr)
    await asyncio.to_thread(sd.wait)
    return True
