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

# Custom tool
@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """An example custom tool.
    Args:
        arg1: The first argument.
        arg2: The second argument.
    """
    return "Add your own magic here!"

# Get current local time in a specific timezone
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Gets the current time in a specific 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 time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

# Integration of DuckDuckGoSearchTool
@tool
def web_search(query: str) -> str:
    """Performs a search on DuckDuckGo and returns the results.
    Args:
        query: The search query.
    """
    search_tool = DuckDuckGoSearchTool()
    results = search_tool.search(query)
    return results

# Image generation tool (Placeholder example, replace with an actual API if needed)
@tool
def image_generation_tool(prompt: str) -> str:
    """Generates an image based on a given text description.
    Args:
        prompt: Description of the image to be generated.
    """
    return f"Image generated with description: {prompt}"

# Initialize the agent with available tools
model = HfApiModel("Qwen/Qwen2.5-Coder-32B-Instruct")  # Replace with the desired model
agent = CodeAgent(model=model, tools=[my_custom_tool, get_current_time_in_timezone, web_search, image_generation_tool])

# Create Gradio UI
ui = GradioUI(agent)

# Launch the UI
if __name__ == "__main__":
    ui.launch()