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


# Load HF image generation tool
hf_image_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)


@tool
def generate_image(prompt: str) -> str:
    """
    Generate an image from a text prompt.

    Args:
        prompt (str): Description of the image.

    Returns:
        str: File path of the generated image.
    """
    img = hf_image_tool(prompt=prompt)
    path = "/tmp/generated_image.png"
    img.save(path)
    return path


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Get current time for a timezone.

    Args:
        timezone (str): Example 'Asia/Kolkata'.

    Returns:
        str: Current time in that timezone.
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"Time in {timezone}: {local_time}"
    except Exception as e:
        return str(e)


# Mandatory final answer tool
final_answer = FinalAnswerTool()


# LLM for reasoning
model = HfApiModel(
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    max_tokens=2048,
    temperature=0.5
)


# Prompt templates
with open("prompts.yaml", "r") as f:
    prompt_templates = yaml.safe_load(f)


# Agent
agent = CodeAgent(
    model=model,
    tools=[
        generate_image,
        DuckDuckGoSearchTool(),
        get_current_time_in_timezone,
        final_answer
    ],
    max_steps=6,
    verbosity_level=2,
    name="multitool_agent",
    description="Agent that can generate images, search the web, and get timezone time.",
    prompt_templates=prompt_templates
)


# Launch UI
GradioUI(agent).launch()