import os
import gradio as gr
from langchain_groq import ChatGroq
from langchain_tavily import TavilySearch

GROQ_API_KEY = os.getenv("GROQ_API_KEY")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")

if not GROQ_API_KEY or not TAVILY_API_KEY:
    raise ValueError("Set GROQ_API_KEY and TAVILY_API_KEY in Hugging Face Secrets")

llm = ChatGroq(
    model_name="openai/gpt-oss-120b",
    temperature=0,
    groq_api_key=GROQ_API_KEY
)

search = TavilySearch(
    max_results=5,
    tavily_api_key=TAVILY_API_KEY
)

def google_style_search(query):
    if not query.strip():
        return "Please enter a search query."

    res = search.invoke(query)
    context = "\n".join([r["content"] for r in res.get("results", [])])

    prompt = f"""
Using the following web search information:

{context}

Question: {query}
Answer clearly and concisely like Google Search:
"""

    response = llm.invoke(prompt)
    return response.content.strip()

google_css = """
@import url('https://fonts.googleapis.com/css2?family=Google+Sans:wght@400;500;700&display=swap');

body {
    background-color: #ffffff;
    font-family: 'Google Sans', Arial, sans-serif;
}

.gradio-container {
    max-width: 900px !important;
    margin-top: 40px;
}

#logo {
    font-size: 96px;          
    font-weight: 700;
    text-align: center;
    margin-bottom: 45px;
    letter-spacing: -2px;
    line-height: 1.1;
}

#logo span:nth-child(1) { color: #4285F4; }
#logo span:nth-child(2) { color: #DB4437; }
#logo span:nth-child(3) { color: #F4B400; }
#logo span:nth-child(4) { color: #4285F4; }
#logo span:nth-child(5) { color: #0F9D58; }
#logo span:nth-child(6) { color: #DB4437; }

/* Search box */
.search-box textarea {
    border-radius: 24px !important;
    border: 1px solid #dfe1e5 !important;
    padding: 16px 22px !important;
    font-size: 17px !important;
}

.search-box textarea:focus {
    box-shadow: 0 1px 6px rgba(32,33,36,.28);
    border-color: transparent !important;
}

/* Button */
.search-btn button {
    border-radius: 6px !important;
    background-color: #f8f9fa !important;
    color: #202124 !important;
    border: 1px solid #f8f9fa !important;
    font-weight: 500;
    padding: 10px 18px;
}

.search-btn button:hover {
    border: 1px solid #dadce0 !important;
}

/* Answer box */
.answer-box textarea {
    border-radius: 12px !important;
    background: #f8f9fa !important;
    font-size: 15px !important;
    line-height: 1.7;
}
"""

with gr.Blocks(css=google_css) as demo:

    gr.Markdown(
        """
        <div id="logo">
            <span>G</span><span>o</span><span>o</span>
            <span>g</span><span>l</span><span>e</span>
        </div>
        """,
        elem_id="logo"
    )

    query_input = gr.Textbox(
        placeholder="Search anything...",
        lines=1,
        show_label=False,
        elem_classes="search-box"
    )

    search_button = gr.Button(
        "Search",
        elem_classes="search-btn"
    )

    answer_output = gr.Textbox(
        label="Search Result",
        lines=12,
        interactive=False,
        elem_classes="answer-box"
    )

    search_button.click(
        fn=google_style_search,
        inputs=query_input,
        outputs=answer_output
    )

demo.launch()