from langchain_community.tools import DuckDuckGoSearchRun
from langchain_community.document_loaders import (
    WikipediaLoader,
    ArxivLoader,
    ImageCaptionLoader,
)
from langchain_core.tools import tool
from typing import List
import math

# Web_search Tool
search_tool = DuckDuckGoSearchRun()


# Calulator Tool
@tool(
    "calculator",
    description="Performs arithmetic calculations. Use this for any math problems.",
)
def calc(expression: str) -> str:
    """Evaluate mathematical expressions.
    Args:
        expression: expression to evaluate
    """
    try:
        cleaned = expression.replace("^", "**").replace(",", "")
        safe_ns = {k: getattr(math, k) for k in dir(math) if not k.startswith("_")}
        safe_ns["__builtins__"] = {}
        result = eval(cleaned, safe_ns)
        return str(result)
    except Exception as e:
        return f"Calculation error: {e}"


@tool
def wiki_search(query: str) -> str:
    """
    Search Wikipedia for a query and return maximum 2 results.
    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
        ]
    )
    return {"wiki_results": 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["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ]
    )
    return {"arxiv_results": formatted_search_docs}


@tool
def run_python_code(code_str: str) -> str:
    """Executes the provided Python code string and returns the result."""
    try:
        exec_globals = {}
        exec(code_str, exec_globals)
        return str(exec_globals.get("result", "Execution successful"))
    except Exception as e:
        return f"Error: {str(e)}"


@tool
def get_image_captioning(list_image_urls: List[str]) -> str:
    """
    generate captions for images

    Args:
        list_image_urls: list of image urls to process
    """
    try:
        loader = ImageCaptionLoader(images=list_image_urls)
        list_docs = loader.load()
        result = "Captions for given images: \n"
        for doc in list_docs:
            result += f"Image: {doc.metadata.__str__()}, Caption:{doc.page_content}"
        return result
    except Exception as e:
        return f"Image captioning failied: {str(e)}"