import pytesseract
from PIL import Image
from dotenv import load_dotenv
from langchain.tools import tool
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.tools import TavilySearchResults


@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) -> float:
    """Divide two numbers.
    
    Args:
        a: first int
        b: second int
    """
    try:
        return a / b
    except ZeroDivisionError:
        raise ValueError("Cannot divide by zero.")

@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 up to 3 results.

    Args:
        query: query to search.
    """

    docs = WikipediaLoader(query, load_max_docs=3).load()
    results = "wiki_search results:\n\n"
    results += '\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 docs
    ])
    return results

@tool
def web_search(query: str) -> str:
    """Search the web for a query and return up to 3 results.

    Args:
        query: query to search.
    """
    docs = TavilySearchResults(max_results=3).invoke(query)
    results = "web_search results:\n\n"
    results += '\n\n---\n\n'.join([
        f'<Document source="{doc["url"]}"/>\n{doc["content"]}\n</Document>'
        for doc in docs
    ])
    return results

@tool
def extract_text_from_image(image_path: str) -> str:
    """Extract text from a image.

    Args:
        image_path: path to the image

    Returns:
        extracted text from the image
    """
    try:
        image = Image.open(image_path)

        text = pytesseract.image_to_string(image)
        return f'Extracted text: {text}'
    except Exception as err:
        return f'Error extracting text from the image {image_path}: {str(err)}'

all_tools = [
    multiply,
    add,
    subtract,
    divide,
    modulus,
    wiki_search,
    web_search,
    extract_text_from_image,
]

def main():
    load_dotenv()

    results = wiki_search('What is Dijkstra algorithm?')
    print(results)
    print('-'*80)

    results = web_search('What is TypedDict in python')
    print(results)
    print('-'*80)


if __name__ == '__main__':
    main()