import os
import mimetypes
import struct
import tempfile
import gradio as gr
from dotenv import load_dotenv
from google import genai
from google.genai import types

load_dotenv()
api_key=os.getenv("GEMINI_API_KEY")
if not api_key:
    raise RuntimeError("Set GEMINI_API_KEY in Hugging Face Secrets.")

client=genai.Client(api_key=api_key)

VOICE_GENDER={"Charon":"Male","Fenrir":"Male","Kore":"Female","Aoede":"Female","Leda":"Female"}
CONTEXT_MATRIX={
"Funny":{"Scene":"Stand-up comedy club or a lively local coffee shop discussion","Context":"{gender} voice, highly expressive, comedic timing, laughing naturally inside sentences, playful tone"},
"Serious":{"Scene":"Formal presentation room","Context":"{gender} voice, authoritative, professional tone"},
"Calm":{"Scene":"Meditation sanctuary","Context":"{gender} voice, calm and soothing"},
"Excited":{"Scene":"Sports stadium","Context":"{gender} voice, energetic and enthusiastic"},
}

def wav(raw):
    bps=16;sr=24000;ch=1
    ds=len(raw);ba=ch*(bps//8);br=sr*ba
    head=struct.pack("<4sI4s4sIHHIIHH4sI",b"RIFF",36+ds,b"WAVE",b"fmt ",16,1,ch,sr,br,ba,bps,b"data",ds)
    return head+raw

def matrix(v,e,p,a):
    g=VOICE_GENDER.get(v,"Neutral")
    m=CONTEXT_MATRIX[e]
    s=m["Scene"]
    c=m["Context"].format(gender=g)
    if p=="Fast": c+=", fast pace"
    elif p=="Slow": c+=", slow pace"
    if a!="Neutral":
        c+=f", {a} accent"
    return s,c

def generate(transcript,voice,emotion,pace,accent,temp):
    scene,context=matrix(voice,emotion,pace,accent)
    prompt=f"""Vocal Environment:{scene}
Primary Acting Style:{emotion}
Speed:{pace}
Accent:{accent}
Guide:{context}

Transcript:
{transcript}"""
    cfg=types.GenerateContentConfig(
        temperature=float(temp),
        response_modalities=["audio"],
        speech_config=types.SpeechConfig(
            voice_config=types.VoiceConfig(
                prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name=voice)
            )
        )
    )
    buf=bytearray();native=False
    for chunk in client.models.generate_content_stream(
        model="gemini-2.5-flash-preview-tts",
        contents=[types.Content(role="user",parts=[types.Part.from_text(text=prompt)])],
        config=cfg):
        if chunk.parts and chunk.parts[0].inline_data:
            d=chunk.parts[0].inline_data
            if d.data:
                buf.extend(d.data)
                ext=mimetypes.guess_extension(d.mime_type or "")
                if ext in [".wav",".mp3"]:
                    native=True
    audio=bytes(buf) if native else wav(bytes(buf))
    f=tempfile.NamedTemporaryFile(delete=False,suffix=".wav")
    f.write(audio);f.close()
    return scene,context,f.name

with gr.Blocks(title="Gemini AI TTS Studio Pro") as demo:
    gr.Markdown("# Gemini AI TTS Studio Pro")
    txt=gr.Textbox(lines=8,label="Transcript")
    with gr.Row():
        voice=gr.Dropdown(list(VOICE_GENDER.keys()),value="Aoede",label="Voice")
        emotion=gr.Dropdown(list(CONTEXT_MATRIX.keys()),value="Funny",label="Emotion")
    with gr.Row():
        pace=gr.Dropdown(["Slow","Normal","Fast"],value="Normal",label="Pace")
        accent=gr.Dropdown(["Neutral","American","British","Australian","Jawa","Sunda"],value="Neutral",label="Accent")
    temp=gr.Slider(0,2,value=1,label="Temperature")
    btn=gr.Button("Generate TTS")
    scene=gr.Textbox(label="Scene")
    context=gr.Textbox(label="Context")
    audio=gr.Audio(label="Audio",type="filepath")
    btn.click(generate,[txt,voice,emotion,pace,accent,temp],[scene,context,audio])

if __name__=="__main__":
    demo.launch()