from smolagents import HfApiModel, CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, Tool, LiteLLMModel
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
import os, time, math, pandas, numpy


class add(Tool):
    name = "add"
    description = """
        This tool adds two integers together and returns an integer.

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        return a + b


class subtract(Tool):
    name = "subtract"
    description = """
        This tool subtracts two integers and returns an integer.

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        return a - b

class multiply(Tool):
    name = "multiply"
    description = """
        This tool multiplies two integers and returns an integer.

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        return a * b

class divide(Tool):
    name = "divide"
    description = """
        This tool divides two integers and returns an integer.

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        if b == 0:
            raise ValueError("Cannot divide by zero.")
        return a / b

class modulo(Tool):
    name = "modulo"
    description = """
        This tool returns the modulus of two integers

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        return a % b

class WikipediaSearchTool(Tool):
    name = "wikipedia_search_tool"
    
    description = """
        Search Wikipedia for a query and return top 2 results.
    
        Args:
            query: the search query.
    """
    
    inputs = {
        "query":{
            "type":"string",
            "description":"the search query"
        }
    }
    output_type = "string"

    def forward(self, query: str) -> str:
        documents = WikipediaLoader(query=query, load_max_docs=2).load()
        condensed_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 documents
        ])
        return {"wikipedia_results": condensed_docs}

class TavilySearchTool(Tool):
    name = "tavily_search_tool"
    
    description = """
        Search the internet using the browser Tavily and return the top 3 results.
    
        Args:
            query: the search query.
    """
    
    inputs = {
        "query":{
            "type":"string",
            "description":"the search query"
        }
    }
    output_type = "string"

    def forward(self, query: str) -> str:
        documents = TavilySearchResults(max_results=3).invoke(query=query)
        condensed_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 documents
        ])
        return {"web_search_results": condensed_docs}


class ArvixSearchTool(Tool):
    name = "arvix_search_tool"
    
    description = """
        Search arxiv for a query and return maximum 3 result.
    
        Args:
            query: The search query.
        """
    
    inputs = {
        "query":{
            "type":"string",
            "description":"the search query"
        }
    }
    output_type = "string"

    def forward(self, query: str) -> str:
        documents = ArxivLoader(query=query, load_max_docs=3).load()
        condensed_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 documents
        ])
        return {"arvix_search_results": condensed_docs}


tools=[
    add(),
    subtract(),
    multiply(),
    divide(),
    modulo(),
    WikipediaSearchTool(),
    ArvixSearchTool(),
    TavilySearchTool()
]

model = LiteLLMModel(model_id="gemini/gemini-2.0-flash-exp",
                     api_key=os.getenv("GEMINI_API_KEY"))

def delay_execution_10(pagent, **kwargs) -> bool:
    """
    Delays the execution for 10 seconds.
    """
    time.sleep(10)
    return True

def create_agent():
    agent = CodeAgent(
        model = model,
        tools = tools,
        max_steps=10,
        verbosity_level=2,
        additional_authorized_imports=['*'],
        planning_interval=5,  
        step_callbacks=[delay_execution_10]
    )

    return agent