import gradio as gr
import pandas as pd
from langchain_community.utilities import SearxSearchWrapper

from utils.helpers import map_results

from .base import BaseInterfaceWrapper


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

    Args:
        query (str): The search query.
        searxng_host (str): The SearxNG host URL.
        searxng_headers (dict): The SearxNG headers for additional metadata and authentication.
        max_results (int, optional): The number of maximum results to return. Defaults to 5.

    Returns:
        dict: The search results from SearxNG Search Engine. Fields:
            headers (list): The headers of the columns in the search results. Contains "title", "link", "body".
            data (list): The data in the search results. Each row contains the title, link, and body of the search result.
    """
    searxng = SearxSearchWrapper(
        searx_host=searxng_host,
        k=max_results,
    )
    results = await searxng.aresults(
        query,
        num_results=max_results,
    )
    mapping = {
        "title": "title",
        "link": "link",
        "snippet": "body",
    }
    results = await map_results(results, mapping)
    return results


class SearxNGInterfaceWrapper(BaseInterfaceWrapper):
    def __init__(self):
        super().__init__(
            fn=searxng_search,
            inputs=[
                gr.Textbox(label="Search Query"),
                gr.Slider(
                    minimum=1, maximum=10, step=1, value=5, label="Number of Results"
                ),
            ],
            title="SearxNG Search",
            description="Search the web using SearxNG Search Engine.",
        )