import json

from langchain_community.document_loaders import WikipediaLoader
from langchain_core.tools import tool


@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers.

    Args:
        a: first int
        b: second int
    """
    return a * b


@tool
def add(a: int, b: int) -> int:
    """
    Add two numbers.

    Args:
        a: first int
        b: second int
    """
    return a + b


@tool
def subtract(a: int, b: int) -> int:
    """
    Subtract two numbers.

    Args:
        a: first int
        b: second int
    """
    return a - b


@tool
def divide(a: int, b: int) -> int:
    """
    Divide two numbers.

    Args:
        a: numerator
        b: denominator
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b


@tool
def modulus(a: int, b: int) -> int:
    """
    Get the modulus of two numbers.

    Args:
        a: first int
        b: second int
    """
    return a % b


@tool
def wiki_search(query: str) -> str:
    """
    Search Wikipedia for a query and return maximum 2 results. The results
    are formatted as a JSON string with the source and content of each
    document.

    Args:
        query: The search query."""
    search_docs = WikipediaLoader(
        query=query,
        load_max_docs=2,
    ).load()
    formatted_search_docs = [
        {
            "source": doc.metadata["source"],
            "content": doc.page_content,
        }
        for doc in search_docs
    ]
    return json.dumps({"wiki_results": formatted_search_docs})