from flask import Flask, request, jsonify, send_file
from google import genai
from google.genai import types
import os
import wave
import tempfile

# ─── 1. Read & validate API key from environment ────────────────────────────────
api_key = os.getenv("GOOGLE_API_KEY")
if not api_key:
    raise ValueError("Environment variable 'GOOGLE_API_KEY' not found.")
client = genai.Client(api_key=api_key)

# ─── 2. Helper to write raw PCM bytes into a .wav file ────────────────────────────
def wave_file(pcm_bytes: bytes, channels: int = 1, rate: int = 24000, sample_width: int = 2) -> str:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
        with wave.open(tmp_file.name, "wb") as wf:
            wf.setnchannels(channels)
            wf.setsampwidth(sample_width)
            wf.setframerate(rate)
            wf.writeframes(pcm_bytes)
        return tmp_file.name

# ─── 3. Main TTS Function ────────────────────────────────────────────────────────
def generate_audio(
    model_choice: str,
    language: str,
    accent: str,
    tone: str,
    speed: str,
    clarity: str,
    style: str,
    instructions: str,
    text: str,
    voice_name: str,
    custom_additions: str,
) -> str:
    prompt = f"""
[Language]: {language}
[Accent]: {accent}
[Tone / Emotion]: {tone}
[Speech Speed]: {speed}
[Clarity]: {clarity}
[Style]: {style}
[Instructions to speaker]: {instructions}
{f"[Custom Attributes]: {custom_additions}" if custom_additions.strip() else ""}

[Text]:
{text}
"""
    response = client.models.generate_content(
        model=model_choice,
        contents=prompt,
        config=types.GenerateContentConfig(
            response_modalities=["AUDIO"],
            speech_config=types.SpeechConfig(
                voice_config=types.VoiceConfig(
                    prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name=voice_name)
                )
            )
        ),
    )
    pcm_data = response.candidates[0].content.parts[0].inline_data.data
    return wave_file(pcm_data)

# ─── 4. Flask REST API ───────────────────────────────────────────────────────────
app = Flask(__name__)

@app.route("/api/tts", methods=["POST"])
def tts_api():
    try:
        data = request.get_json()
        file_path = generate_audio(
            data["model_choice"],
            data["language"],
            data["accent"],
            data["tone"],
            data["speed"],
            data["clarity"],
            data["style"],
            data["instructions"],
            data["text"],
            data["voice_name"],
            data["custom_additions"]
        )
        # Return the .wav file as attachment
        return send_file(file_path, mimetype="audio/wav", as_attachment=True)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    # Listen on 0.0.0.0:7860 so HF Spaces can route traffic here
    app.run(host="0.0.0.0", port=7860)