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

from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def select_random_image_style()-> str: #it's import to specify the return type
    #Keep this format for the description / args / args description but feel free to modify the tool
    """Select a random image style to pass to image generation together with the user's input
    Args:
    """
    image_styles = [
    {
        "name": "Flat Design",
        "description": "Ultra-simplified, two-dimensional visuals with no depth cues.",
        "key_elements": [
            "Solid, un-textured color blocks",
            "Minimal or zero gradients/shadows",
            "Crisp vector edges",
            "High foreground–background contrast",
        ],
    },
    {
        "name": "Low-Poly",
        "description": "Imagery built from large, faceted polygons, echoing early 3-D graphics.",
        "key_elements": [
            "20–200 visible polygon faces",
            "Hard edges, no smoothing",
            "One flat color per face",
            "Directional lighting to emphasize facets",
        ],
    },
    {
        "name": "Cyberpunk Neon",
        "description": "Futuristic, dystopian cityscapes drenched in glowing signage.",
        "key_elements": [
            "High-saturation magentas, cyans, blues",
            "Rim lighting and bloom effects",
            "Rain-slicked reflections",
            "Dense vertical composition with cluttered signage",
        ],
    },
    {
        "name": "Studio Portrait (Rembrandt Lighting)",
        "description": "Classical studio portrait with the signature cheek-triangle of light.",
        "key_elements": [
            "Single softbox at ~45°",
            "Triangle of light under shadow-side eye",
            "Low-key background, gentle fall-off",
            "Natural skin tones, minimal retouch",
        ],
    },
    {
        "name": "Watercolor Wash",
        "description": "Loose, translucent paint strokes and bleeds on textured paper.",
        "key_elements": [
            "Soft edges and color bleeding",
            "Paper grain visible",
            "Low-contrast outlines or none",
            "Wet-on-wet and drybrush variation",
        ],
    },
    {
        "name": "Art Deco Poster",
        "description": "1920s–30s luxury aesthetic with bold geometry and metallic accents.",
        "key_elements": [
            "Symmetry, stepped forms, sunbursts",
            "Narrow stylized typography",
            "Gold, black, cream, teal palette",
            "Elongated figures or transport motifs",
        ],
    },
    {
        "name": "Noir Film Still",
        "description": "High-contrast B&W cinematography evoking 1940s crime drama.",
        "key_elements": [
            "Hard key plus side/back lights",
            "Deep blacks, clipped highlights",
            "Venetian-blind or window-frame shadows",
            "Urban set dressing, cigarette haze",
        ],
    },
    {
        "name": "Ghibli-Inspired Pastel Fantasy",
        "description": "Soft, whimsical landscapes reminiscent of Studio Ghibli backgrounds.",
        "key_elements": [
            "Pastel greens, blues, creams",
            "Painterly visible brush strokes",
            "Nature focus: rolling hills, fluffy clouds",
            "Tiny architectural details for scale",
        ],
    },
    {
        "name": "Brutalist Collage",
        "description": "Raw photomontage reflecting Brutalist concrete heft.",
        "key_elements": [
            "Greyscale concrete textures",
            "Blocky, asymmetric composition",
            "Bold Helvetica-style typography",
            "Occasional harsh primary accent",
        ],
    },
    {
        "name": "Glitch Vaporwave",
        "description": "Retro-digital aesthetic mixing ’80s nostalgia with data corruption artifacts.",
        "key_elements": [
            "CRT scanlines, RGB channel shifts, datamosh blocks",
            "Sunset gradients (purple→pink→orange)",
            "Greek busts, palm trees, wireframe grids",
            "Oversized Japanese katakana/kanji overlays",
        ],
    },
]
    return image_styles[random.randint(0,len(image_styles))]

@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='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, image_generation_tool, select_random_image_style], ## 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()