import cmath
import json

from langchain_community.document_loaders import ArxivLoader
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader


@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 page links.
    Args:
        query: The search query."""
    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
    #     ]
    # )
    formatted_search_docs = [
        {"source": doc.metadata.get("source", ""), "page": doc.metadata.get("page", "")}
        for doc in search_docs
    ]
    return {"wiki_page_links": json.dumps(formatted_search_docs)}


@tool
def web_search(query: str) -> str:
    """Search Tavily for a query and return maximum 3 page links.
    Args:
        query: The search query."""
    search_docs = TavilySearchResults(
        max_results=3,
        tavily_api_key="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",
    ).invoke(query)
    # formatted_search_docs = "\n\n---\n\n".join(
    #     [
    #         f'<Document source="{doc.get("url", "")}" title="{doc.get("title", "")}"/>\n{doc.get("content", "")}\n</Document>'
    #         for doc in search_docs
    #     ]
    # )
    formatted_search_docs = [
        {"source": doc.get("url", ""), "page": doc.get("title", "")}
        for doc in search_docs
    ]
    return {"web_page_links": json.dumps(formatted_search_docs)}

@tool
def visit_webpage(url: str) -> str:
    """Retrieve content from a webpage using page links. A good option when you need detailed information from a specific webpage.
    Args:
        url: The URL to retrieve content from."""
    search_docs = WebBaseLoader(url).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata.get("source", "")}" title="{doc.metadata.get("title", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ]
    )
    return {"web_page_content": formatted_search_docs}


@tool
def arxiv_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 result.
    Args:
        query: The search query."""
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["Title"]}" page="{doc.metadata.get("Summary", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ]
    )
    return {"arxiv_results": formatted_search_docs}

@tool
def translate_to_english(text: str) -> str:
    """Translate the given text to English.
    Args:
        text: The text to translate to English."""
    from langchain_ollama import ChatOllama

    llm = ChatOllama(
        model="tinyllama:1.1b",
        api_base="http://localhost:11434",
    )

    prompt = f"Translate the following text to English:\n\n{text}"
    response = llm.invoke(
        [
            {"role": "system", "content": "You are a helpful assistant that translates text to English."},
            {"role": "user", "content": prompt},
        ]
    )
    return {"translated_text": response.content}

@tool
def multiply(a: float, b: float) -> float:
    """
    Multiplies two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a * b


@tool
def add(a: float, b: float) -> float:
    """
    Adds two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a + b


@tool
def subtract(a: float, b: float) -> int:
    """
    Subtracts two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a - b


@tool
def divide(a: float, b: float) -> float:
    """
    Divides two numbers.
    Args:
        a (float): the first float number
        b (float): the second float number
    """
    if b == 0:
        raise ValueError("Cannot divided by zero.")
    return a / b


@tool
def modulus(a: int, b: int) -> int:
    """
    Get the modulus of two numbers.
    Args:
        a (int): the first number
        b (int): the second number
    """
    return a % b


@tool
def power(a: float, b: float) -> float:
    """
    Get the power of two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a**b


@tool
def square_root(a: float) -> float | complex:
    """
    Get the square root of a number.
    Args:
        a (float): the number to get the square root of
    """
    if a >= 0:
        return a**0.5
    return cmath.sqrt(a)