import gradio as gr
from phi.assistant import Assistant
from phi.llm.openai import OpenAIChat
from phi.llm.groq import Groq
from phi.llm.ollama import Ollama
from phi.tools.duckduckgo import DuckDuckGo

# 模型字典
model_choices = {
    "OpenAI Model (gpt-3.5-turbo)": {
        "class": OpenAIChat,
        "params": {
            "model": "gpt-3.5-turbo"
        }
    },
    "Groq Model (llama3-70b-8192)": {
        "class": Groq,
        "params": {
            "model": "llama3-70b-8192"
        }
    },
    "Ollama Model (llama3)": {
        "class": Ollama,
        "params": {
            "model": "llama3"
        }
    }
}

def ai_search(api_key, model_choice, base_url, query):
    if model_choice in ["OpenAI Model (gpt-3.5-turbo)", "Groq Model (llama3-70b-8192)"] and not api_key:
        return "API key is required."
    
    if model_choice == "OpenAI Model (gpt-3.5-turbo)" and not base_url:
        return "Base URL is required for OpenAI Model."

    # 动态创建模型实例
    model_info = model_choices[model_choice]
    llm_class = model_info["class"]
    llm_params = model_info["params"]

    if model_choice != "Ollama Model (llama3)":
        llm_params["api_key"] = api_key

    if model_choice == "OpenAI Model (gpt-3.5-turbo)":
        llm_params["base_url"] = base_url

    llm_instance = llm_class(**llm_params)

    # 创建 Assistant 实例
    assistant = Assistant(llm=llm_instance, tools=[DuckDuckGo()], show_tool_calls=True)

    # 使用 DuckDuckGo 进行搜索
    search_results_generator = assistant.run(f"Search: {query}")

    if not search_results_generator:
        return "No results found."
    
    # 迭代生成器并合并结果
    search_results = "".join(list(search_results_generator))
    
    # 打印搜索结果以检查其结构
    print(search_results)

    # 格式化搜索结果为 Markdown
    markdown_results = f"### Search Results for: {query}\n\n"
    markdown_results += "Here are some results:\n\n"
    markdown_results += f"{search_results}\n"
    
    return markdown_results

def create_interface():
    with gr.Blocks() as iface:
        gr.Markdown("## AI Search Engine")
        
        model_choice_input = gr.Dropdown(label="Choose Model", choices=list(model_choices.keys()))
        api_key_input = gr.Textbox(label="API Key", placeholder="Enter your API key here", type="password", visible=False)
        base_url_input = gr.Textbox(label="Base URL", placeholder="Enter the base URL here", visible=False)
        query_input = gr.Textbox(label="Search Query", placeholder="Enter your search query here")
        
        search_button = gr.Button("Search", interactive=False)
        
        output = gr.Textbox(label="Response", interactive=False)
        
        def update_inputs(model_choice):
            if model_choice == "OpenAI Model (gpt-3.5-turbo)":
                return gr.update(visible=True), gr.update(visible=True), "", ""
            elif model_choice == "Groq Model (llama3-70b-8192)":
                return gr.update(visible=True), gr.update(visible=False), "", ""
            else:
                return gr.update(visible=False), gr.update(visible=False), "", ""

        def check_button_state(api_key, base_url, model_choice, query):
            # print(api_key)
            # print(base_url)
            # print(model_choice)
            # if model_choice == "Ollama Model (llama3)":
            #     return gr.update(interactive=bool(query))
            # elif model_choice == "Groq Model (llama3-70b-8192)":
            #     return gr.update(interactive=True)
            # elif model_choice == "OpenAI Model (gpt-3.5-turbo)":
            #     return gr.update(interactive=bool(api_key and base_url and query))
            # else:
            #     return gr.update(interactive=True)
            return gr.update(interactive=True)

        def combined_update(model_choice, api_key, base_url, query):
            updates = update_inputs(model_choice)
            button_state = check_button_state(api_key, base_url, model_choice, query)
            return updates + (button_state,)

        model_choice_input.change(
            lambda model_choice: combined_update(model_choice, api_key_input.value, base_url_input.value, query_input.value),
            inputs=model_choice_input,
            outputs=[api_key_input, base_url_input, api_key_input, base_url_input, search_button]
        )
        api_key_input.change(
            lambda api_key: check_button_state(api_key, base_url_input.value, model_choice_input.value, query_input.value),
            inputs=api_key_input,
            outputs=search_button
        )
        base_url_input.change(
            lambda base_url: check_button_state(api_key_input.value, base_url, model_choice_input.value, query_input.value),
            inputs=base_url_input,
            outputs=search_button
        )
        query_input.change(
            lambda query: check_button_state(api_key_input.value, base_url_input.value, model_choice_input.value, query),
            inputs=query_input,
            outputs=search_button
        )
        
        search_button.click(
            fn=ai_search,
            inputs=[api_key_input, model_choice_input, base_url_input, query_input],
            outputs=output
        )
        
    return iface

iface = create_interface()
iface.launch()