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


@tool
def get_color_palette_based_on_emotions(emotions: str) -> str:
    """A tool that creates the color palette of a given emotions.
    Args:
        emotions: the emotions to create the color palette of.
    """
    model = HfApiModel(
        max_tokens=2096,
        temperature=0.5,
        model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
        custom_role_conversions=None,
    )

    messages = [
        {
            "role": "system",
            "content": (
                "ROLE: you are a cute color palette generator. "
                "TASK: take the emotions and create a color palette which will complement the emotions. Use only pastel colors. "
                "OUTPUT: a list of 6 colors in hex format. "
                "CONSTRAINTS: Do not add any words or explanations. Just return the list of colors."
            ),
        },
        {
            "role": "user",
            "content": emotions,
        },
    ]
    response = model(messages, stop_sequences=["END"])
    return response.content


@tool
def get_postcard_prompt_based_on_color_palette_and_greeting(
    color_palette: str, greeting: str
) -> str:
    """A tool that creates the image prompt for post card based on the color palette and greeting.
    Args:
        color_palette: 6 colors in hex format to create the image prompt of.
        greeting: the greeting to create the image prompt of.
    """
    model = HfApiModel(
        max_tokens=2096,
        temperature=0.5,
        model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
        custom_role_conversions=None,
    )

    messages = [
        {
            "role": "system",
            "content": (
                "ROLE: you are a cute postcard image prompt generator. "
                "TASK: take the color palette and greeting and create a postcard image prompt which simbolses the greeting. "
                "Use simple vector shapes, and pastel colors from the color palette. "
                "OUTPUT: a postcard image prompt. "
                "CONSTRAINTS: Use only the colors from the color palette. Include the greeting into the image prompt so it would be rendered as a text. The prompt should be properly formatted as a string and contain escape characters for new lines, quotes and special characters. "
            ),
        },
        {
            "role": "user",
            "content": f"Color palette: {color_palette}, Greeting: {greeting}",
        },
    ]
    response = model(messages, stop_sequences=["END"])
    return response.content


@tool
def get_comic_story_series_prompts_based_on_greeting(greeting: str) -> str:
    """A tool that generates a series of 4 comic-style postcard prompts based on a greeting, creating a sequential story to generate 4 postcards images.
    Args:
        greeting: the greeting to create the image prompts of.
    """
    model = HfApiModel(
        max_tokens=2096,
        temperature=0.5,
        model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
        custom_role_conversions=None,
    )

    messages = [
        {
            "role": "system",
            "content": (
                "ROLE: you are a cute story writer for a postcard series. "
                "TASK: take the greeting and create a series of 4 postcard image prompts which simbolses the greeting. "
                "Use simple vector shapes, and pastel colors. "
                "OUTPUT: a series of 4 postcard image prompts. "
                "CONSTRAINTS: Use only the colors from the color palette. Use simple texts to create the story. The prompt should be properly formatted as a string and contain escape characters for new lines, quotes and special characters. "
            ),
        },
        {
            "role": "user",
            "content": f"Greeting: {greeting}",
        },
    ]
    response = model(messages, stop_sequences=["END"])
    return response.content


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="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,
        get_color_palette_based_on_emotions,
        get_postcard_prompt_based_on_color_palette_and_greeting,
        get_comic_story_series_prompts_based_on_greeting,
        image_generation_tool,
    ],  ## 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()