import gradio as gr
import os
import google.generativeai as genai
import requests
import json


# ------------------- MODEL MANAGER -------------------

def get_available_gemini_models(api_key):
    url = "https://generativelanguage.googleapis.com/v1beta/models"
    response = requests.get(f"{url}?key={api_key}")
    if response.status_code == 200:
        models = response.json().get("models", [])
        return [m["name"] for m in models if "generateContent" in m.get("supportedGenerationMethods", [])]
    return []

def pick_best_gemini_model(models):
    preferred = [
        "gemini-1.5-pro",
        "gemini-1.5-flash",
        "gemini-1.0-pro-vision-latest",
        "gemini-1.5-flash-002",
        "gemini-2.5-pro",
        "gemini-2.5-flash"
    ]
    for name in preferred:
        for model in models:
            if name in model:
                return model
    return models[0] if models else None

def pick_best_groq_model(models):
    # prioritize based on Groq's known models
    preferred = ["llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma-7b-it"]
    for name in preferred:
        for model in models:
            if name in model:
                return model
    return models[0] if models else None

# ------------------- MAIN FUNCTION -------------------

def generate_reply(user_input, api_key, model_type):
    if model_type == "Gemini":
        genai.configure(api_key=api_key)
        models = get_available_gemini_models(api_key)
        selected_model = pick_best_gemini_model(models)

        if not selected_model:
            return "❌ No compatible Gemini model found."

        model = genai.GenerativeModel(selected_model)
        try:
            response = model.generate_content(user_input)
            return response.text if hasattr(response, "text") else str(response)
        except Exception as e:
            return f"❌ Gemini API Error: {str(e)}"

    elif model_type == "Groq":
        client = Groq(api_key=api_key)
        try:
            models = [m.id for m in client.models.list().data]
            selected_model = pick_best_groq_model(models)

            if not selected_model:
                return "❌ No compatible Groq model found."

            response = client.chat.completions.create(
                model=selected_model,
                messages=[{"role": "user", "content": user_input}],
            )
            return response.choices[0].message.content
        except Exception as e:
            return f"❌ Groq API Error: {str(e)}"

    else:
        return "❌ Invalid model type."

# ------------------- GRADIO UI -------------------

with gr.Blocks() as demo:
    gr.Markdown("## 🔊 Voice-to-LLM Assistant")
    with gr.Row():
        api_key = gr.Textbox(label="🔑 API Key (Gemini / Groq)", type="password")
        model_type = gr.Dropdown(choices=["Gemini", "Groq"], label="🤖 Choose Model Type", value="Gemini")

    mic_input = gr.Audio(type="filepath", label="🎤 Speak Now", interactive=True)
    transcribed_text = gr.Textbox(label="📝 Transcribed Text")
    response_text = gr.Textbox(label="💡 AI Response")

    def transcribe_and_generate(audio, api_key, model_type):
        import whisper
        model = whisper.load_model("base")
        result = model.transcribe(audio)
        text = result["text"]
        reply = generate_reply(text, api_key, model_type)
        return text, reply

    btn = gr.Button("🚀 Generate")
    btn.click(fn=transcribe_and_generate, inputs=[mic_input, api_key, model_type], outputs=[transcribed_text, response_text])

demo.launch()