import gradio as gr

# Import your existing functions
from src import search_on_web, search_custom_sites, search_custom_domain


def run_search_on_web(query: str):
    try:
        result = search_on_web(query)
        return result.response_str, [item.model_dump() for item in result.citation]
    except Exception as e:
        return str(e), []

def run_search_custom_sites(query: str, sites: str):
    try:
        site_list = [s.strip() for s in sites.split(",") if s.strip()]
        result = search_custom_sites(query, site_list)
        return result.response_str, [item.model_dump() for item in result.citation]
    except Exception as e:
        return str(e), []


def run_search_custom_domains(query: str, domains: str):
    try:
        domain_list = [d.strip() for d in domains.split(",") if d.strip()]
        result = search_custom_domain(query, domain_list)
        return result.response_str, [item.model_dump() for item in result.citation]
    except Exception as e:
        return str(e), []

with gr.Blocks() as demo:
    gr.Markdown("# 🔎 MCP Search Panel")

    with gr.Tab("General Web Search"):
        query_input_web = gr.Textbox(label="Query")
        output_web_str = gr.Textbox(label="Answer")
        output_web_cite = gr.JSON(label="Citations")
        search_btn_web = gr.Button("Search")
        search_btn_web.click(
            fn=run_search_on_web,
            inputs=query_input_web,
            outputs=[output_web_str, output_web_cite],
        )

    with gr.Tab("Custom Sites Search"):
        query_input_sites = gr.Textbox(label="Query")
        sites_input = gr.Textbox(label="Sites (comma-separated)")
        output_sites_str = gr.Textbox(label="Answer")
        output_sites_cite = gr.JSON(label="Citations")
        search_btn_sites = gr.Button("Search")
        search_btn_sites.click(
            fn=run_search_custom_sites,
            inputs=[query_input_sites, sites_input],
            outputs=[output_sites_str, output_sites_cite],
        )

    with gr.Tab("Custom Domains Search"):
        query_input_domains = gr.Textbox(label="Query")
        domains_input = gr.Textbox(label="Domains (comma-separated, e.g., edu, gov)")
        output_domains_str = gr.Textbox(label="Answer")
        output_domains_cite = gr.JSON(label="Citations")
        search_btn_domains = gr.Button("Search")
        search_btn_domains.click(
            fn=run_search_custom_domains,
            inputs=[query_input_domains, domains_input],
            outputs=[output_domains_str, output_domains_cite],
        )

if __name__ == "__main__":
    demo.launch(mcp_server=True)