from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    HfApiModel,
    load_tool,
    tool,
    FinalAnswerTool,
    GradioUI
)
from typing import Optional
import datetime
import requests
import pytz
import yaml

@tool
def get_travel_duration(start_location: str, destination_location: str, transportation_mode: Optional[str] = None) -> str:
    """
    Get travel time between two places.

    Args:
        start_location: Place from which you start your drive.
        destination_location: Place of arrival.
        transportation_mode: Transportation mode, e.g. 'walking', 'driving', 'bicycling', 'transit','motobike'.
    """
    import os
    import googlemaps
    from datetime import datetime

    gmaps = googlemaps.Client(os.getenv("AIzaSyAFjYCOmBgyvaNNFSk1rW71hX1nWbrR6Ew"))

    if transportation_mode  is None:
        transportation_mode = "motobike"
    
    try:
        directions_result = gmaps.directions(
            start_location,
            destination_location,
            mode = transportation_mode,
            departure_time = datetime(2026, 6, 6, 11, 0)

        )
        if len(directions_result) ==0:
            return "No way found between these places with  the required transportation mode."
        return directions_result[0]['legs'][0]["duration"]["text"]
    except Exception as e:
        print(e)
        return 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 get_current_weather(location: str) -> str:
    """A tool that fetches the current weather for a given city/location.
    Args:
        location: The name of the city or location (e.g., 'Ho Chi Minh City').
    """
    try:
        # Geocode
        geo_response = requests.get(
            f"https://geocoding-api.open-meteo.com/v1/search?name={location}&count=1"
        ).json()

        if not geo_response.get("results"):
            return f"Could not find location: {location}"

        result = geo_response["results"][0]
        lat, lon = result["latitude"], result["longitude"]
        city_name = result["name"]
        country = result.get("country", "")

        # Fetch weather
        weather_response = requests.get(
            f"https://api.open-meteo.com/v1/forecast?"
            f"latitude={lat}&longitude={lon}"
            f"&current_weather=true"
        ).json()
        current = weather_response["current_weather"]

        weather_codes = {
            0: "Clear sky", 1: "Mainly clear", 2: "Partly cloudy", 3: "Overcast",
            45: "Foggy", 48: "Icy fog", 51: "Light drizzle", 53: "Moderate drizzle",
            61: "Slight rain", 63: "Moderate rain", 65: "Heavy rain",
            71: "Slight snow", 73: "Moderate snow", 75: "Heavy snow",
            80: "Slight showers", 81: "Moderate showers", 82: "Violent showers",
            95: "Thunderstorm", 96: "Thunderstorm with hail",
        }
        description = weather_codes.get(current["weathercode"], "Unknown")

        return (
            f"Weather in {city_name}, {country}:\n"
            f"  🌡️  Temperature     : {current['temperature']}°C\n"
            f"  💨  Wind Speed      : {current['windspeed']} km/h\n"
            f"  🧭  Wind Direction  : {current['winddirection']}°\n"
            f"  🌤️  Condition       : {description}\n"
        )

    except Exception as e:
        return f"Error fetching weather for '{location}': {str(e)}"


final_answer = FinalAnswerTool()

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

image_generation_tool = load_tool("m-ric/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,
        get_current_time_in_timezone,
        get_current_weather,          # ✅ new weather tool
        get_travel_duration,
        DuckDuckGoSearchTool(),
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()