# from google import genai
# from google.genai.types import Tool, GenerateContentConfig, GoogleSearch
import os
from dotenv import load_dotenv
from smolagents import CodeAgent, LiteLLMModel, DuckDuckGoSearchTool, VisitWebpageTool, FinalAnswerTool
from tools import visitWikipedia, visitWikipediaTable
import time

load_dotenv()
google = os.getenv("GOOGLE")

# client = genai.Client(api_key=google)
# google_search_tool = Tool(
#     google_search = GoogleSearch()
# )
# response = client.models.generate_content(
#     model="gemini-2.0-flash",
#     contents = prompt,
#     config=GenerateContentConfig(
#         tools=[google_search_tool],
#         response_modalities=["TEXT"],
#     )
# )
# print(response.text)

model = LiteLLMModel(
    model_id="gemini/gemini-2.0-flash",
    api_key=google,
    max_tokens=1000,
    temperature=0.0,
)
tools = [FinalAnswerTool(), DuckDuckGoSearchTool(), VisitWebpageTool(), visitWikipedia, visitWikipediaTable]
agent = CodeAgent(
    tools=tools, 
    model=model,
    additional_authorized_imports=["math", "datetime", "re", "json", "pandas", "numpy"],
    max_steps=15,
    )

prompt = """
Do not assume your code is correct. Always check your code and the output of your code.
You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the FinalAnswerTool(). YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Question:
"""

def run(question):
    try: 
        return agent.run(task=prompt+"\n"+question)
    except Exception as e:
        time.sleep(60)
        return agent.run(task=prompt+"\n"+question)

question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."

# answer = run(question)
# print(answer)