import os
import serpapi
import gradio as gr
from langchain_google_genai import ChatGoogleGenerativeAI

from utils.utils import get_top_results, get_prompt


def main(search_query, search_location, serp_api_key):
    search_response = serpapi.search(
        q=search_query,
        location=search_location,
        api_key=serp_api_key
    ) # engine="google", location="Austin, Texas", hl="en", gl="us"
    
    top_results = get_top_results(search_response)

    prompt = get_prompt(search_query, top_results)

    chat_model = ChatGoogleGenerativeAI(
        model="gemini-pro",
        convert_system_message_to_human=True,
        temperature=0.0, # temperature=0.7 (default)
        # top_p=0.5,
        ) 
    response = chat_model.invoke(prompt)

    return response.content


if __name__ == "__main__":
    try:
        app = gr.Interface(
            fn=main,
            inputs=[
                gr.Textbox(label="Search Query (検索クエリ)"),
                gr.Textbox(label="Search Location (検索を行う国を英語で入力してください)"),
                gr.Textbox(label="SERP API Key (https://serpapi.com/manage-api-key)")
            ],
            outputs=[gr.Textbox(label="Search Result (検索結果)")],
            title="Google Search enhanced by LLM"
        )
        app.launch(share=True)
    except Exception as e:
        raise e