from smolagents import CodeAgent, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import os
from PIL import Image, ImageDraw, ImageFont
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Ensure the latest smolagents version: `pip install --upgrade smolagents`
# Ensure prompts.yaml matches the provided structure with system_prompt, final_answer, planning, and managed_agent
# Ensure Pillow is installed: `pip install Pillow`
# Set your Hugging Face API token as an environment variable if required
HF_TOKEN = os.environ.get("HF_TOKEN")

@tool
def get_motivational_quote(category: str = "inspirational") -> str:
    """A tool that fetches a random motivational quote from the Quotable API.
    Args:
        category: A string representing the quote category (e.g., 'inspirational', 'life'). Defaults to 'inspirational'.
    """
    try:
        response = requests.get(f"https://api.quotable.io/random?tags={category}")
        response.raise_for_status()
        data = response.json()
        quote = data['content']
        author = data['author']
        return f"\"{quote}\" - {author}"
    except Exception as e:
        return f"Error fetching quote: {str(e)}"

@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)}"

@tool
def overlay_quote_on_image(image: object, quote: str) -> str:
    """A tool that overlays a motivational quote on an image and saves it to a file.
    Args:
        image: An AgentImage object from text-to-image tool.
        quote: A string containing the quote to overlay.
    Returns:
        A string representing the file path of the saved image.
    """
    try:
        # Convert AgentImage to PIL Image (assuming AgentImage has a to_pil() method or similar)
        pil_image = image.to_pil() if hasattr(image, 'to_pil') else Image.open(image)
        draw = ImageDraw.Draw(pil_image)
        
        # Load a default font (or specify a path to a .ttf file if available)
        try:
            font = ImageFont.truetype("arial.ttf", 40)
        except:
            font = ImageFont.load_default()
        
        # Calculate text size and position
        text = quote
        text_width, text_height = draw.textsize(text, font=font)
        image_width, image_height = pil_image.size
        text_position = ((image_width - text_width) // 2, image_height - text_height - 50)
        
        # Draw text with a black outline and white fill
        draw.text(text_position, text, font=font, fill="white", stroke_width=2, stroke_fill="black")
        
        # Save the image to a file
        output_path = "/tmp/output_image_with_quote.png"
        pil_image.save(output_path)
        return output_path
    except Exception as e:
        return f"Error overlaying quote on image: {str(e)}"

final_answer = FinalAnswerTool()
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)

# Restructure prompt_templates to match CodeAgent expectations
restructured_templates = {
    "system_prompt": prompt_templates.get("system_prompt", ""),
    "final_answer": prompt_templates.get("final_answer", ""),
    "planning": {
        "initial_facts": prompt_templates.get("planning", {}).get("initial_facts", ""),
        "initial_plan": prompt_templates.get("planning", {}).get("initial_plan", ""),
        "update_facts_pre_messages": prompt_templates.get("planning", {}).get("update_facts_pre_messages", ""),
        "update_facts_post_messages": prompt_templates.get("planning", {}).get("update_facts_post_messages", ""),
        "update_plan_pre_messages": prompt_templates.get("planning", {}).get("update_plan_pre_messages", ""),
        "update_plan_post_messages": prompt_templates.get("planning", {}).get("update_plan_post_messages", "")
    },
    "managed_agent": {
        "task": prompt_templates.get("managed_agent", {}).get("task", ""),
        "report": prompt_templates.get("managed_agent", {}).get("report", "")
    }
}

model = HfApiModel(
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    token=HF_TOKEN,
    max_tokens=2096,
    temperature=0.5,
    custom_role_conversions=None
)

agent = CodeAgent(
    model=model,
    tools=[get_motivational_quote, get_current_time_in_timezone, image_generation_tool, overlay_quote_on_image, final_answer],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=restructured_templates
)

GradioUI(agent).launch()