import math
import os
import re

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

hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"


@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"


@tool
def download_attached_files(task_id: str) -> str:
    """
    Get the first attached file if exists and save it. Returns the filename.
    Args:
        task_id: this id of the task
    """
    api_url = DEFAULT_API_URL
    file_url = f"{api_url}/files/{task_id}"
    response = requests.get(file_url, timeout=15)
    if response.status_code == 200:
        cd = response.headers.get("content-disposition")
        if cd:
            fname = re.findall("filename=(.+)", cd)
            if len(fname) > 0:
                with open(fname[0], "wb") as file:
                    file.write(response.content)
                return f"{fname[0]} saved "
    return "no file attached to this task"


tools = [
    calculator,
    arxiv_search_tool,
    download_attached_files,
    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"],
)