from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import random

from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI

@tool
def get_trivia_question(category: str, difficulty: str) -> str:
    """Fetches a random trivia question from the Open Trivia Database.

    Args:
        category: The category of the question. Must be one of:
            - 'General Knowledge'
            - 'Books'
            - 'Film'
            - 'Music'
            - 'Video Games'
            - 'Computers'
            - 'Mathematics'
            - 'Mythology'
            - 'Sports'
        difficulty: The difficulty level of the question. Must be 'easy', 'medium', or 'hard'.

    Returns:
        A formatted trivia question with possible answers.
    """
    category_map = {
        "General Knowledge": 9,
        "Books": 10,
        "Film": 11,
        "Music": 12,
        "Video Games": 15,
        "Computers": 18,
        "Mathematics": 19,
        "Mythology": 20,
        "Sports": 21
    }
    if category not in category_map:
        return f"Invalid category. Must be one of: {', '.join(category_map.keys())}"
    if difficulty not in ["easy", "medium", "hard"]:
        return "Invalid difficulty. Must be 'easy', 'medium', or 'hard'."
    question_type = random.choice(["multiple", "boolean"])
    params = {
        "amount": 1,
        "category": category_map[category],
        "difficulty": difficulty,
        "type": question_type
    }
    response = requests.get("https://opentdb.com/api.php", params=params)
    if response.status_code != 200:
        return "An error occurred while fetching a trivia question."
    data = response.json()
    if data["response_code"] != 0:
        return "No question found for the given criteria."
    question_data = data["results"][0]
    question = question_data["question"]
    correct_answer = question_data["correct_answer"]
    incorrect_answers = question_data["incorrect_answers"]
    if question_type == "multiple":
        options = incorrect_answers + [correct_answer]
        random.shuffle(options)
        options_text = "\n".join([f"- {opt}" for opt in options])
        return f"Trivia Question: {question}\n\nPossible answers:\n{options_text}\n\nWhat is your answer?"
    else:
        return f"Trivia Question: {question}\n\nAnswer: True or False?"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

final_answer = FinalAnswerTool()
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-7B-Instruct-1M',
    custom_role_conversions=None,
)

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)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_trivia_question, get_current_time_in_timezone],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()