# Example: Copy tool implementations from sample_agent.tools here

# Math tools
def multiply(a: int, b: int) -> int:
    """Multiply two numbers.

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

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

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

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

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

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

    Args:
        a: first int
        b: second int
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

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

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

# Wikipedia search tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 results.

    Args:
        query: The search query."""
    try:
        from langchain_community.document_loaders import WikipediaLoader
        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
        formatted_search_docs = "\n\n---\n\n".join(
            [
                f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
                for doc in search_docs
            ])
        return formatted_search_docs
    except Exception as e:
        return f"Error in wiki_search: {e}"

# Web search tool
def web_search(query: str) -> str:
    """Search Tavily for a query and return maximum 3 results.

    Args:
        query: The search query."""
    try:
        from langchain_community.tools.tavily_search import TavilySearchResults
        search_tool = TavilySearchResults(max_results=3)
        search_docs = search_tool.invoke({"query": query})
        # Each doc is a dict, not an object with .metadata/.page_content
        formatted_search_docs = "\n\n---\n\n".join(
            [
                f'<Document source="{doc.get("source", "")}" page="{doc.get("page", "")}"/>\n{doc.get("content", "")}\n</Document>'
                for doc in search_docs
            ])
        return formatted_search_docs
    except Exception as e:
        return f"Error in web_search: {e}"

# Arxiv search tool
def arvix_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 result.

    Args:
        query: The search query."""
    try:
        from langchain_community.document_loaders import ArxivLoader
        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
        formatted_search_docs = "\n\n---\n\n".join(
            [
                f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
                for doc in search_docs
            ])
        return formatted_search_docs
    except Exception as e:
        return f"Error in arvix_search: {e}"

TOOL_REGISTRY = {
    "multiply": {
        "description": "Multiply two numbers. Usage: multiply(a, b)",
        "function": multiply,
    },
    "add": {
        "description": "Add two numbers. Usage: add(a, b)",
        "function": add,
    },
    "subtract": {
        "description": "Subtract two numbers. Usage: subtract(a, b)",
        "function": subtract,
    },
    "divide": {
        "description": "Divide two numbers. Usage: divide(a, b)",
        "function": divide,
    },
    "modulus": {
        "description": "Get the modulus of two numbers. Usage: modulus(a, b)",
        "function": modulus,
    },
    "wiki_search": {
        "description": "Search Wikipedia for a query and return up to 2 results. Usage: wiki_search(query)",
        "function": wiki_search,
    },
    "web_search": {
        "description": "Search Tavily for a query and return up to 3 results. Usage: web_search(query)",
        "function": web_search,
    },
    "arvix_search": {
        "description": "Search Arxiv for a query and return up to 3 results. Usage: arvix_search(query)",
        "function": arvix_search,
    },
}