from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from io import BytesIO
from PIL import Image
import base64, os, uuid

from Gradio_UI import GradioUI

_CITY_TZ = {
    "berlin": "Europe/Berlin",
    "munich": "Europe/Berlin",
    "frankfurt": "Europe/Berlin",
    "london": "Europe/London",
    "paris": "Europe/Paris",
    "rome": "Europe/Rome",
    "madrid": "Europe/Madrid",
    "vienna": "Europe/Vienna",
    "zurich": "Europe/Zurich",
    "new york": "America/New_York",
    "nyc": "America/New_York",
    "san francisco": "America/Los_Angeles",
    "los angeles": "America/Los_Angeles",
    "sf": "America/Los_Angeles",
    "tokyo": "Asia/Tokyo",
    "singapore": "Asia/Singapore",
    "hong kong": "Asia/Hong_Kong",
    "shanghai": "Asia/Shanghai",
    "beijing": "Asia/Shanghai",
    "seoul": "Asia/Seoul",
    "sydney": "Australia/Sydney",
    "melbourne": "Australia/Melbourne",
    "delhi": "Asia/Kolkata",
    "mumbai": "Asia/Kolkata",
    "toronto": "America/Toronto",
    "vancouver": "America/Vancouver",
}

@tool
def city_to_timezone(city: str) -> str:
    """Map a city name to its IANA timezone string.
    Args:
        city: City name (e.g. 'Munich', 'San Francisco', 'Tokyo').
    """
    key = city.strip().lower()
    if key in _CITY_TZ:
        return _CITY_TZ[key]
    # fallback: try direct match in pytz
    for tz in pytz.all_timezones:
        if tz.lower().endswith("/" + key.replace(" ", "_")):
            return tz
    return f"Unknown city: {city}"
    
from smolagents import tool
import datetime as _dt
import pytz

@tool
def cat_image_prompt_generator(local_time: str, timezone: str = "Europe/Berlin") -> str:
    """Generate a cute cat image prompt based on local time.
    Args:
        local_time: A string 'YYYY-MM-DD HH:MM'
        timezone: IANA timezone string (default: 'Europe/Berlin')
    """
    try:
        tz = pytz.timezone(timezone)
    except Exception:
        return f"Unknown timezone: {timezone}"

    # Parse datetime
    if local_time.strip().lower() == "now":
        now = _dt.datetime.now(tz)
    else:
        try:
            now = _dt.datetime.strptime(local_time.strip(), "%Y-%m-%d %H:%M")
            now = tz.localize(now)
        except Exception:
            return "Invalid datetime format. Use 'YYYY-MM-DD HH:MM' or 'now'."

    minutes = now.hour * 60 + now.minute

    def in_range(start, end):
        return start <= minutes < end

    # Decide state
    if in_range(7*60, 8*60) or in_range(12*60, 13*60) or in_range(18*60, 19*60):
        cat_state = "eat"
    elif in_range(21*60, 24*60) or in_range(0, 2*60) or in_range(5*60, 7*60):
        cat_state = "play"
    else:
        cat_state = "sleep"

    # Prompt templates
    prompts = {
        "eat":  "2 cute cats are happily eating from small bowls indoors, warm soft lighting, cozy atmosphere, photorealistic.",
        "play": "2 playful cats are chasing each other with colorful toys on a wooden floor, bright daylight, motion blur, joyful atmosphere.",
        "sleep":"2 cute cats are sleeping curled up on a soft blanket, peaceful, warm ambient light, cozy bedroom."
    }

    return {prompts[cat_state]}
        

@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")
        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=[city_to_timezone, get_current_time_in_timezone, cat_image_prompt_generator, image_generation_tool, final_answer], ## 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()