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
from diffusers import StableDiffusionPipeline
import torch
from io import BytesIO
import base64

# Below is an example of a tool that does nothing. Amaze us with your creativity !

# @tool

class ImageGenerator:
    """A class to generate images from text prompts using Stable Diffusion."""
    def __init__(self, model_id="runwayml/stable-diffusion-v1-5", device="cuda" if torch.cuda.is_available() else "cpu"):
        """Initializes the ImageGenerator with the Stable Diffusion pipeline."""
        self.pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32).to(device)
        self.device = device

    def generate_image(self, prompt, num_inference_steps=25, guidance_scale=7.5):
        """Generates an image from a text prompt.
        Args:
            prompt (str): The text prompt to generate the image from.
            num_inference_steps (int): The number of inference steps.
            guidance_scale (float): The guidance scale.
        Returns:
            PIL.Image.Image: The generated image.
        """
        image = self.pipeline(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
        return image

    def generate_base64_image(self, prompt, num_inference_steps=25, guidance_scale=7.5):
        """Generates a base64 encoded image from a text prompt.
        Args:
            prompt (str): The text prompt to generate the image from.
            num_inference_steps (int): The number of inference steps.
            guidance_scale (float): The guidance scale.
        Returns:
            str: The base64 encoded image.
        """
        image = self.generate_image(prompt, num_inference_steps, guidance_scale)
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        return img_str

def generate_image_tool(image_generator):
    """Creates a tool function for image generation."""
    def image_generation_tool(prompt):
        """Generates an image from a prompt."""
        return image_generator.generate_base64_image(prompt)
    return image_generation_tool

# Initialize the ImageGenerator and tool
image_generator = ImageGenerator()
image_generation_tool_function = generate_image_tool(image_generator)

@tool
def generate_image_from_prompt(prompt: str) -> str:
    """Generates an image from a text prompt and embeds it in an HTML img tag.
    Args:
        prompt: The text prompt to generate the image from.
    """
    base64_image = image_generation_tool_function(prompt)
    return f'<img src="data:image/png;base64,{base64_image}" alt="Generated Image" />'

@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:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that 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()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
# model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
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)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, generate_image_from_prompt, get_current_time_in_timezone], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()