from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import pytz
import random
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Load external tool before function declarations
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

@tool
def tell_joke() -> str:
    """Returns a random joke."""
    print("DEBUG: tell_joke() was called")
    jokes = [
        "Why don't skeletons fight each other? They don't have the guts.",
        "I told my wife she was drawing her eyebrows too high. She looked surprised.",
        "Why don’t scientists trust atoms? Because they make up everything!"
    ]
    return random.choice(jokes)

@tool
def random_fact() -> str:
    """Fetches a random fun fact."""
    print("DEBUG: random_fact() was called")
    facts = [
        "Honey never spoils.",
        "Octopuses have three hearts.",
        "Bananas are berries, but strawberries are not."
    ]
    return random.choice(facts)

@tool
def generate_image_with_prompt(prompt: str) -> str:
    """Generates an image based on the provided prompt."""
    print(f"DEBUG: generate_image_with_prompt() was called with prompt: {prompt}")
    try:
        image_url = image_generation_tool.run(prompt)
        return f"Here's your generated image: {image_url}"
    except Exception as e:
        return f"Error generating image: {str(e)}"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Fetches the current local time in a specified timezone."""
    print(f"DEBUG: get_current_time_in_timezone() was called with timezone: {timezone}")
    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-Coder-32B-Instruct'
)

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

agent = CodeAgent(
    model=model,
    tools=[
        tell_joke,
        random_fact,
        generate_image_with_prompt,
        get_current_time_in_timezone,
        final_answer,
        DuckDuckGoSearchTool()
    ],
    max_steps=6,
    verbosity_level=2,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()