from langchain_community.document_loaders import ArxivLoader, WikipediaLoader
from langchain_core.tools import tool
from langchain_tavily import TavilySearch


@tool(description="search-web-tool", parse_docstring=True)
async def search_web(query: str, max_results: int = 5) -> str:
    """Searches the web for a query and return formatted results.

    Args:
        query (str): The search query.
        max_results (int, optional): Maximum number of results to return.
                                     Defaults to 5.

    Returns:
        str: A string containing the formatted search results. Each result
             includes a document source, page information, and content.
    """
    # Configure tavily search tool
    tavily_search_tool = TavilySearch(
        max_results=max_results,
        topic="general",
        search_depth="advanced"
    )

    # Execute the search
    search_docs = await tavily_search_tool.arun(query)

    # Format results
    formatted_search_docs = "\n\n---\n\n".join(
        [
            (f'<Document source="{doc.get("url", "No URL")}"/>'
             f'{doc.get("title", "No Title")}\n'
             f'{doc.get("content", "")}\n</Document>')
            for doc in search_docs.get('results', [])
        ]
    )

    return formatted_search_docs


@tool("search-wikipedia-tool", parse_docstring=True)
async def search_wikipedia(query: str, max_results: int = 5) -> str:
    """Search Wikipedia for articles matching the query and
    return formatted results.

    Args:
        query (str): The search query.
        max_results (int, optional): Maximum number of results to return.
                                     Defaults to 5.

    Returns:
        str: A string containing the formatted search results. Each result
             includes a document source, page information, and content.
    """
    search_docs = await WikipediaLoader(
        query=query,
        load_max_docs=max_results,
        load_all_available_meta=True,
        doc_content_chars_max=100000
    ).aload()

    formatted_search_docs = "\n\n---\n\n".join(
        [
            (f'<Document source="{doc.metadata["source"]}"'
             f' page="{doc.metadata.get("page", "")}"/>'
             f'\n{doc.page_content}\n</Document>')
            for doc in search_docs
        ]
    )

    return formatted_search_docs


@tool("search-arxiv-tool", parse_docstring=True)
async def search_arxiv(query: str, max_results: int = 5) -> str:
    """Search arXiv for academic papers matching the query and
    return formatted results.

    Args:
        query (str): The search query.
        max_results (int, optional): Maximum number of results to return.
                                     Defaults to 5.

    Returns:
        str: A string containing the formatted search results. Each result
             includes a document source, page information, and content.
    """
    search_docs = await ArxivLoader(
        query=query,
        load_max_docs=max_results
    ).aload()

    formatted_search_docs = "\n\n---\n\n".join(
        [
            (f'<Document source="{doc.metadata.get("source", "")}"'
             f' page="{doc.metadata.get("page", "")}"/>'
             f'\n{doc.page_content}\n</Document>')
            for doc in search_docs
        ]
    )

    return formatted_search_docs