import math
import os

import numexpr
from langchain_community.document_loaders import ArxivLoader
from smolagents import (
    CodeAgent,
    GoogleSearchTool,
    InferenceClientModel,
    WikipediaSearchTool,
    tool,
)

hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")


@tool
def calculator(expression: str) -> str:
    """Calculate expression using Python's numexpr library.

    Expression should be a single line mathematical expression
    that solves the problem.

    Examples:
        "37593 * 67" for "37593 times 67"
        "37593**(1/5)" for "37593^(1/5)"
    Args:
        expression: the expression to be compute
    """
    local_dict = {"pi": math.pi, "e": math.e}
    return str(
        numexpr.evaluate(
            expression.strip(),
            global_dict={},  # restrict access to globals
            local_dict=local_dict,  # add common mathematical functions
        )
    )


@tool
def arxiv_search_tool(query: str) -> str:
    """
    Searches Arxiv and returns a summary or full text of the given topic of the 5 most likely related papers, along with the page URL.
    Args:
        query: the topic to be searched
    """
    docs = ArxivLoader(query=query, load_max_docs=5).load()
    if len(docs) > 0:
        return "\n".join(
            [
                f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
                for doc in docs
            ]
        )

    return "No content found"


tools = [
    calculator,
    arxiv_search_tool,
    GoogleSearchTool(provider="serper"),
    WikipediaSearchTool(),
]

model_id = "Qwen/Qwen3-235B-A22B"
model = InferenceClientModel(
    model_id=model_id, token=hf_token
)  # You can choose to not pass any model_id to InferenceClientModel to use a default model
agent = CodeAgent(
    model=model,
    tools=tools,
    add_base_tools=True,
    additional_authorized_imports=["pandas", "numpy", "csv", "subprocess"],
)