import wikipedia
from langchain_core.tools import tool
from loguru import logger


@tool("wikipedia_search_tool", parse_docstring=True)
def wikipedia_search(query: str) -> str:
    """
    Searches Wikipedia for the given query.

    Args:
        query (str): The search query to look up on Wikipedia.

    Returns:
        str: A formatted string with the search results, page title and url.
    """
    logger.info(f"use wikipedia_search_tool with param: {query}")

    search_results = wikipedia.search(query, results=5)

    if not search_results:
        return "No results found for the query."

    result_text = ""
    try:
        for i, title in enumerate(search_results, 1):
            page = wikipedia.page(search_results[i - 1], auto_suggest=False)
            result_text += f"{i}. [{title}]({page.url})\n"

        return result_text

    except wikipedia.DisambiguationError as e:
        return "Disambiguation page found. Possible matches:\n{}".format('\n'.join(e.options))
    except wikipedia.PageError as e:
        return f"Page not found. Try another search term."
    except Exception as e:
        return f"An error occurred: {str(e)}"


if __name__ == "__main__":
    print(wikipedia_search.invoke("Mercedes Sosa discography"))