from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
from google import genai
from google.genai import types
import wave
import io
import os
from typing import Optional, List
from pydantic import BaseModel
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

app = FastAPI(
    title="Google GenAI TTS API with Multi-Key Quota Fallback",
    description="TTS API using Google GenAI with multiple API keys and 429 quota handling.",
    version="1.3.0",
    docs_url="/docs",
    redoc_url=None
)

class TTSRequest(BaseModel):
    text: str
    voice_name: Optional[str] = "Kore"
    cheerful: Optional[bool] = True
    sample_rate: Optional[int] = 24000
    channels: Optional[int] = 1
    sample_width: Optional[int] = 2

def get_api_keys() -> List[str]:
    api_keys = os.getenv("GEMINI_API_KEYS")
    if not api_keys:
        raise ValueError("No API keys found in GEMINI_API_KEYS environment variable.")
    return [key.strip() for key in api_keys.split(",") if key.strip()]

def generate_wave_bytes(pcm_data: bytes, channels: int, rate: int, sample_width: int) -> bytes:
    with io.BytesIO() as wav_buffer:
        with wave.open(wav_buffer, "wb") as wf:
            wf.setnchannels(channels)
            wf.setsampwidth(sample_width)
            wf.setframerate(rate)
            wf.writeframes(pcm_data)
        return wav_buffer.getvalue()

def try_generate_tts(api_key, request):
    client = genai.Client(api_key=api_key)
    text_to_speak = f"Say cheerfully: {request.text}" if request.cheerful else request.text

    response = client.models.generate_content(
        model="gemini-2.5-flash-preview-tts",
        contents=text_to_speak,
        config=types.GenerateContentConfig(
            response_modalities=["AUDIO"],
            speech_config=types.SpeechConfig(
                voice_config=types.VoiceConfig(
                    prebuilt_voice_config=types.PrebuiltVoiceConfig(
                        voice_name=request.voice_name,
                    )
                )
            ),
        )
    )

    if not response.candidates or not response.candidates[0].content.parts:
        raise HTTPException(status_code=500, detail="No audio data received from GenAI.")

    audio_data = response.candidates[0].content.parts[0].inline_data.data

    wav_bytes = generate_wave_bytes(
        audio_data,
        channels=request.channels,
        rate=request.sample_rate,
        sample_width=request.sample_width
    )

    return wav_bytes

@app.post("/api/generate-tts/")
async def generate_tts(request: TTSRequest):
    api_keys = get_api_keys()
    last_error = None

    for api_key in api_keys:
        try:
            print(f"Trying API key: {api_key[:5]}...")
            wav_bytes = try_generate_tts(api_key, request)
            return StreamingResponse(
                io.BytesIO(wav_bytes),
                media_type="audio/wav",
                headers={"Content-Disposition": "attachment; filename=generated_audio.wav"}
            )

        except Exception as e:
            error_msg = str(e)
            print(f"Key {api_key[:5]} failed: {error_msg}")

            if "RESOURCE_EXHAUSTED" in error_msg or "quota" in error_msg.lower() or "429" in error_msg:
                # Quota exhausted, try next key
                last_error = error_msg
                continue
            else:
                # Some other error, don't continue
                return JSONResponse(
                    {"status": "error", "message": error_msg},
                    status_code=500
                )

    return JSONResponse(
        {"status": "error", "message": f"All API keys exhausted or invalid. Last error: {last_error}"},
        status_code=429
    )

@app.get("/")
async def root():
    return {"message": "Google GenAI TTS API is running"}

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)