import os

from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import json
import yaml
from tools.final_answer import FinalAnswerTool
import wikipediaapi
from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def search_with_duckduckgo(query: str)-> str: #it's import to specify the return type
    #Keep this format for the description / args / args description but feel free to modify the tool
    """DuckDuckGo Search Engine 
    Args:
        query: The query you want to be passed to the search engine.
    """

    search = DuckDuckGoSearchTool()
    
    return search.forward(query)

@tool
def search_wikipedia(query: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        query: Search Wikipedia for a term or phrase.
    """
    wiki = wikipediaapi.Wikipedia(user_agent='Chuck Conway (charles.conway@mailme.com)',
                                  language='en')

    try:
        print(f"Searching Wikipedia for: {query}")
        page = wiki.page(query)

        if page.exists():
            # Create a dictionary with query, title, and summary
            result = {
                "query": query,
                "title": page.title,
                "summary": page.summary
            }

            return json.dumps(result, ensure_ascii=False, indent=2)
        else:
            print(f"No results found for query: {query}")
            return None

    except Exception as e:
        print(f"An error occurred while processing the Wikipedia query: {e}")
        return None


# Import tool from Hub
#image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

model = HfApiModel(
    model_id="mistralai/Mistral-7B-instruct-v0.3",
    token=os.getenv("HF_TOKEN"),
    max_tokens=5000,
)
model.last_input_token_count = 0
model.last_output_token_count = 0
    
agent = CodeAgent(
    model=model,
    tools=[FinalAnswerTool(), search_with_duckduckgo, search_wikipedia], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()