import os
import gradio as gr
from google import genai
from google.genai import types
from google.genai.types import GenerateContentConfig, GoogleSearch, Tool

# =========================================================
# Gemini Setup
# =========================================================
API_KEY = os.getenv("Gemini_API_Key")

if not API_KEY:
    raise RuntimeError(
        "❌ Gemini_API_Key is not set. "
        "Add it in Hugging Face → Settings → Secrets."
    )

client = genai.Client(api_key=API_KEY)

# ✅ Valid model
MODEL_ID = "gemini-2.0-flash"

# =========================================================
# Custom CSS + JS
# =========================================================
custom_css = """
#search-btn {
    background-color: #4f46e5 !important;
    color: white !important;
    font-size: 16px;
    border-radius: 10px;
    padding: 10px 18px;
    transition: opacity 0.6s ease, transform 0.2s ease;
}
#search-btn:active {
    opacity: 0.5;
    transform: scale(0.97);
}
"""

custom_js = """
() => {
    const clap = new Audio("https://www.soundjay.com/human/applause-8.mp3");
    clap.play();
}
"""

# =========================================================
# AI Function (BULLETPROOF)
# =========================================================
def google_search_query(question):

    # ---------- 1️⃣ Normalize safely ----------
    if question is None:
        question = ""
    else:
        question = str(question)

    question = question.strip()

    if question == "":
        return "Please type a question above 👆", ""

    try:
        # ---------- 2️⃣ Define Search Tool ----------
        google_search_tool = Tool(
            google_search=GoogleSearch()
        )

        # ---------- 3️⃣ Generate Response ----------
        response = client.models.generate_content(
            model=MODEL_ID,
            contents=[
                types.Content(
                    role="user",
                    parts=[types.Part.from_text(question)]
                )
            ],
            config=GenerateContentConfig(
                tools=[google_search_tool]
            ),
        )

        # ---------- 4️⃣ Extract AI Text Safely ----------
        ai_response = ""

        if hasattr(response, "text") and response.text:
            ai_response = response.text
        elif response.candidates:
            try:
                ai_response = response.candidates[0].content.parts[0].text
            except Exception:
                ai_response = "No AI response generated."
        else:
            ai_response = "No AI response generated."

        # ---------- 5️⃣ Extract Search Grounding Safely ----------
        search_results = ""

        try:
            candidate = response.candidates[0]
            if (
                hasattr(candidate, "grounding_metadata")
                and candidate.grounding_metadata
                and candidate.grounding_metadata.search_entry_point
            ):
                search_results = (
                    candidate
                    .grounding_metadata
                    .search_entry_point
                    .rendered_content
                )
        except Exception:
            search_results = ""

        return ai_response, search_results

    except Exception as e:
        return f"❌ Error: {str(e)}", ""


# =========================================================
# Gradio App
# =========================================================
with gr.Blocks(css=custom_css) as app:
    gr.Markdown("## 🔍 Google Search with Gemini AI")

    question = gr.Textbox(
        lines=2,
        label="Ask a Question",
        placeholder="e.g. What are types of machine learning?"
    )

    search_btn = gr.Button("🔍 Search", elem_id="search-btn")

    ai_output = gr.Textbox(label="AI Response")
    search_output = gr.HTML(label="Search Results")

    # Button click
    search_btn.click(
        fn=google_search_query,
        inputs=[question],
        outputs=[ai_output, search_output],
        js=custom_js
    )

    # Enter key submit
    question.submit(
        fn=google_search_query,
        inputs=[question],
        outputs=[ai_output, search_output],
    )

# =========================================================
# Launch (HF SAFE)
# =========================================================
if __name__ == "__main__":
    app.queue().launch()