import gradio as gr
import pandas as pd
from langchain_tavily import TavilySearch

from utils.helpers import map_results


async def tavily_search(query: str, api_key: str, max_results: int = 5) -> pd.DataFrame:
    """
    Given a search query, returns the search results from Tavily.

    Args:
        query (str): The search query.
        api_key (str): The API key for Tavily Search.
        max_results (int, optional): The number of maximum results to return. Defaults to 5.

    Returns:
        str: The search results from Tavily Search Engine.
    """
    tavily = TavilySearch(
        tavily_api_key=api_key,
    )
    tavily.max_results = max_results
    results = await tavily.ainvoke(query)
    mapping = {
        "title": "title",
        "url": "link",
        "content": "body",
    }
    results = await map_results(results["results"], mapping)
    return results


tavily_interface = gr.Interface(
    fn=tavily_search,
    inputs=[
        gr.Textbox(label="Search Query"),
        gr.Textbox(label="API Key", type="password"),
        gr.Slider(minimum=1, maximum=10, step=1, value=5, label="Number of Results"),
    ],
    outputs=gr.Dataframe(
        label="Search Results",
        headers=["title", "body", "link"],
        show_fullscreen_button=True,
        show_row_numbers=True,
        show_copy_button=True,
        wrap=True,
    ),
    title="Tavily Search",
    description="Search the web using Tavily Search Engine.",
)