"""Quick test: TTS via generate_content with AUDIO modality."""
import os
import wave
import struct

from google import genai
from google.genai import types

client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

TEXT = """
The year was 2049, and the last lighthouse keeper stood alone
on the edge of the world, watching the storm roll in. She had
been warned that no one was coming. But she lit the lamp anyway.
"""

# Try a few voices
for voice_name in ["Kore", "Charon", "Peri", "Fenrir"]:
    print(f"\n--- Generating with voice: {voice_name} ---")
    try:
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=f"Read this text aloud with feeling:\n\n{TEXT.strip()}",
            config=types.GenerateContentConfig(
                response_modalities=["AUDIO"],
                speech_config=types.SpeechConfig(
                    voice_config=types.VoiceConfig(
                        prebuilt_voice_config=types.PrebuiltVoiceConfig(
                            voice_name=voice_name
                        )
                    )
                ),
            ),
        )

        # Find the audio part
        for part in response.candidates[0].content.parts:
            if hasattr(part, "inline_data") and part.inline_data and part.inline_data.data:
                data = part.inline_data.data
                mime = part.inline_data.mime_type
                print(f"  Got {len(data)} bytes, mime: {mime}")

                # Save raw audio
                raw_path = f"/tmp/genmedia/audio_test_{voice_name}.raw"
                os.makedirs("/tmp/genmedia", exist_ok=True)
                with open(raw_path, "wb") as f:
                    f.write(data)

                # If PCM, wrap in WAV for playback
                if "pcm" in (mime or ""):
                    wav_path = f"/tmp/genmedia/audio_test_{voice_name}.wav"
                    # Gemini PCM is typically 24kHz 16-bit mono little-endian
                    sample_rate = 24000
                    if "rate=" in (mime or ""):
                        rate_str = mime.split("rate=")[1].split(";")[0]
                        sample_rate = int(rate_str)
                    with wave.open(wav_path, "wb") as wf:
                        wf.setnchannels(1)
                        wf.setsampwidth(2)  # 16-bit
                        wf.setframerate(sample_rate)
                        wf.writeframes(data)
                    print(f"  Saved WAV: {wav_path} ({sample_rate}Hz)")
                else:
                    ext = mime.split("/")[-1] if mime else "bin"
                    out_path = f"/tmp/genmedia/audio_test_{voice_name}.{ext}"
                    with open(out_path, "wb") as f:
                        f.write(data)
                    print(f"  Saved: {out_path}")
                break
        else:
            print("  No audio data in response!")
            if response.candidates:
                parts = response.candidates[0].content.parts
                print(f"  Parts: {[type(p).__name__ for p in parts]}")

    except Exception as e:
        print(f"  Error: {e}")
