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 my_cutom_tool(arg1: str, arg2: int) -> str:
    """A magical tool that conjures a mystical incantation.
    
    Args:
        arg1: A mystical phrase or keyword.
        arg2: An intensity level of magic.
    
    Returns:
        A randomly generated incantation string.
    """
    magical_templates = [
        f"By the light of the moon and stars, '{arg1}' awakens with intensity {arg2}!",
        f"Summoning ancient powers: '{arg1}' transforms into cosmic energy of level {arg2}.",
        f"Let the magic flow! '{arg1}' reverberates with a force of {arg2} magnitudes."
    ]
    return random.choice(magical_templates)

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Fetches the current local time in a specified timezone.
    
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    
    Returns:
        A string with the current time in the given timezone or an error message.
    """
    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)}"

@tool
def get_random_joke() -> str:
    """Fetches a random joke from an online API.
    
    Returns:
        A string containing a joke.
    """
    try:
        response = requests.get("https://official-joke-api.appspot.com/random_joke")
        data = response.json()
        joke = f"{data['setup']} ... {data['punchline']}"
        return joke
    except Exception as e:
        return f"Error fetching joke: {str(e)}"

@tool
def get_motivational_quote() -> str:
    """Fetches a motivational quote from an online API.
    
    Returns:
        A string containing an inspirational quote.
    """
    try:
        response = requests.get("https://zenquotes.io/api/random")
        data = response.json()
        if isinstance(data, list) and len(data) > 0:
            quote = data[0].get('q', 'Stay positive!')
            author = data[0].get('a', 'Unknown')
            return f"\"{quote}\" - {author}"
        return "No quote found."
    except Exception as e:
        return f"Error fetching quote: {str(e)}"

# This tool always has to remain in your toolkit
final_answer = FinalAnswerTool()

# Set up the model with your desired parameters
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='https://wxknx1kg971u7k1n.us-east-1.aws.endpoints.huggingface.cloud',  # Note: This model may be overloaded at times.
    custom_role_conversions=None,
)

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

# Load your prompt templates from a YAML file
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# Create the awesome agent with a suite of magical, informative, and fun tools.
agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        my_cutom_tool,
        get_current_time_in_timezone,
        image_generation_tool,
        get_random_joke,
        get_motivational_quote,
        DuckDuckGoSearchTool(),  # Enables web search capabilities
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="AwesomeAgent",
    description="A highly creative and magical agent with time, image, search, and humor capabilities.",
    prompt_templates=prompt_templates
)

# Launch the agent in a Gradio interface
GradioUI(agent).launch()